BRON_internal_links.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:23:58.663738Z",
     "iopub.status.busy": "2022-05-30T09:23:58.660537Z",
     "iopub.status.idle": "2022-05-30T09:24:01.286098Z",
     "shell.execute_reply": "2022-05-30T09:24:01.284737Z"
    }
   },
   "outputs": [],
   "source": [
    "from sklearn.feature_extraction.text import CountVectorizer\n",
    "import numpy as np\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import accuracy_score, roc_auc_score, f1_score, confusion_matrix\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from transformers import AutoModel, AutoTokenizer, AutoModelForMaskedLM\n",
    "import json\n",
    "import spacy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:01.293666Z",
     "iopub.status.busy": "2022-05-30T09:24:01.292458Z",
     "iopub.status.idle": "2022-05-30T09:24:01.297331Z",
     "shell.execute_reply": "2022-05-30T09:24:01.296343Z"
    }
   },
   "outputs": [],
   "source": [
    "negative_example_ids = []\n",
    "negative_example_json = json.load(open(\"negative_examples_100_2.json\"))\n",
    "for example in negative_example_json:\n",
    "    negative_example_ids.append((example[\"ap\"], example[\"technique\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:01.319057Z",
     "iopub.status.busy": "2022-05-30T09:24:01.307988Z",
     "iopub.status.idle": "2022-05-30T09:24:01.367929Z",
     "shell.execute_reply": "2022-05-30T09:24:01.367258Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(\"w_dict.json\", \"r\") as f:\n",
    "    w_dict = json.load(f)\n",
    "with open(\"ap_dict.json\", \"r\") as f:\n",
    "    ap_dict = json.load(f)\n",
    "with open(\"technique_dict.json\", \"r\") as f:\n",
    "    technique_dict = json.load(f)\n",
    "with open(\"tactic_dict.json\", \"r\") as f:\n",
    "    tactic_dict = json.load(f)\n",
    "with open(\"cwe_names.json\", \"r\") as f:\n",
    "    cwe_names = json.load(f)\n",
    "with open(\"ap_names.json\", \"r\") as f:\n",
    "    ap_names = json.load(f)\n",
    "with open(\"technique_names.json\", \"r\") as f:\n",
    "    technique_names = json.load(f)\n",
    "with open(\"tactic_names.json\", \"r\") as f:\n",
    "    tactic_names = json.load(f)\n",
    "with open(\"cwe_short_descriptions.json\", \"r\") as f:\n",
    "    cwe_short_descriptions = json.load(f)\n",
    "with open(\"ap_short_descriptions.json\", \"r\") as f:\n",
    "    ap_short_descriptions = json.load(f)\n",
    "with open(\"technique_short_descriptions.json\", \"r\") as f:\n",
    "    technique_short_descriptions = json.load(f)\n",
    "with open(\"tactic_short_descriptions.json\", \"r\") as f:\n",
    "    tactic_short_descriptions = json.load(f)\n",
    "with open(\"cwe_descriptions.json\", \"r\") as f:\n",
    "    cwe_descriptions = json.load(f)\n",
    "with open(\"ap_descriptions.json\", \"r\") as f:\n",
    "    ap_descriptions = json.load(f)\n",
    "with open(\"technique_descriptions.json\", \"r\") as f:\n",
    "    technique_descriptions = json.load(f)\n",
    "with open(\"tactic_descriptions.json\", \"r\") as f:\n",
    "    tactic_descriptions = json.load(f)\n",
    "\n",
    "with open(\"ap_mitigation_descriptions.json\", \"r\") as f:\n",
    "    ap_mitigation_descriptions = json.load(f)\n",
    "with open(\"cwe_mitigation_descriptions.json\", \"r\") as f:\n",
    "    cwe_mitigation_descriptions = json.load(f)\n",
    "with open(\"tech_mitigation_names.json\", \"r\") as f:\n",
    "    tech_mitigation_names = json.load(f)\n",
    "with open(\"tech_detection_names.json\", \"r\") as f:\n",
    "    tech_detection_names = json.load(f)\n",
    "with open(\"ap_detection_descriptions.json\", \"r\") as f:\n",
    "    ap_detection_descriptions = json.load(f)\n",
    "with open(\"cwe_detection_descriptions.json\", \"r\") as f:\n",
    "    cwe_detection_descriptions = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:01.379145Z",
     "iopub.status.busy": "2022-05-30T09:24:01.377973Z",
     "iopub.status.idle": "2022-05-30T09:24:01.413876Z",
     "shell.execute_reply": "2022-05-30T09:24:01.414676Z"
    }
   },
   "outputs": [],
   "source": [
    "f = open(\"cwe_mitigation_ids_temp.json\")\n",
    "w_mitigation = json.load(f)\n",
    "\n",
    "f = open(\"capec_mitigation_temp.json\")\n",
    "ap_mitigation = json.load(f)\n",
    "\n",
    "f = open(\"technique_mitigation_temp.json\")\n",
    "technique_mitigation = json.load(f)\n",
    "\n",
    "f = open(\"technique_detection_temp.json\")\n",
    "technique_detection = json.load(f)\n",
    "\n",
    "f = open(\"capec_detection_temp.json\")\n",
    "ap_detection = json.load(f)\n",
    "\n",
    "f = open(\"cwe_detection_temp.json\")\n",
    "w_detection = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:01.422624Z",
     "iopub.status.busy": "2022-05-30T09:24:01.421630Z",
     "iopub.status.idle": "2022-05-30T09:24:01.427079Z",
     "shell.execute_reply": "2022-05-30T09:24:01.425910Z"
    }
   },
   "outputs": [],
   "source": [
    "positive_example_ids = []\n",
    "for ap in ap_dict:\n",
    "    for technique in ap_dict[ap][\"techniques\"]:\n",
    "        positive_example_ids.append((ap, technique))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:01.438394Z",
     "iopub.status.busy": "2022-05-30T09:24:01.434037Z",
     "iopub.status.idle": "2022-05-30T09:24:01.444269Z",
     "shell.execute_reply": "2022-05-30T09:24:01.443235Z"
    }
   },
   "outputs": [],
   "source": [
    "f = open(\"capec_capec.json\")\n",
    "capec_capec = json.load(f)\n",
    "\n",
    "f = open(\"cwe_cwe.json\")\n",
    "cwe_cwe = json.load(f)\n",
    "\n",
    "f = open(\"technique_technique.json\")\n",
    "technique_technique = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:01.453369Z",
     "iopub.status.busy": "2022-05-30T09:24:01.452081Z",
     "iopub.status.idle": "2022-05-30T09:24:01.454794Z",
     "shell.execute_reply": "2022-05-30T09:24:01.455726Z"
    }
   },
   "outputs": [],
   "source": [
    "for ap in ap_dict:\n",
    "    ap_dict[ap][\"internals\"] = []\n",
    "\n",
    "for link in capec_capec:\n",
    "    ap_dict[link[\"_from\"]][\"internals\"].append(link[\"_to\"])\n",
    "    ap_dict[link[\"_to\"]][\"internals\"].append(link[\"_from\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:01.464976Z",
     "iopub.status.busy": "2022-05-30T09:24:01.463846Z",
     "iopub.status.idle": "2022-05-30T09:24:01.468211Z",
     "shell.execute_reply": "2022-05-30T09:24:01.467072Z"
    }
   },
   "outputs": [],
   "source": [
    "for cwe in w_dict:\n",
    "    w_dict[cwe][\"internals\"] = []\n",
    "\n",
    "for link in cwe_cwe:\n",
    "    w_dict[link[\"_from\"]][\"internals\"].append(link[\"_to\"])\n",
    "    w_dict[link[\"_to\"]][\"internals\"].append(link[\"_from\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:01.476620Z",
     "iopub.status.busy": "2022-05-30T09:24:01.475425Z",
     "iopub.status.idle": "2022-05-30T09:24:01.479289Z",
     "shell.execute_reply": "2022-05-30T09:24:01.478406Z"
    }
   },
   "outputs": [],
   "source": [
    "for technique in technique_dict:\n",
    "    technique_dict[technique][\"internals\"] = []\n",
    "\n",
    "for link in technique_technique:\n",
    "    technique_dict[link[\"_from\"]][\"internals\"].append(link[\"_to\"])\n",
    "    technique_dict[link[\"_to\"]][\"internals\"].append(link[\"_from\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:01.485494Z",
     "iopub.status.busy": "2022-05-30T09:24:01.484422Z",
     "iopub.status.idle": "2022-05-30T09:24:01.488322Z",
     "shell.execute_reply": "2022-05-30T09:24:01.487366Z"
    }
   },
   "outputs": [],
   "source": [
    "example_ids = positive_example_ids + negative_example_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:01.496063Z",
     "iopub.status.busy": "2022-05-30T09:24:01.494966Z",
     "iopub.status.idle": "2022-05-30T09:24:06.389042Z",
     "shell.execute_reply": "2022-05-30T09:24:06.389898Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.seq_relationship.bias', 'cls.seq_relationship.weight', 'cls.predictions.decoder.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.bias', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.weight']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "device = \"cpu\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\")\n",
    "pretrained_model = AutoModel.from_pretrained(\"bert-base-uncased\").to(device)\n",
    "\n",
    "model_path = \"bert_base\"\n",
    "finetuned_model = AutoModelForMaskedLM.from_pretrained(model_path).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:06.395829Z",
     "iopub.status.busy": "2022-05-30T09:24:06.394947Z",
     "iopub.status.idle": "2022-05-30T09:24:08.591505Z",
     "shell.execute_reply": "2022-05-30T09:24:08.589933Z"
    }
   },
   "outputs": [],
   "source": [
    "encode = spacy.load(\"en_core_web_lg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:08.599907Z",
     "iopub.status.busy": "2022-05-30T09:24:08.598563Z",
     "iopub.status.idle": "2022-05-30T09:24:08.601478Z",
     "shell.execute_reply": "2022-05-30T09:24:08.602908Z"
    }
   },
   "outputs": [],
   "source": [
    "results = {\n",
    "    \"1\": {},\n",
    "    \"2\": {},\n",
    "    \"3\": {},\n",
    "    \"4\": {},\n",
    "    \"5\": {},\n",
    "    \"6\": {},\n",
    "    \"7\": {},\n",
    "    \"8\": {},\n",
    "    \"9\": {},\n",
    "    \"10\": {},\n",
    "    \"11\": {},\n",
    "    \"12\": {},\n",
    "    \"13\": {},\n",
    "    \"14\": {},\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:08.617008Z",
     "iopub.status.busy": "2022-05-30T09:24:08.615729Z",
     "iopub.status.idle": "2022-05-30T09:24:09.044105Z",
     "shell.execute_reply": "2022-05-30T09:24:09.044970Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CountVectorizer()"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ap_name_vectorizer = CountVectorizer()\n",
    "ap_name_vectorizer.fit(ap_names)\n",
    "\n",
    "technique_name_vectorizer = CountVectorizer()\n",
    "technique_name_vectorizer.fit(technique_names)\n",
    "\n",
    "cwe_name_vectorizer = CountVectorizer()\n",
    "cwe_name_vectorizer.fit(cwe_names)\n",
    "\n",
    "tactic_name_vectorizer = CountVectorizer()\n",
    "tactic_name_vectorizer.fit(tactic_names)\n",
    "\n",
    "ap_mitigation_vectorizer = CountVectorizer()\n",
    "cwe_mitigation_vectorizer = CountVectorizer()\n",
    "tech_mitigation_vectorizer = CountVectorizer()\n",
    "\n",
    "ap_mitigation_vectorizer.fit(ap_mitigation_descriptions)\n",
    "cwe_mitigation_vectorizer.fit(cwe_mitigation_descriptions)\n",
    "tech_mitigation_vectorizer.fit(tech_mitigation_names)\n",
    "\n",
    "ap_detection_vectorizer = CountVectorizer()\n",
    "cwe_detection_vectorizer = CountVectorizer()\n",
    "tech_detection_vectorizer = CountVectorizer()\n",
    "\n",
    "ap_detection_vectorizer.fit(ap_detection_descriptions)\n",
    "cwe_detection_vectorizer.fit(cwe_detection_descriptions)\n",
    "tech_detection_vectorizer.fit(tech_detection_names)\n",
    "\n",
    "ap_description_vectorizer = CountVectorizer()\n",
    "ap_description_vectorizer.fit(ap_descriptions)\n",
    "\n",
    "cwe_description_vectorizer = CountVectorizer()\n",
    "cwe_description_vectorizer.fit(cwe_descriptions)\n",
    "\n",
    "technique_description_vectorizer = CountVectorizer()\n",
    "technique_description_vectorizer.fit(technique_descriptions)\n",
    "\n",
    "tactic_description_vectorizer = CountVectorizer()\n",
    "tactic_description_vectorizer.fit(tactic_descriptions)\n",
    "\n",
    "ap_short_description_vectorizer = CountVectorizer()\n",
    "ap_short_description_vectorizer.fit(ap_short_descriptions)\n",
    "\n",
    "cwe_short_description_vectorizer = CountVectorizer()\n",
    "cwe_short_description_vectorizer.fit(cwe_short_descriptions)\n",
    "\n",
    "technique_short_description_vectorizer = CountVectorizer()\n",
    "technique_short_description_vectorizer.fit(technique_short_descriptions)\n",
    "\n",
    "tactic_short_description_vectorizer = CountVectorizer()\n",
    "tactic_short_description_vectorizer.fit(tactic_short_descriptions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:09.056612Z",
     "iopub.status.busy": "2022-05-30T09:24:09.055122Z",
     "iopub.status.idle": "2022-05-30T09:24:09.061898Z",
     "shell.execute_reply": "2022-05-30T09:24:09.060875Z"
    }
   },
   "outputs": [],
   "source": [
    "cwe_mitigations_dict = {}\n",
    "ap_mitigations_dict = {}\n",
    "technique_mitigations_dict = {}\n",
    "\n",
    "for cwe_mit in w_mitigation:\n",
    "    cwe_mitigations_dict[cwe_mit[\"_id\"]] = cwe_mit[\"metadata\"][\"Description\"]\n",
    "\n",
    "for ap_mit in ap_mitigation:\n",
    "    ap_mitigations_dict[ap_mit[\"_id\"]] = ap_mit[\"metadata\"]\n",
    "\n",
    "for tech_mit in technique_mitigation:\n",
    "    technique_mitigations_dict[tech_mit[\"_id\"]] = tech_mit[\"name\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:09.072956Z",
     "iopub.status.busy": "2022-05-30T09:24:09.071944Z",
     "iopub.status.idle": "2022-05-30T09:24:09.077757Z",
     "shell.execute_reply": "2022-05-30T09:24:09.076891Z"
    }
   },
   "outputs": [],
   "source": [
    "cwe_detections_dict = {}\n",
    "ap_detections_dict = {}\n",
    "technique_detections_dict = {}\n",
    "\n",
    "for cwe_det in w_detection:\n",
    "    cwe_detections_dict[cwe_det[\"_id\"]] = cwe_det[\"metadata\"][\"Description\"]\n",
    "\n",
    "for ap_det in ap_detection:\n",
    "    ap_detections_dict[ap_det[\"_id\"]] = ap_det[\"metadata\"]\n",
    "\n",
    "for tech_det in technique_detection:\n",
    "    technique_detections_dict[tech_det[\"_id\"]] = tech_det[\"name\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:09.121703Z",
     "iopub.status.busy": "2022-05-30T09:24:09.083225Z",
     "iopub.status.idle": "2022-05-30T09:24:09.187541Z",
     "shell.execute_reply": "2022-05-30T09:24:09.186107Z"
    }
   },
   "outputs": [],
   "source": [
    "def vector_encoding(\n",
    "    encoding_type, text, vectorizer=None, bert_output_type=None, bert_finetuned=False\n",
    "):\n",
    "    if encoding_type == \"None\":\n",
    "        return text\n",
    "    elif encoding_type == \"BoW\":\n",
    "        return vectorizer_transform(text, vectorizer)\n",
    "    elif encoding_type == \"spaCy\":\n",
    "        return spaCy_vector(text)\n",
    "    elif encoding_type == \"BERT\":\n",
    "        if bert_finetuned:\n",
    "            model = finetuned_model\n",
    "        else:\n",
    "            model = pretrained_model\n",
    "\n",
    "        if bert_output_type == \"pooler_output\":\n",
    "            return get_pooler_output(model, text)\n",
    "        elif bert_output_type == \"hidden_state\":\n",
    "            return get_hidden_state(model, text)\n",
    "\n",
    "\n",
    "def vectorizer_transform(input_to_BoW, vectorizer):\n",
    "    return vectorizer.transform([input_to_BoW])[0].toarray().flatten()\n",
    "\n",
    "\n",
    "def spaCy_vector(text):\n",
    "    return encode(text).vector\n",
    "\n",
    "\n",
    "def get_pooler_output(model, text):\n",
    "    inputs = tokenizer(text.lower(), truncation=True, return_tensors=\"pt\").to(device)\n",
    "    outputs = model(**inputs)\n",
    "    pooled_output = outputs.pooler_output\n",
    "    return pooled_output.detach().cpu().numpy().flatten()\n",
    "\n",
    "\n",
    "def get_hidden_state(model, text):\n",
    "    inputs = tokenizer(text.lower(), truncation=True, return_tensors=\"pt\").to(device)\n",
    "    outputs = model(**inputs, output_hidden_states=True)\n",
    "    hidden_states = outputs.hidden_states\n",
    "    return hidden_states[-1][:, 0, :].detach().cpu().numpy().flatten()\n",
    "\n",
    "\n",
    "def append_data(\n",
    "    encoding_type,\n",
    "    data_combo,\n",
    "    ap,\n",
    "    technique,\n",
    "    bert_output_type=None,\n",
    "    bert_finetuned=False,\n",
    "):\n",
    "    output = []\n",
    "    vectorizer = CountVectorizer()\n",
    "\n",
    "    ap_text = ap_names\n",
    "    technique_text = technique_names\n",
    "    cwe_text = cwe_names\n",
    "    tactic_text = tactic_names\n",
    "\n",
    "    if data_combo == \"A0\":\n",
    "        vectorizer.fit(ap_text + technique_text)\n",
    "\n",
    "    elif data_combo == \"A1\":\n",
    "        vectorizer.fit(ap_text + technique_text + cwe_text + tactic_text)\n",
    "\n",
    "    elif data_combo == \"A1 + MI\":\n",
    "        vectorizer.fit(\n",
    "            ap_text\n",
    "            + technique_text\n",
    "            + cwe_text\n",
    "            + tactic_text\n",
    "            + cwe_mitigation_descriptions\n",
    "            + ap_mitigation_descriptions\n",
    "            + tech_mitigation_names\n",
    "        )\n",
    "\n",
    "    elif data_combo == \"A1 + D\":\n",
    "        vectorizer.fit(\n",
    "            ap_text\n",
    "            + technique_text\n",
    "            + cwe_text\n",
    "            + tactic_text\n",
    "            + cwe_detection_descriptions\n",
    "            + ap_detection_descriptions\n",
    "            + tech_detection_names\n",
    "        )\n",
    "\n",
    "    elif data_combo == \"A1 + MI + D\":\n",
    "        vectorizer.fit(\n",
    "            ap_text\n",
    "            + technique_text\n",
    "            + cwe_text\n",
    "            + tactic_text\n",
    "            + cwe_mitigation_descriptions\n",
    "            + ap_mitigation_descriptions\n",
    "            + tech_mitigation_names\n",
    "            + cwe_detection_descriptions\n",
    "            + ap_detection_descriptions\n",
    "            + tech_detection_names\n",
    "        )\n",
    "\n",
    "    output.append(ap_dict[ap][\"name\"])\n",
    "    output.append(technique_dict[technique][\"name\"])\n",
    "\n",
    "    for internal in ap_dict[ap][\"internals\"]:\n",
    "        output.append(ap_dict[internal][\"name\"])\n",
    "\n",
    "    for internal in technique_dict[technique][\"internals\"]:\n",
    "        output.append(technique_dict[internal][\"name\"])\n",
    "\n",
    "    if \"A1\" in data_combo:\n",
    "        for cwe in ap_dict[ap][\"cwes\"]:\n",
    "            output.append(w_dict[cwe][\"name\"])\n",
    "            for w in w_dict[cwe][\"internals\"]:\n",
    "                output.append(w_dict[w][\"name\"])\n",
    "\n",
    "        for tac in technique_dict[technique][\"tactics\"]:\n",
    "            output.append(tactic_dict[tac][\"name\"])\n",
    "\n",
    "    if data_combo in [\"A1 + MI\", \"A1 + MI + D\"]:\n",
    "        for cwe in ap_dict[ap][\"cwes\"]:\n",
    "            for mitigation in w_dict[cwe][\"mitigations\"]:\n",
    "                output.append(cwe_mitigations_dict[mitigation])\n",
    "            for w in w_dict[cwe][\"internals\"]:\n",
    "                for mitigation in w_dict[w][\"mitigations\"]:\n",
    "                    output.append(cwe_mitigations_dict[mitigation])\n",
    "\n",
    "        for mitigation in ap_dict[ap][\"mitigations\"]:\n",
    "            output.append(ap_mitigations_dict[mitigation])\n",
    "        for internal in ap_dict[ap][\"internals\"]:\n",
    "            for mitigation in ap_dict[internal][\"mitigations\"]:\n",
    "                output.append(ap_mitigations_dict[mitigation])\n",
    "\n",
    "        for mitigation in technique_dict[technique][\"mitigations\"]:\n",
    "            output.append(technique_mitigations_dict[mitigation])\n",
    "        for internal in technique_dict[technique][\"internals\"]:\n",
    "            for mitigation in technique_dict[internal][\"mitigations\"]:\n",
    "                output.append(technique_mitigations_dict[mitigation])\n",
    "\n",
    "    if data_combo in [\"A1 + D\", \"A1 + MI + D\"]:\n",
    "        for cwe in ap_dict[ap][\"cwes\"]:\n",
    "            for detection in w_dict[cwe][\"detections\"]:\n",
    "                output.append(cwe_detections_dict[detection])\n",
    "            for w in w_dict[cwe][\"internals\"]:\n",
    "                for detection in w_dict[w][\"detections\"]:\n",
    "                    output.append(cwe_detections_dict[detection])\n",
    "\n",
    "        for detection in ap_dict[ap][\"detections\"]:\n",
    "            output.append(ap_detections_dict[detection])\n",
    "        for internal in ap_dict[ap][\"internals\"]:\n",
    "            for detection in ap_dict[internal][\"detections\"]:\n",
    "                output.append(ap_detections_dict[detection])\n",
    "\n",
    "        for detection in technique_dict[technique][\"detections\"]:\n",
    "            output.append(technique_detections_dict[detection])\n",
    "        for internal in technique_dict[technique][\"internals\"]:\n",
    "            for detection in technique_dict[internal][\"detections\"]:\n",
    "                output.append(technique_detections_dict[detection])\n",
    "\n",
    "    output = \" \".join(output)\n",
    "    return vector_encoding(\n",
    "        encoding_type, output, vectorizer, bert_output_type, bert_finetuned\n",
    "    )\n",
    "\n",
    "\n",
    "def handle_data(\n",
    "    encoding_type,\n",
    "    data_combo,\n",
    "    ap,\n",
    "    technique,\n",
    "    bert_output_type=None,\n",
    "    bert_finetuned=False,\n",
    "):\n",
    "    example = []\n",
    "\n",
    "    aps = []\n",
    "    aps.append(ap_dict[ap][\"name\"])\n",
    "    for internal in ap_dict[ap][\"internals\"]:\n",
    "        aps.append(ap_dict[internal][\"name\"])\n",
    "    aps = \" \".join(aps)\n",
    "    example.append(\n",
    "        vector_encoding(\n",
    "            encoding_type, aps, ap_name_vectorizer, bert_output_type, bert_finetuned\n",
    "        )\n",
    "    )\n",
    "\n",
    "    techniques = []\n",
    "    techniques.append(technique_dict[technique][\"name\"])\n",
    "    for internal in technique_dict[technique][\"internals\"]:\n",
    "        techniques.append(technique_dict[internal][\"name\"])\n",
    "    techniques = \" \".join(techniques)\n",
    "    example.append(\n",
    "        vector_encoding(\n",
    "            encoding_type,\n",
    "            techniques,\n",
    "            technique_name_vectorizer,\n",
    "            bert_output_type,\n",
    "            bert_finetuned,\n",
    "        )\n",
    "    )\n",
    "\n",
    "    if \"A1\" in data_combo:\n",
    "        tactics = []\n",
    "        for tac in technique_dict[technique][\"tactics\"]:\n",
    "            tactics.append(tactic_dict[tac][\"name\"])\n",
    "        tactics = \" \".join(tactics)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                tactics,\n",
    "                tactic_name_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        cwes = []\n",
    "        for cwe in ap_dict[ap][\"cwes\"]:\n",
    "            cwes.append(w_dict[cwe][\"name\"])\n",
    "            for w in w_dict[cwe][\"internals\"]:\n",
    "                cwes.append(w_dict[w][\"name\"])\n",
    "        cwes = \" \".join(cwes)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                cwes,\n",
    "                cwe_name_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "    if data_combo in [\"A1 + MI\", \"A1 + MI + D\"]:\n",
    "        cwe_mitigations = []\n",
    "        for cwe in ap_dict[ap][\"cwes\"]:\n",
    "            for mitigation in w_dict[cwe][\"mitigations\"]:\n",
    "                cwe_mitigations.append(cwe_mitigations_dict[mitigation])\n",
    "            for w in w_dict[cwe][\"internals\"]:\n",
    "                for mitigation in w_dict[w][\"mitigations\"]:\n",
    "                    cwe_mitigations.append(cwe_mitigations_dict[mitigation])\n",
    "        cwe_mitigations = \" \".join(cwe_mitigations)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                cwe_mitigations,\n",
    "                cwe_mitigation_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        capec_mitigations = []\n",
    "        for mitigation in ap_dict[ap][\"mitigations\"]:\n",
    "            capec_mitigations.append(ap_mitigations_dict[mitigation])\n",
    "        for internal in ap_dict[ap][\"internals\"]:\n",
    "            for mitigation in ap_dict[internal][\"mitigations\"]:\n",
    "                capec_mitigations.append(ap_mitigations_dict[mitigation])\n",
    "        capec_mitigations = \" \".join(capec_mitigations)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                capec_mitigations,\n",
    "                ap_mitigation_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        tech_mitigations = []\n",
    "        for mitigation in technique_dict[technique][\"mitigations\"]:\n",
    "            tech_mitigations.append(technique_mitigations_dict[mitigation])\n",
    "        for internal in technique_dict[technique][\"internals\"]:\n",
    "            for mitigation in technique_dict[internal][\"mitigations\"]:\n",
    "                tech_mitigations.append(technique_mitigations_dict[mitigation])\n",
    "        tech_mitigations = \" \".join(tech_mitigations)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                tech_mitigations,\n",
    "                tech_mitigation_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "    if data_combo in [\"A1 + D\", \"A1 + MI + D\"]:\n",
    "        cwe_detections = []\n",
    "        for cwe in ap_dict[ap][\"cwes\"]:\n",
    "            for detection in w_dict[cwe][\"detections\"]:\n",
    "                cwe_detections.append(cwe_detections_dict[detection])\n",
    "            for w in w_dict[cwe][\"internals\"]:\n",
    "                for detection in w_dict[w][\"detections\"]:\n",
    "                    cwe_detections.append(cwe_detections_dict[detection])\n",
    "\n",
    "        cwe_detections = \" \".join(cwe_detections)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                cwe_detections,\n",
    "                cwe_detection_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        capec_detections = []\n",
    "        for detection in ap_dict[ap][\"detections\"]:\n",
    "            capec_detections.append(ap_detections_dict[detection])\n",
    "        for internal in ap_dict[ap][\"internals\"]:\n",
    "            for detection in ap_dict[internal][\"detections\"]:\n",
    "                capec_detections.append(ap_detections_dict[detection])\n",
    "        capec_detections = \" \".join(capec_detections)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                capec_detections,\n",
    "                ap_detection_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        tech_detections = []\n",
    "        for detection in technique_dict[technique][\"detections\"]:\n",
    "            tech_detections.append(technique_detections_dict[detection])\n",
    "        for internal in technique_dict[technique][\"internals\"]:\n",
    "            for detection in technique_dict[internal][\"detections\"]:\n",
    "                tech_detections.append(technique_detections_dict[detection])\n",
    "        tech_detections = \" \".join(tech_detections)\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                tech_detections,\n",
    "                tech_detection_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "    return np.hstack(example)\n",
    "\n",
    "\n",
    "def encode_data(\n",
    "    encoding_type,\n",
    "    data_combo,\n",
    "    ap,\n",
    "    technique,\n",
    "    bert_output_type=None,\n",
    "    bert_finetuned=False,\n",
    "):\n",
    "    example = []\n",
    "    example.append(\n",
    "        vector_encoding(\n",
    "            encoding_type,\n",
    "            ap_dict[ap][\"name\"],\n",
    "            ap_name_vectorizer,\n",
    "            bert_output_type,\n",
    "            bert_finetuned,\n",
    "        )\n",
    "    )\n",
    "    for internal in ap_dict[ap][\"internals\"]:\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                ap_dict[internal][\"name\"],\n",
    "                ap_name_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "    example.append(\n",
    "        vector_encoding(\n",
    "            encoding_type,\n",
    "            technique_dict[technique][\"name\"],\n",
    "            technique_name_vectorizer,\n",
    "            bert_output_type,\n",
    "            bert_finetuned,\n",
    "        )\n",
    "    )\n",
    "    for internal in technique_dict[technique][\"internals\"]:\n",
    "        example.append(\n",
    "            vector_encoding(\n",
    "                encoding_type,\n",
    "                technique_dict[internal][\"name\"],\n",
    "                technique_name_vectorizer,\n",
    "                bert_output_type,\n",
    "                bert_finetuned,\n",
    "            )\n",
    "        )\n",
    "\n",
    "    if \"A1\" in data_combo:\n",
    "        for cwe in ap_dict[ap][\"cwes\"]:\n",
    "            example.append(\n",
    "                vector_encoding(\n",
    "                    encoding_type,\n",
    "                    w_dict[cwe][\"name\"],\n",
    "                    cwe_name_vectorizer,\n",
    "                    bert_output_type,\n",
    "                    bert_finetuned,\n",
    "                )\n",
    "            )\n",
    "            for internal in w_dict[cwe][\"internals\"]:\n",
    "                example.append(\n",
    "                    vector_encoding(\n",
    "                        encoding_type,\n",
    "                        w_dict[internal][\"name\"],\n",
    "                        cwe_name_vectorizer,\n",
    "                        bert_output_type,\n",
    "                        bert_finetuned,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "        for tac in technique_dict[technique][\"tactics\"]:\n",
    "            example.append(\n",
    "                vector_encoding(\n",
    "                    encoding_type,\n",
    "                    tactic_dict[tac][\"name\"],\n",
    "                    tactic_name_vectorizer,\n",
    "                    bert_output_type,\n",
    "                    bert_finetuned,\n",
    "                )\n",
    "            )\n",
    "\n",
    "    if data_combo in [\"A1 + MI\", \"A1 + MI + D\"]:\n",
    "        for cwe in ap_dict[ap][\"cwes\"]:\n",
    "            for mitigation in w_dict[cwe][\"mitigations\"]:\n",
    "                example.append(\n",
    "                    vector_encoding(\n",
    "                        encoding_type,\n",
    "                        cwe_mitigations_dict[mitigation],\n",
    "                        cwe_mitigation_vectorizer,\n",
    "                        bert_output_type,\n",
    "                        bert_finetuned,\n",
    "                    )\n",
    "                )\n",
    "            for internal in w_dict[cwe][\"internals\"]:\n",
    "                for mitigation in w_dict[internal][\"mitigations\"]:\n",
    "                    example.append(\n",
    "                        vector_encoding(\n",
    "                            encoding_type,\n",
    "                            cwe_mitigations_dict[mitigation],\n",
    "                            cwe_mitigation_vectorizer,\n",
    "                            bert_output_type,\n",
    "                            bert_finetuned,\n",
    "                        )\n",
    "                    )\n",
    "\n",
    "        for mitigation in ap_dict[ap][\"mitigations\"]:\n",
    "            example.append(\n",
    "                vector_encoding(\n",
    "                    encoding_type,\n",
    "                    ap_mitigations_dict[mitigation],\n",
    "                    ap_mitigation_vectorizer,\n",
    "                    bert_output_type,\n",
    "                    bert_finetuned,\n",
    "                )\n",
    "            )\n",
    "        for internal in ap_dict[ap][\"internals\"]:\n",
    "            for mitigation in ap_dict[internal][\"mitigations\"]:\n",
    "                example.append(\n",
    "                    vector_encoding(\n",
    "                        encoding_type,\n",
    "                        ap_mitigations_dict[mitigation],\n",
    "                        ap_mitigation_vectorizer,\n",
    "                        bert_output_type,\n",
    "                        bert_finetuned,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "        for mitigation in technique_dict[technique][\"mitigations\"]:\n",
    "            example.append(\n",
    "                vector_encoding(\n",
    "                    encoding_type,\n",
    "                    technique_mitigations_dict[mitigation],\n",
    "                    tech_mitigation_vectorizer,\n",
    "                    bert_output_type,\n",
    "                    bert_finetuned,\n",
    "                )\n",
    "            )\n",
    "        for internal in technique_dict[technique][\"internals\"]:\n",
    "            for mitigation in technique_dict[internal][\"mitigations\"]:\n",
    "                example.append(\n",
    "                    vector_encoding(\n",
    "                        encoding_type,\n",
    "                        technique_mitigations_dict[mitigation],\n",
    "                        tech_mitigation_vectorizer,\n",
    "                        bert_output_type,\n",
    "                        bert_finetuned,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "    if data_combo in [\"A1 + D\", \"A1 + MI + D\"]:\n",
    "        for cwe in ap_dict[ap][\"cwes\"]:\n",
    "            for detection in w_dict[cwe][\"detections\"]:\n",
    "                example.append(\n",
    "                    vector_encoding(\n",
    "                        encoding_type,\n",
    "                        cwe_detections_dict[detection],\n",
    "                        cwe_detection_vectorizer,\n",
    "                        bert_output_type,\n",
    "                        bert_finetuned,\n",
    "                    )\n",
    "                )\n",
    "            for internal in w_dict[cwe][\"internals\"]:\n",
    "                for detection in w_dict[internal][\"detections\"]:\n",
    "                    example.append(\n",
    "                        vector_encoding(\n",
    "                            encoding_type,\n",
    "                            cwe_detections_dict[detection],\n",
    "                            cwe_detection_vectorizer,\n",
    "                            bert_output_type,\n",
    "                            bert_finetuned,\n",
    "                        )\n",
    "                    )\n",
    "\n",
    "        for detection in ap_dict[ap][\"detections\"]:\n",
    "            example.append(\n",
    "                vector_encoding(\n",
    "                    encoding_type,\n",
    "                    ap_detections_dict[detection],\n",
    "                    ap_detection_vectorizer,\n",
    "                    bert_output_type,\n",
    "                    bert_finetuned,\n",
    "                )\n",
    "            )\n",
    "        for internal in ap_dict[ap][\"internals\"]:\n",
    "            for detection in ap_dict[internal][\"detections\"]:\n",
    "                example.append(\n",
    "                    vector_encoding(\n",
    "                        encoding_type,\n",
    "                        ap_detections_dict[detection],\n",
    "                        ap_detection_vectorizer,\n",
    "                        bert_output_type,\n",
    "                        bert_finetuned,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "        for detection in technique_dict[technique][\"detections\"]:\n",
    "            example.append(\n",
    "                vector_encoding(\n",
    "                    encoding_type,\n",
    "                    technique_detections_dict[detection],\n",
    "                    tech_detection_vectorizer,\n",
    "                    bert_output_type,\n",
    "                    bert_finetuned,\n",
    "                )\n",
    "            )\n",
    "        for internal in technique_dict[technique][\"internals\"]:\n",
    "            for detection in technique_dict[internal][\"detections\"]:\n",
    "                example.append(\n",
    "                    vector_encoding(\n",
    "                        encoding_type,\n",
    "                        technique_detections_dict[detection],\n",
    "                        tech_detection_vectorizer,\n",
    "                        bert_output_type,\n",
    "                        bert_finetuned,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "    return np.hstack(example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:24:09.222913Z",
     "iopub.status.busy": "2022-05-30T09:24:09.208248Z",
     "iopub.status.idle": "2022-05-30T09:45:17.803475Z",
     "shell.execute_reply": "2022-05-30T09:45:17.802036Z"
    }
   },
   "outputs": [],
   "source": [
    "encodings_to_use = [6]\n",
    "\n",
    "max_length_BoW = 0\n",
    "max_length_spaCy = 0\n",
    "max_length_BERT = 0\n",
    "for data_combo in [\"A0\", \"A1\", \"A1 + MI\", \"A1 + D\", \"A1 + MI + D\"]:\n",
    "    for encoding in encodings_to_use:\n",
    "        accuracies = []\n",
    "        aucs = []\n",
    "        f1s = []\n",
    "        fps = []\n",
    "        labels = []\n",
    "        examples = []\n",
    "        for i, (ap, technique) in enumerate(example_ids):\n",
    "            if encoding == 1:\n",
    "                examples.append(append_data(\"BoW\", data_combo, ap, technique))\n",
    "\n",
    "            elif encoding == 2:\n",
    "                examples.append(handle_data(\"BoW\", data_combo, ap, technique))\n",
    "\n",
    "            elif encoding == 3:\n",
    "                examples.append(encode_data(\"BoW\", data_combo, ap, technique))\n",
    "                max_length_BoW = max(max_length_BoW, len(examples[-1]))\n",
    "\n",
    "            elif encoding == 4:\n",
    "                examples.append(append_data(\"spaCy\", data_combo, ap, technique))\n",
    "\n",
    "            elif encoding == 5:\n",
    "                examples.append(handle_data(\"spaCy\", data_combo, ap, technique))\n",
    "\n",
    "            elif encoding == 6:\n",
    "                examples.append(encode_data(\"spaCy\", data_combo, ap, technique))\n",
    "                max_length_spaCy = max(max_length_spaCy, len(examples[-1]))\n",
    "\n",
    "            elif encoding == 7:\n",
    "                examples.append(\n",
    "                    encode_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        technique,\n",
    "                        bert_output_type=\"pooler_output\",\n",
    "                        bert_finetuned=False,\n",
    "                    )\n",
    "                )\n",
    "                max_length_BERT = max(max_length_BERT, len(examples[-1]))\n",
    "\n",
    "            elif encoding == 8:\n",
    "                examples.append(\n",
    "                    append_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        technique,\n",
    "                        bert_output_type=\"pooler_output\",\n",
    "                        bert_finetuned=False,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "            elif encoding == 9:\n",
    "                examples.append(\n",
    "                    encode_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        technique,\n",
    "                        bert_output_type=\"hidden_state\",\n",
    "                        bert_finetuned=False,\n",
    "                    )\n",
    "                )\n",
    "                max_length_BERT = max(max_length_BERT, len(examples[-1]))\n",
    "\n",
    "            elif encoding == 10:\n",
    "                examples.append(\n",
    "                    append_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        technique,\n",
    "                        bert_output_type=\"hidden_state\",\n",
    "                        bert_finetuned=False,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "            elif encoding == 11:\n",
    "                examples.append(\n",
    "                    encode_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        technique,\n",
    "                        bert_output_type=\"hidden_state\",\n",
    "                        bert_finetuned=True,\n",
    "                    )\n",
    "                )\n",
    "                max_length_BERT = max(max_length_BERT, len(examples[-1]))\n",
    "\n",
    "            elif encoding == 12:\n",
    "                examples.append(\n",
    "                    append_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        technique,\n",
    "                        bert_output_type=\"hidden_state\",\n",
    "                        bert_finetuned=True,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "            elif encoding == 13:\n",
    "                examples.append(\n",
    "                    handle_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        technique,\n",
    "                        bert_output_type=\"hidden_state\",\n",
    "                        bert_finetuned=False,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "            elif encoding == 14:\n",
    "                examples.append(\n",
    "                    handle_data(\n",
    "                        \"BERT\",\n",
    "                        data_combo,\n",
    "                        ap,\n",
    "                        technique,\n",
    "                        bert_output_type=\"hidden_state\",\n",
    "                        bert_finetuned=True,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "            if i < len(example_ids) / 2:\n",
    "                labels.append(1)\n",
    "            else:\n",
    "                labels.append(0)\n",
    "\n",
    "        if encoding == 3:\n",
    "            for i in range(len(examples)):\n",
    "                examples[i] = np.pad(\n",
    "                    examples[i],\n",
    "                    pad_width=(0, max_length_BoW - len(examples[i])),\n",
    "                    mode=\"constant\",\n",
    "                )\n",
    "\n",
    "        if encoding == 6:\n",
    "            for i in range(len(examples)):\n",
    "                examples[i] = np.pad(\n",
    "                    examples[i],\n",
    "                    pad_width=(0, max_length_spaCy - len(examples[i])),\n",
    "                    mode=\"constant\",\n",
    "                )\n",
    "\n",
    "        if encoding == 7 or encoding == 9 or encoding == 11:\n",
    "            for i in range(len(examples)):\n",
    "                examples[i] = np.pad(\n",
    "                    examples[i],\n",
    "                    pad_width=(0, max_length_BERT - len(examples[i])),\n",
    "                    mode=\"constant\",\n",
    "                )\n",
    "\n",
    "        for i in range(100):\n",
    "            # training/testing classifier\n",
    "            X_train, X_test, y_train, y_test = train_test_split(\n",
    "                examples, labels, test_size=0.3, random_state=i\n",
    "            )\n",
    "            clf = RandomForestClassifier(random_state=i, class_weight={0: 1, 1: 1})\n",
    "            clf.fit(X_train, y_train)\n",
    "            y_pred = clf.predict(X_test)\n",
    "\n",
    "            confusion = confusion_matrix(y_test, y_pred)\n",
    "            fp_rate = confusion[0][1] / (confusion[0][1] + confusion[0][0])\n",
    "            fps.append(fp_rate)\n",
    "\n",
    "            accuracies.append(accuracy_score(y_test, y_pred))\n",
    "            aucs.append(roc_auc_score(y_test, clf.predict_proba(X_test)[:, 1]))\n",
    "            f1s.append(f1_score(y_test, y_pred))\n",
    "\n",
    "        results[str(encoding)][data_combo] = {\n",
    "            \"fp\": fps,\n",
    "            \"acc\": accuracies,\n",
    "            \"auc\": aucs,\n",
    "            \"f1\": f1s,\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-30T09:45:17.818090Z",
     "iopub.status.busy": "2022-05-30T09:45:17.816906Z",
     "iopub.status.idle": "2022-05-30T09:45:17.822346Z",
     "shell.execute_reply": "2022-05-30T09:45:17.821301Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(\"results_100_with_spaCy_one_step_internal_links.json\", \"w\") as f:\n",
    "    json.dump(results, f)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "98b0a9b7b4eaaa670588a142fd0a9b87eaafe866f1db4228be72b4211d12040f"
  },
  "kernelspec": {
   "display_name": "Python 3.7.10 64-bit ('base': conda)",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}