add script for validation/test split for reproducibility

1 files changed, 73 insertions, 0 deletions

diff --git a/notebooks/validation_test_split.ipynb b/notebooks/validation_test_split.ipynb
new file mode 100644
index 00000000..1ee335d7
--- /dev/null
+++ b/notebooks/validation_test_split.ipynb
@@ -0,0 +1,73 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "source": [
+    "from sklearn.model_selection import train_test_split\n",
+    "from sec_certs.dataset.common_criteria import CCDataset\n",
+    "import json"
+   ],
+   "outputs": [],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "source": [
+    "cc_dset: CCDataset = CCDataset.from_json('/Users/adam/phd/projects/certificates/datasets/cc_full_dataset/CommonCriteria_dataset.json')\n",
+    "\n",
+    "cve_rich_certs = [x for x in cc_dset if x.heuristics.related_cves]\n",
+    "cve_free_certs = [x for x in cc_dset if not x.heuristics.related_cves]\n",
+    "\n",
+    "x_valid_cve_rich, x_test_cve_rich = train_test_split(cve_rich_certs, test_size=0.5)\n",
+    "x_valid_cve_free, x_test_cve_free = train_test_split(cve_free_certs, test_size=0.5)\n",
+    "\n",
+    "validation_set = x_valid_cve_rich + x_valid_cve_free\n",
+    "test_set = x_test_cve_rich + x_test_cve_free\n",
+    "\n",
+    "validation_set = {x.dgst: x.name for x in validation_set}\n",
+    "test_set = {x.dgst: x.name for x in test_set}"
+   ],
+   "outputs": [],
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "source": [
+    "with open('./validation_set.json', 'w') as handle:\n",
+    "    json.dump(validation_set, handle, indent=4)\n",
+    "\n",
+    "with open('./test_set.json', 'w') as handle:\n",
+    "    json.dump(test_set, handle, indent=4)"
+   ],
+   "outputs": [],
+   "metadata": {}
+  }
+ ],
+ "metadata": {
+  "orig_nbformat": 4,
+  "language_info": {
+   "name": "python",
+   "version": "3.8.1",
+   "mimetype": "text/x-python",
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "pygments_lexer": "ipython3",
+   "nbconvert_exporter": "python",
+   "file_extension": ".py"
+  },
+  "kernelspec": {
+   "name": "python3",
+   "display_name": "Python 3.8.1 64-bit ('3.8.1': pyenv)"
+  },
+  "interpreter": {
+   "hash": "9c2c998cb2293d5512c6ce922cbfe08310964e398b5d5fc2b51147c001724304"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
\ No newline at end of file