Pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e5452609-26f5-4ff9-8740-03b41421b589",
   "metadata": {},
   "source": [
    "# Evaluating the retrival pipeline\n",
    "\n",
    "## On Top Level MITRE ATT&CK Techniques"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "746651e5-8af7-443d-9e0c-dcc4a8840524",
   "metadata": {
    "is_executing": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "from libs.pygaggle.data.relevance import RelevanceExample\n",
    "from libs.pygaggle.model import StepEvaluator\n",
    "\n",
    "from libs import resources as res, rank"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0887a976-555b-4399-9161-df4c3f230b98",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "if device == 'cuda':\n",
    "    cuda_device = 7\n",
    "    torch.cuda.set_device(cuda_device)\n",
    "\n",
    "cache_dir = 'cache'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87563809-030f-4f49-b5e5-b4aa4babb7b0",
   "metadata": {},
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d500882-37dd-420a-8238-4d610829a0b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "sentences = res.load_annotated()\n",
    "sentences = sentences[sentences.tech_id.str.len() == 5]\n",
    "\n",
    "sentences.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36e306eb-88a5-4fe5-bd5f-ad40ad6ba2a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = res.load_mitre_kb()\n",
    "dataset = dataset[dataset.tech_id.str.len() == 5]\n",
    "\n",
    "corpus = pd.DataFrame([{\n",
    "    'tech_id': tech_id,\n",
    "    'text': name + ' ' + ' '.join(g['text'].values)\n",
    "} for (tech_id, name), g in dataset.groupby(['tech_id', 'tech_name'])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "198b49cb-05cf-48ee-8bf2-d32d1ba47c4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "queries = sentences[sentences['tech_id'].isin(dataset.tech_id.values)]\n",
    "len(queries['tech_id'].drop_duplicates())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbcb7746-f3a7-44ea-b4d9-6a71338dc3bf",
   "metadata": {},
   "source": [
    "## Rerank"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09a2c76e-6cfe-4994-8ad3-531c8b6d353a",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = ['recall@3', 'recall@5', 'recall@10', 'recall@20', 'recall@40', 'recall@50', 'recall@100', 'mrr']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e11b3ee-fbd3-4585-98ba-965b28d4f420",
   "metadata": {},
   "source": [
    "### Generate Initial Examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ba23aaa-2cc6-4185-9621-3ea574aa1740",
   "metadata": {},
   "outputs": [],
   "source": [
    "texts, label_map = rank.get_texts(corpus)\n",
    "queries = rank.get_queries(queries, label_col='tech_id')\n",
    "\n",
    "examples = [RelevanceExample(\n",
    "    query,\n",
    "    texts,\n",
    "    [(True if label_map[query.metadata['label']] == i else False) for i in range(len(texts))]\n",
    ") for query in queries]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46a5bf5f-dd32-4f0e-89b9-829ebb4c1860",
   "metadata": {},
   "source": [
    "### First Stage\n",
    "#### BM25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21faeb24-0d92-420a-bba0-a54b517caa69",
   "metadata": {},
   "outputs": [],
   "source": [
    "def stage1_runner(examples):\n",
    "    bm25_reranker = rank.construct_bm25()\n",
    "    bm25_eval = StepEvaluator(bm25_reranker, metrics, n_hits=100)\n",
    "    return bm25_eval.evaluate(examples)\n",
    "\n",
    "\n",
    "ranker_suffix = f'bm25__100'\n",
    "stage1_cache_file = f'../../data/cache/top_level__{ranker_suffix}.pkl'\n",
    "bm25_exmp, bm25_metrics = rank.load_cache_or_run(stage1_cache_file, stage1_runner, examples=examples)\n",
    "\n",
    "for metric in bm25_metrics:\n",
    "    print(f'{metric.name} = {metric.value:.5}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24809167-b7f1-4584-87d4-91505ff5b52d",
   "metadata": {},
   "source": [
    "### Second Stage\n",
    "#### SentSecBert"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "310f9a13-932c-455b-a523-9e62ae2c08fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def stage2_runner(examples, seg_size, stride):\n",
    "    sentsecbert_sim_reranker = rank.construct_sentsecbert('../../models/SentSecBert_10k_AllDataSplit')\n",
    "    sentsecbert_eval = StepEvaluator(sentsecbert_sim_reranker, metrics, n_hits=50)\n",
    "    return sentsecbert_eval.evaluate_by_segments(\n",
    "        examples,\n",
    "        seg_size=seg_size,\n",
    "        stride=stride,\n",
    "        aggregate_method='max'\n",
    "    )\n",
    "\n",
    "\n",
    "seg_size = 14\n",
    "overlap = 0.25\n",
    "stride = seg_size - int(seg_size * overlap)\n",
    "\n",
    "ranker_suffix = f'sentsecbert__50__{seg_size}__{str(overlap).replace(\".\", \"_\")}.pkl'\n",
    "stage2_cache_file = f'../../data/cache/top_level__{ranker_suffix}.pkl'\n",
    "sentsecbert_exmp, sentsecbert_metrics = rank.load_cache_or_run(stage2_cache_file, stage2_runner,\n",
    "                                                               examples=bm25_exmp, seg_size=seg_size, stride=stride)\n",
    "\n",
    "for metric in sentsecbert_metrics:\n",
    "    print(f'{metric.name} = {metric.value:.5}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c82e5d13-5bc1-470d-8995-d4524da7db31",
   "metadata": {},
   "source": [
    "### Third Stage\n",
    "#### MonoT5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e664a673-7163-4e28-890a-119e6db3da46",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def stage3_runner(examples, seg_size, stride):\n",
    "    monot5_tram_reranker = rank.construct_monot5('../../models/monot5_AllDataSplit')\n",
    "    monot5_tram_eval = StepEvaluator(monot5_tram_reranker, metrics, n_hits=10)\n",
    "    return monot5_tram_eval.evaluate_by_segments(\n",
    "        examples,\n",
    "        seg_size=seg_size,\n",
    "        stride=stride,\n",
    "        aggregate_method='max'\n",
    "    )\n",
    "\n",
    "\n",
    "seg_size = 14\n",
    "overlap = 0.25\n",
    "stride = seg_size - int(seg_size * overlap)\n",
    "\n",
    "ranker_suffix = f'monot5__10__{seg_size}__{str(overlap).replace(\".\", \"_\")}.pkl'\n",
    "stage3_cache_file = f'../../data/cache/top_level__{ranker_suffix}.pkl'\n",
    "_, monot5_tram_metrics = rank.load_cache_or_run(stage3_cache_file, stage3_runner,\n",
    "                                                examples=sentsecbert_exmp, seg_size=seg_size, stride=stride)\n",
    "\n",
    "for metric in monot5_tram_metrics:\n",
    "    print(f'{metric.name} = {metric.value:.5}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}