2_Baseline_Comparison.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Baseline Comparison & Evaluation\n",
    "\n",
    "This notebook recreates the results of the baseline comparison and evaluation in the paper.\n",
    "\n",
    "The previous notebook must be ran first to generate the necessary corpora, etc., for this notebook to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "import collections as cll\n",
    "import datetime as dt\n",
    "import gc\n",
    "import gzip\n",
    "import itertools as it\n",
    "import multiprocessing as mp\n",
    "import multiprocessing.pool as mp_pool\n",
    "import os\n",
    "import pickle\n",
    "import random\n",
    "import sys\n",
    "import typing as T\n",
    "\n",
    "import numpy as np\n",
    "import numpy.typing as T_np\n",
    "\n",
    "import sklearn.metrics as metrics\n",
    "import sklearn.utils\n",
    "\n",
    "from sklearn.preprocessing import MultiLabelBinarizer\n",
    "from iterstrat.ml_stratifiers import MultilabelStratifiedKFold"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Global notebook options and variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "# The path containing the corpora and other artefacts from the previous Jupyter\n",
    "# notebook.\n",
    "ARTEFACTS_PATH = \"./Artefacts/\"\n",
    "\n",
    "# This is the seed we used for our experiments. If you wish to change this,\n",
    "# then be sure to delete any cached results you may have already generated.\n",
    "# It will have no serious impact on the results at all.\n",
    "SEED = 463\n",
    "\n",
    "K_FOLDS = 5  # This controls the number of folds used in the cross-validation.\n",
    "K_FOLD_LIMIT = K_FOLDS  # A debug option to investigate one fold or few folds.\n",
    "TOP_CLASSES: T.Optional[int] = 10  # Top-supported classes outputted (summary).\n",
    "AVERAGING_METHODS: T.Tuple[  # The paper used both `macro` and `micro`.\n",
    "    T.Literal[\"macro\", \"micro\", \"weighted\", \"samples\"], ...\n",
    "] = (\"macro\", \"micro\", \"weighted\", \"samples\")\n",
    "\n",
    "# Useful for debugging the pipelines with a small slice of data.\n",
    "SAMPLE_LIMIT: T.Optional[slice] = None\n",
    "\n",
    "# Our experiments were conducted with the majority class trimmed to the size of\n",
    "# the second largest class. This can be disabled, but it is most useful when\n",
    "# the majority class dwarfs all other classes.\n",
    "TRIM_MAJORITY_CLASS = True\n",
    "\n",
    "# While we tried numerous classifiers, we reported k-NN results due to its\n",
    "# simplicity and acceptable performance for how little hyperparameter tuning it\n",
    "# requires. That said, we found XGBoost to outperform k-NN by a small but\n",
    "# noticeable margin, and excluded it from the paper due to space constraints.\n",
    "# See the `build_classifier` function for hyperparameter details.\n",
    "CLASSIFIER_NAME: T.Literal[\n",
    "    \"k-NN\",  # 10-NN with distance weighting in the paper.\n",
    "    \"XGBoost\",\n",
    "    \"Dummy\",  # Not to be confused with the \"Dummy\" approach in the paper.\n",
    "] = \"k-NN\"\n",
    "\n",
    "# In case a crash happens, we don't need to redo a fold as long as the seed is\n",
    "# identical.\n",
    "USE_LAST_CACHED_SUMMARY = True\n",
    "\n",
    "# Type aliases.\n",
    "Score = T.Dict[str, T.Optional[float]]\n",
    "ClassScores = T.Dict[str, Score]\n",
    "Scores = T.Dict[str, T.Union[T.Optional[float], Score, ClassScores]]\n",
    "Summary = T.Dict[str, T.Union[str, int, dt.datetime, Scores]]\n",
    "\n",
    "# An extremely poor attempt at taking advantage of Copy-on-Write behaviour in\n",
    "# GNU/Linux, as multiprocessing is used; thus, `fork()` is involved. Reference\n",
    "# counting often defeats this. These are populated during preparation of\n",
    "# cross-validation for a run of an approach.\n",
    "x_train_folds: T.List[T.List[T.Sequence[T.Any]]] = []\n",
    "y_train_folds: T.List[T.List[T.Sequence[T.Any]]] = []\n",
    "x_test_folds: T.List[T.List[T.Sequence[T.Any]]] = []\n",
    "y_test_folds: T.List[T.List[T.Sequence[T.Any]]] = []\n",
    "\n",
    "# Resolves the set bit of a label (our samples have multiple labels) to the\n",
    "# actual label name (IXP Scrubber filtering rule hashes in the paper's case).\n",
    "column_label_targets: T.List[str] = []"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Helper functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "# This is a hack to make the default multiprocessing pool not spawn daemon\n",
    "# processes; hence, we can then use pools within child processes made by a\n",
    "# pool. Will cause a lot of zombie processes or fail to exit cleanly if\n",
    "# something ancestral crashes, as they're no longer marked daemonic.\n",
    "class NoDaemonProcess(mp.Process):\n",
    "    @property\n",
    "    def daemon(self):\n",
    "        return False\n",
    "\n",
    "    @daemon.setter\n",
    "    def daemon(self, _):\n",
    "        pass\n",
    "\n",
    "\n",
    "class CustomPool(mp_pool.Pool):  # See above.\n",
    "    def Process(self, *args, **kwargs):\n",
    "        proc = super().Process(*args, **kwargs)  # type: ignore\n",
    "        proc.__class__ = NoDaemonProcess\n",
    "        return proc\n",
    "\n",
    "\n",
    "LOG_LOCK = mp.Lock()\n",
    "\n",
    "\n",
    "def log(*args, **kwargs):\n",
    "    with LOG_LOCK:\n",
    "        print(\n",
    "            f\"[{dt.datetime.now().strftime('%d-%m %I:%M:%S %p')}]:\",\n",
    "            *args,\n",
    "            **kwargs,\n",
    "        )\n",
    "\n",
    "\n",
    "def load_pickle(path: str, compressed=True) -> T.Any:\n",
    "    try:\n",
    "        gc.disable()\n",
    "        if compressed:\n",
    "            assert path.endswith(\".gz\")\n",
    "            with gzip.open(path, \"rb\") as file:\n",
    "                return pickle.load(file)\n",
    "        else:\n",
    "            with open(path, \"rb\") as file:\n",
    "                return pickle.load(file)\n",
    "    except Exception as exception:\n",
    "        raise exception from None\n",
    "    finally:\n",
    "        gc.enable()\n",
    "\n",
    "\n",
    "def save_pickle(object: T.Any, path: str, compress=True):\n",
    "    if compress:\n",
    "        assert path.endswith(\".gz\")\n",
    "        with gzip.open(path, \"wb\", compresslevel=6) as file:\n",
    "            pickle.dump(object, file, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "    else:\n",
    "        with open(path, \"wb\") as file:\n",
    "            pickle.dump(object, file, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "\n",
    "\n",
    "def average_scores(\n",
    "    fold_scores: T.Sequence[Scores],\n",
    "    class_labels: T.Optional[T.Union[int, T.Sequence[str]]] = None,\n",
    ") -> T.Dict[str, T.Tuple[float, float]]:\n",
    "    averaged_scores: T.DefaultDict[str, T.List[float]] = cll.defaultdict(list)\n",
    "    for average in AVERAGING_METHODS:\n",
    "        average = average.title()\n",
    "        for scores in fold_scores:\n",
    "            for metric, value in scores[\n",
    "                f\"Average ({average})\"\n",
    "            ].items():  # type: ignore\n",
    "                if isinstance(value, float):\n",
    "                    averaged_scores[f\"[{average}] {metric}\"].append(value)\n",
    "\n",
    "    if class_labels:\n",
    "        if isinstance(class_labels, int):\n",
    "            top_classes: T.DefaultDict[str, float] = cll.defaultdict(float)\n",
    "            for scores in fold_scores:\n",
    "                class_scores: ClassScores = scores[\"Classes\"]  # type: ignore\n",
    "                for class_label, metrics in class_scores.items():\n",
    "                    if \"Support\" in metrics and isinstance(\n",
    "                        metrics[\"Support\"], (int, float)\n",
    "                    ):\n",
    "                        top_classes[class_label] += float(metrics[\"Support\"])\n",
    "            class_labels = tuple(\n",
    "                class_label[0]\n",
    "                for class_label in sorted(\n",
    "                    top_classes.items(), key=lambda x: x[1], reverse=True\n",
    "                )[0:class_labels]\n",
    "            )\n",
    "\n",
    "        for scores in fold_scores:\n",
    "            class_scores: ClassScores = scores[\"Classes\"]  # type: ignore\n",
    "            for class_label in class_labels:\n",
    "                if class_label in class_scores:\n",
    "                    for metric, value in class_scores[class_label].items():\n",
    "                        if isinstance(value, float):\n",
    "                            averaged_scores[\n",
    "                                f\"|{class_label}| {metric}\"\n",
    "                            ].append(value)\n",
    "\n",
    "    return {\n",
    "        metric: (float(np.mean(values)), float(np.std(values)))\n",
    "        for metric, values in averaged_scores.items()\n",
    "    }\n",
    "\n",
    "\n",
    "def print_scores(averaged_scores: T.Dict[str, T.Tuple[float, float]]):\n",
    "    prefixes = (\"[\", \"|\")\n",
    "    for prefix in (*prefixes, None):\n",
    "        previous: T.Optional[str] = None\n",
    "        for metric, score in averaged_scores.items():\n",
    "            if (\n",
    "                prefix is None and metric[0] not in prefixes\n",
    "            ) or prefix == metric[0]:\n",
    "                current = metric.split(\" \", 1)[0]\n",
    "                if previous is None or current != previous:\n",
    "                    print(\"\\t\" + \"-\" * 63)\n",
    "                    previous = current\n",
    "\n",
    "                metric += \":\"\n",
    "                print(f\"\\t{metric:40s}\\t{score[0]:.4f} ± {score[1]:.4f}\")\n",
    "    print(\"\\t\" + \"-\" * 63)\n",
    "\n",
    "\n",
    "def print_summaries(\n",
    "    summaries: T.Iterable[Summary],\n",
    "    class_labels: T.Optional[T.Union[int, T.Sequence[str]]] = None,\n",
    "    include_training=True,\n",
    "):\n",
    "    print()\n",
    "    for data_set in (\n",
    "        (\"Training\", \"Testing\") if include_training else (\"Testing\",)\n",
    "    ):\n",
    "        log(f\"{data_set} scores:\")\n",
    "        print_scores(\n",
    "            average_scores(\n",
    "                tuple(\n",
    "                    summary[f\"{data_set} Scores\"]  # type: ignore\n",
    "                    for summary in summaries\n",
    "                ),\n",
    "                class_labels,\n",
    "            )\n",
    "        )\n",
    "        print()\n",
    "\n",
    "    for time_name, calculator in (\n",
    "        (\"Training\", lambda x: x[\"End Fit Transform Time\"] - x[\"Start Time\"]),\n",
    "        (\"Testing\", lambda x: x[\"End Time\"] - x[\"Start Predict Score Time\"]),\n",
    "    ):\n",
    "        fold_times = tuple(map(calculator, summaries))\n",
    "        for index, fold_time in enumerate(fold_times, start=1):\n",
    "            print(f\"{time_name} fold {index} time: {fold_time}\")\n",
    "        print(\n",
    "            f\"{time_name} average fold time:\",\n",
    "            dt.timedelta(\n",
    "                seconds=float(\n",
    "                    np.mean(\n",
    "                        [fold_time.total_seconds() for fold_time in fold_times]\n",
    "                    )\n",
    "                )\n",
    "            ),\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "def label_counts(\n",
    "    y_tags: T.List[str],\n",
    "    labels: T.Dict[str, T.Dict[T.Optional[str], int]],\n",
    "    print_top_n: T.Optional[int] = None,\n",
    "):\n",
    "    counts: T.DefaultDict[T.Optional[str], int] = cll.defaultdict(int)\n",
    "    for tag in y_tags:\n",
    "        for label in labels[tag]:\n",
    "            counts[label] += 1\n",
    "    sorted_counts = sorted(  # Highest to lowest.\n",
    "        counts.items(), key=lambda x: x[1], reverse=True\n",
    "    )\n",
    "    if print_top_n is not None:\n",
    "        for index, (label, count) in enumerate(sorted_counts):\n",
    "            if index >= print_top_n:\n",
    "                print(\n",
    "                    \"Rest of the labels \"\n",
    "                    f\"({len(counts) - print_top_n}/{len(counts)}) not shown: \"\n",
    "                    f\"{sum(i[1] for i in sorted_counts[print_top_n:]):,}\",\n",
    "                )\n",
    "                break\n",
    "            if label is None:\n",
    "                label = \"-\"\n",
    "            print(f'Label \"{label}\" occurs {count:,} times.')\n",
    "\n",
    "    return dict(counts), sorted_counts\n",
    "\n",
    "\n",
    "def trimmer(\n",
    "    x: T.List[T.Any],\n",
    "    y_tags: T.List[str],\n",
    "    labels: T.Dict[str, T.Dict[T.Optional[str], int]],\n",
    "    random_seed: int = SEED,\n",
    "):\n",
    "    x, y_tags = sklearn.utils.shuffle(\n",
    "        x, y_tags, random_state=random_seed\n",
    "    )  # type: ignore\n",
    "\n",
    "    def match_counts():\n",
    "        no_matches = 0\n",
    "        some_matches = 0\n",
    "        all_matches = 0\n",
    "        for tag in y_tags:\n",
    "            if None in labels[tag]:\n",
    "                if len(labels[tag]) > 1:\n",
    "                    some_matches += 1\n",
    "                else:\n",
    "                    no_matches += 1\n",
    "            elif labels[tag]:\n",
    "                all_matches += 1\n",
    "\n",
    "        print(f\"{no_matches:,} documents with no filtering rule matches.\")\n",
    "        print(f\"{some_matches:,} documents with some filtering rule matches.\")\n",
    "        print(f\"{all_matches:,} documents with only filtering rule matches.\")\n",
    "        return no_matches, some_matches, all_matches\n",
    "\n",
    "    log(\"Before trimming:\")\n",
    "    no_matches, some_matches, all_matches = match_counts()\n",
    "    counts, sorted_counts = label_counts(\n",
    "        y_tags, labels, print_top_n=TOP_CLASSES\n",
    "    )\n",
    "\n",
    "    # Remove classes with truly insufficient support. Ideally, we would remove\n",
    "    # even more of these ultra rare classes to prevent the macro-averaged\n",
    "    # metric scores from heavily decreasing based on these miniscule\n",
    "    # appearances, but we keep them for the sake of completeness.\n",
    "    pop_off = 0\n",
    "    for label, count in reversed(sorted_counts):\n",
    "        if count >= K_FOLDS * 2:\n",
    "            break\n",
    "        pop_off += 1\n",
    "        print(\n",
    "            f'\"{label}\" has less than {K_FOLDS * 2} support,',\n",
    "            \"substituting it with the null rule.\",\n",
    "        )\n",
    "        for tag in y_tags:\n",
    "            tag_labels = labels[tag]\n",
    "            if label in tag_labels:\n",
    "                try:\n",
    "                    tag_labels[None] += tag_labels[label]\n",
    "                except KeyError:\n",
    "                    tag_labels[None] = tag_labels[label]\n",
    "                del tag_labels[label]\n",
    "        del counts[label]\n",
    "\n",
    "    try:\n",
    "        for _ in range(pop_off):\n",
    "            sorted_counts.pop()\n",
    "    except IndexError as exception:\n",
    "        raise ValueError(\"All classes lack enough support.\") from exception\n",
    "\n",
    "    log(f\"Label count mean: {np.mean(list(counts.values())):.2f}\")\n",
    "    log(f\"Label count median: {np.median(list(counts.values())):.2f}\")\n",
    "\n",
    "    # NOTE: You can remove this if you still wish to trim the majority class.\n",
    "    # We expect that the majority of documents will match no filtering rules.\n",
    "    if no_matches < some_matches and no_matches < all_matches:\n",
    "        log(\"No-filtering-rule documents is the minority.\")\n",
    "        # This should rarely occur (if ever) for a decently long dataset slice\n",
    "        # based on the current IXP Scrubber rules, at least in our private IXP\n",
    "        # dataset.\n",
    "        return (\n",
    "            x,\n",
    "            y_tags,\n",
    "            [\n",
    "                top[0] if top[0] is not None else \"-\"\n",
    "                for top in sorted_counts[0:TOP_CLASSES]\n",
    "            ],\n",
    "        )\n",
    "\n",
    "    first_label, first_count = sorted_counts[0]\n",
    "    second_label, second_count = sorted_counts[1]\n",
    "\n",
    "    # First try removing the documents with only the largest label.\n",
    "    index = 0\n",
    "    while first_count > second_count and index < len(y_tags):\n",
    "        tag = y_tags[index]\n",
    "        tag_labels = labels[tag]\n",
    "        if first_label in tag_labels and len(tag_labels) == 1:\n",
    "            x[index] = None\n",
    "            y_tags[index] = None  # type: ignore\n",
    "            del labels[tag]\n",
    "            first_count -= 1\n",
    "        index += 1\n",
    "\n",
    "    # Then try removing the documents with only the largest label and the next\n",
    "    # largest label (if the previous step still includes the imbalance).\n",
    "    index = 0\n",
    "    while first_count > second_count and index < len(y_tags):\n",
    "        tag = y_tags[index]\n",
    "        if tag is None:\n",
    "            index += 1\n",
    "            continue  # Was removed in the previous step.\n",
    "        tag_labels = labels[tag]\n",
    "        if (\n",
    "            first_label in tag_labels\n",
    "            and second_label in tag_labels\n",
    "            and len(tag_labels) == 2\n",
    "        ):\n",
    "            x[index] = None\n",
    "            y_tags[index] = None  # type: ignore\n",
    "            del labels[tag]\n",
    "            first_count -= 1\n",
    "            second_count -= 1\n",
    "        index += 1\n",
    "\n",
    "    # If the previous steps still include the imbalance, then try removing any\n",
    "    # documents with the largest label until it is equal in support with the\n",
    "    # second/next largest label.\n",
    "    index = 0\n",
    "    while first_count > second_count and index < len(y_tags):\n",
    "        tag = y_tags[index]\n",
    "        if tag is None:\n",
    "            index += 1\n",
    "            continue\n",
    "        tag_labels = labels[tag]\n",
    "        if first_label in tag_labels:\n",
    "            x[index] = None\n",
    "            y_tags[index] = None  # type: ignore\n",
    "            del labels[tag]\n",
    "            first_count -= 1\n",
    "            if second_label in tag_labels:\n",
    "                second_count -= 1\n",
    "        index += 1\n",
    "\n",
    "    x = [i for i in x if i is not None]\n",
    "    y_tags = [i for i in y_tags if i is not None]\n",
    "\n",
    "    log(\"After trimming:\")\n",
    "    match_counts()\n",
    "\n",
    "    return (\n",
    "        x,\n",
    "        y_tags,\n",
    "        [\n",
    "            top[0] if top[0] is not None else \"-\"\n",
    "            for top in label_counts(y_tags, labels, print_top_n=TOP_CLASSES)[\n",
    "                1\n",
    "            ][0:TOP_CLASSES]\n",
    "        ],\n",
    "    )\n",
    "\n",
    "\n",
    "def build_binariser(\n",
    "    y_tags: T.Iterable[str],\n",
    "    labels: T.Dict[str, T.Dict[T.Optional[str], int]],\n",
    ") -> T.Tuple[MultiLabelBinarizer, T_np.NDArray[T.Any]]:\n",
    "    binariser = MultiLabelBinarizer(sparse_output=False)\n",
    "    binarised_labels = binariser.fit_transform(\n",
    "        (label if label is not None else \"-\" for label in labels[tag])\n",
    "        for tag in y_tags\n",
    "    )\n",
    "    return binariser, binarised_labels  # type: ignore\n",
    "\n",
    "\n",
    "def scorer(\n",
    "    y_truth,\n",
    "    y_predictions,\n",
    "    average: T.Literal[\"macro\", \"micro\", \"samples\", \"weighted\"] = \"weighted\",\n",
    ") -> Scores:\n",
    "    class_scores: ClassScores = {label: {} for label in column_label_targets}\n",
    "\n",
    "    for index, (precision, recall, f1_score, support) in enumerate(\n",
    "        zip(\n",
    "            *np.asarray(\n",
    "                metrics.precision_recall_fscore_support(\n",
    "                    y_truth, y_predictions, average=None, zero_division=0\n",
    "                )\n",
    "            )\n",
    "        )\n",
    "    ):\n",
    "        class_scores[column_label_targets[index]].update(\n",
    "            {\n",
    "                \"Precision\": precision,\n",
    "                \"Recall\": recall,\n",
    "                \"F1 Score\": f1_score,\n",
    "                \"Support\": support,\n",
    "            }\n",
    "        )\n",
    "\n",
    "    for index, jaccard_score in enumerate(\n",
    "        np.asarray(\n",
    "            metrics.jaccard_score(\n",
    "                y_truth, y_predictions, average=None, zero_division=0\n",
    "            )\n",
    "        )\n",
    "    ):\n",
    "        class_scores[column_label_targets[index]][\n",
    "            \"Jaccard Score\"\n",
    "        ] = jaccard_score\n",
    "\n",
    "    for index, confusion_matrix in enumerate(\n",
    "        metrics.multilabel_confusion_matrix(y_truth, y_predictions)\n",
    "    ):\n",
    "        class_scores[column_label_targets[index]].update(\n",
    "            {\n",
    "                \"False Negatives\": confusion_matrix[1][0],\n",
    "                \"True Negatives\": confusion_matrix[0][0],\n",
    "                \"False Positives\": confusion_matrix[0][1],\n",
    "                \"True Positives\": confusion_matrix[1][1],\n",
    "            }\n",
    "        )\n",
    "\n",
    "    scores: Scores = {\"Classes\": class_scores}\n",
    "\n",
    "    # These are repeated within the averaged score computations for simplicity.\n",
    "    exact_match_ratio = float(metrics.accuracy_score(y_truth, y_predictions))\n",
    "    hamming_loss = float(metrics.hamming_loss(y_truth, y_predictions))\n",
    "\n",
    "    for average in AVERAGING_METHODS:\n",
    "        (\n",
    "            precision,\n",
    "            recall,\n",
    "            f1_score,\n",
    "            _,\n",
    "        ) = metrics.precision_recall_fscore_support(\n",
    "            y_truth, y_predictions, average=average, zero_division=0\n",
    "        )\n",
    "\n",
    "        scores[f\"Average ({average.title()})\"] = {\n",
    "            \"Precision\": float(precision),\n",
    "            \"Recall\": float(recall),\n",
    "            \"F1 Score\": float(f1_score),\n",
    "            \"Jaccard Score\": float(\n",
    "                metrics.jaccard_score(\n",
    "                    y_truth,\n",
    "                    y_predictions,\n",
    "                    average=average,\n",
    "                    zero_division=0,\n",
    "                )\n",
    "            ),\n",
    "            \"Exact Match Ratio\": exact_match_ratio,\n",
    "            \"Hamming Loss\": hamming_loss,\n",
    "        }\n",
    "\n",
    "    return scores\n",
    "\n",
    "\n",
    "def load_cached_summary(name: str, fold_index: int) -> T.Optional[Summary]:\n",
    "    summary_file_prefix = f\"{name}_fold-{fold_index}_summary.\"\n",
    "    cached_summaries = [\n",
    "        (\n",
    "            int(file[len(summary_file_prefix) : -10]),\n",
    "            os.path.join(ARTEFACTS_PATH, file),\n",
    "        )\n",
    "        for file in os.listdir(ARTEFACTS_PATH)\n",
    "        if file.startswith(summary_file_prefix) and file.endswith(\".pickle.gz\")\n",
    "    ]\n",
    "\n",
    "    if cached_summaries:  # Get the latest one.\n",
    "        return load_pickle(max(cached_summaries, key=lambda x: x[0])[1])\n",
    "\n",
    "\n",
    "def load_cached_run(name: str):\n",
    "    summaries: T.List[Summary] = []\n",
    "    for fold_index in range(K_FOLDS):\n",
    "        summary = load_cached_summary(name, fold_index)\n",
    "        if summary is not None:\n",
    "            summaries.append(summary)\n",
    "    return summaries\n",
    "\n",
    "\n",
    "def training_and_testing_scores(\n",
    "    name: str,\n",
    "    fold_index: int,\n",
    "    transformer,\n",
    "    classifier,\n",
    "    transformer_fit_transform_args=[],\n",
    "    transformer_fit_transform_kwargs={},\n",
    "    classifier_fit_args=[],\n",
    "    classifier_fit_kwargs={},\n",
    ") -> Summary:\n",
    "    if USE_LAST_CACHED_SUMMARY:\n",
    "        cached_summary = load_cached_summary(name, fold_index)\n",
    "        if cached_summary:\n",
    "            log(f\"WARNING: Using cached summary for fold {fold_index + 1}.\")\n",
    "            return cached_summary\n",
    "\n",
    "    start_time = dt.datetime.now(tz=dt.timezone.utc)\n",
    "\n",
    "    x_train = (\n",
    "        transformer.fit_transform(\n",
    "            x_train_folds[fold_index],\n",
    "            y_train_folds[fold_index],\n",
    "            *transformer_fit_transform_args,\n",
    "            **transformer_fit_transform_kwargs,\n",
    "        )\n",
    "        if transformer is not None\n",
    "        else x_train_folds[fold_index]\n",
    "    )\n",
    "\n",
    "    classifier.fit(\n",
    "        x_train,\n",
    "        y_train_folds[fold_index],\n",
    "        *classifier_fit_args,\n",
    "        **classifier_fit_kwargs,\n",
    "    )\n",
    "\n",
    "    end_fit_transform_time = dt.datetime.now(tz=dt.timezone.utc)\n",
    "\n",
    "    training_scores = scorer(\n",
    "        y_train_folds[fold_index], classifier.predict(x_train)\n",
    "    )\n",
    "\n",
    "    start_predict_score_time = dt.datetime.now(tz=dt.timezone.utc)\n",
    "\n",
    "    testing_scores = scorer(\n",
    "        y_test_folds[fold_index],\n",
    "        classifier.predict(\n",
    "            transformer.transform(x_test_folds[fold_index])\n",
    "            if transformer is not None\n",
    "            else x_test_folds[fold_index]\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    end_time = dt.datetime.now(tz=dt.timezone.utc)\n",
    "\n",
    "    summary: Summary = {\n",
    "        \"Name\": name,\n",
    "        \"Fold\": fold_index,\n",
    "        \"Classifier\": CLASSIFIER_NAME,\n",
    "        \"Start Time\": start_time,\n",
    "        \"End Fit Transform Time\": end_fit_transform_time,\n",
    "        \"Start Predict Score Time\": start_predict_score_time,\n",
    "        \"End Time\": end_time,\n",
    "        \"Training Scores\": training_scores,\n",
    "        \"Testing Scores\": testing_scores,\n",
    "    }\n",
    "\n",
    "    # This is just in case something goes wrong with the Jupyter kernel. It\n",
    "    # caches the results of the current fold so that we can load it again in\n",
    "    # case of a crash, or we can simply reload the results.\n",
    "    try:\n",
    "        save_pickle(\n",
    "            summary,\n",
    "            os.path.join(\n",
    "                ARTEFACTS_PATH,\n",
    "                (\n",
    "                    f\"{name}_fold-{fold_index}_summary.\"\n",
    "                    f\"{round(dt.datetime.now().timestamp())}.pickle.gz\"\n",
    "                ),\n",
    "            ),\n",
    "        )\n",
    "    except Exception as exception:\n",
    "        log(f\"Error saving scores: {exception}\", file=sys.stderr)\n",
    "\n",
    "    return summary\n",
    "\n",
    "\n",
    "def cross_validator(\n",
    "    x,\n",
    "    y,\n",
    "    binariser: MultiLabelBinarizer,\n",
    "    classification_function: T.Callable[[int], Summary],\n",
    "    simultaneous_folds: T.Optional[int] = K_FOLDS,\n",
    "    fold_limit: int = K_FOLD_LIMIT,\n",
    "    worker_mode: T.Literal[\"thread\", \"process\"] = \"process\",\n",
    "):\n",
    "    if SAMPLE_LIMIT is not None:\n",
    "        log(\n",
    "            f\"Warning: Using a sample limit: '{SAMPLE_LIMIT}'.\",\n",
    "            file=sys.stderr,\n",
    "        )\n",
    "\n",
    "    x_train_folds.clear()\n",
    "    y_train_folds.clear()\n",
    "\n",
    "    x_test_folds.clear()\n",
    "    y_test_folds.clear()\n",
    "\n",
    "    column_label_targets.clear()\n",
    "    column_label_targets.extend(binariser.classes_)\n",
    "\n",
    "    random.seed(SEED)\n",
    "    np.random.seed(SEED)\n",
    "\n",
    "    splitter = MultilabelStratifiedKFold(\n",
    "        n_splits=K_FOLDS, shuffle=True, random_state=SEED\n",
    "    )\n",
    "\n",
    "    fold_count = 0\n",
    "    for train_indexes, test_indexes in splitter.split(x, y):\n",
    "        log(f\"Splitting fold {fold_count + 1}...\")\n",
    "\n",
    "        x_train_folds.append([x[index] for index in train_indexes])\n",
    "        y_train_folds.append([y[index] for index in train_indexes])\n",
    "\n",
    "        x_test_folds.append([x[index] for index in test_indexes])\n",
    "        y_test_folds.append([y[index] for index in test_indexes])\n",
    "\n",
    "        fold_count += 1\n",
    "        if fold_count >= fold_limit:\n",
    "            break\n",
    "\n",
    "    log(\"Beginning classification...\")\n",
    "    if simultaneous_folds and simultaneous_folds > 1 and fold_count > 1:\n",
    "        if worker_mode == \"process\":\n",
    "            with CustomPool(simultaneous_folds, maxtasksperchild=1) as pool:\n",
    "                return pool.map(\n",
    "                    classification_function, range(fold_count), chunksize=1\n",
    "                )\n",
    "        elif worker_mode == \"thread\":\n",
    "            with mp_pool.ThreadPool(simultaneous_folds) as pool:\n",
    "                return pool.map(\n",
    "                    classification_function, range(fold_count), chunksize=1\n",
    "                )\n",
    "        else:\n",
    "            raise ValueError(f\"Unknown worker mode: {worker_mode}\")\n",
    "    else:\n",
    "        return list(map(classification_function, range(fold_count)))\n",
    "\n",
    "\n",
    "def build_classifier():\n",
    "    log(f'Using classifier \"{CLASSIFIER_NAME}\".')\n",
    "\n",
    "    if CLASSIFIER_NAME == \"k-NN\":\n",
    "        from sklearn.neighbors import KNeighborsClassifier\n",
    "\n",
    "        return KNeighborsClassifier(\n",
    "            n_neighbors=10, n_jobs=-1, weights=\"distance\"\n",
    "        )\n",
    "    elif CLASSIFIER_NAME == \"XGBoost\":\n",
    "        from xgboost import XGBClassifier\n",
    "\n",
    "        return XGBClassifier(tree_method=\"hist\", n_jobs=-1, random_state=SEED)\n",
    "    elif CLASSIFIER_NAME == \"Dummy\":\n",
    "        from sklearn.dummy import DummyClassifier\n",
    "\n",
    "        return DummyClassifier(random_state=SEED)\n",
    "\n",
    "    raise ValueError(f\"Unknown classifier: {CLASSIFIER_NAME}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dummy\n",
    "Always predicts the null rule label \"-\" (no filtering rules matched) regardless of the input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DUMMY_LABELS_PATH = os.path.join(\n",
    "    ARTEFACTS_PATH, \"month.baseline.labels.pickle.gz\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log(\"Loading labels...\")\n",
    "labels: T.Dict[str, T.Dict[T.Optional[str], int]] = load_pickle(\n",
    "    DUMMY_LABELS_PATH\n",
    ")\n",
    "log(\"Loaded labels.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log(\"Organising data and labels...\")\n",
    "x: T.List[T.Any] = [True for _ in range(len(labels))]\n",
    "y_tags: T.List[str] = list(labels)\n",
    "\n",
    "if SAMPLE_LIMIT is not None:\n",
    "    log(f\"Warning: Using a sample limit: '{SAMPLE_LIMIT}'.\", file=sys.stderr)\n",
    "    x = x[SAMPLE_LIMIT]\n",
    "    y_tags = y_tags[SAMPLE_LIMIT]\n",
    "\n",
    "if TRIM_MAJORITY_CLASS:\n",
    "    log(\"Trimming...\")\n",
    "    x, y_tags, top_classes = trimmer(x, y_tags, labels)\n",
    "else:\n",
    "    top_classes = TOP_CLASSES\n",
    "\n",
    "binariser, y = build_binariser(y_tags, labels)\n",
    "log(\"Organised data and labels.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dummy_classification(fold_index: int) -> Summary:\n",
    "    from sklearn.dummy import DummyClassifier\n",
    "\n",
    "    return training_and_testing_scores(\n",
    "        \"Dummy\",\n",
    "        fold_index,\n",
    "        None,\n",
    "        DummyClassifier(\n",
    "            strategy=\"constant\",\n",
    "            constant=np.array(\n",
    "                [int(label == \"-\") for label in column_label_targets]\n",
    "            ),\n",
    "            random_state=SEED,\n",
    "        ),\n",
    "    )\n",
    "\n",
    "\n",
    "dummy_fold_summaries = cross_validator(x, y, binariser, dummy_classification)\n",
    "\n",
    "print_summaries(dummy_fold_summaries, top_classes)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Baseline Heuristic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": ""
    }
   },
   "outputs": [],
   "source": [
    "BASELINE_DATA_PATH = os.path.join(\n",
    "    ARTEFACTS_PATH, \"month.baseline.data.pickle.gz\"\n",
    ")\n",
    "\n",
    "BASELINE_LABELS_PATH = os.path.join(\n",
    "    ARTEFACTS_PATH, \"month.baseline.labels.pickle.gz\"\n",
    ")\n",
    "\n",
    "# This is based on the domain knowledge we specify in the paper.\n",
    "AMPLIFICATION_PORTS = (\n",
    "    19,\n",
    "    53,\n",
    "    69,\n",
    "    88,\n",
    "    123,\n",
    "    137,\n",
    "    138,\n",
    "    139,\n",
    "    161,\n",
    "    389,\n",
    "    853,\n",
    "    5353,\n",
    "    11211,\n",
    ")\n",
    "\n",
    "AMPLIFICATION_SOURCE_PORTS = tuple(  # Note that this is post-amplification.\n",
    "    map(lambda port: f\"{port}->\", AMPLIFICATION_PORTS)\n",
    ")\n",
    "\n",
    "# We excluded port zero for this \"generic\" range, as it is technically invalid.\n",
    "GENERIC_SYSTEM_SOURCE_PORTS = frozenset(\n",
    "    map(\n",
    "        lambda port: f\"{port}->\",\n",
    "        (port for port in range(1, 1024) if port not in AMPLIFICATION_PORTS),\n",
    "    )\n",
    ")\n",
    "\n",
    "GENERIC_USER_SOURCE_PORTS = frozenset(\n",
    "    map(\n",
    "        lambda port: f\"{port}->\",\n",
    "        (\n",
    "            port\n",
    "            for port in range(1024, 49152)\n",
    "            if port not in AMPLIFICATION_PORTS\n",
    "        ),\n",
    "    )\n",
    ")\n",
    "\n",
    "GENERIC_DYNAMIC_SOURCE_PORTS = frozenset(\n",
    "    map(\n",
    "        lambda port: f\"{port}->\",\n",
    "        (\n",
    "            port\n",
    "            for port in range(49152, 65536)\n",
    "            if port not in AMPLIFICATION_PORTS\n",
    "        ),\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log(\"Loading data...\")\n",
    "data: T.Dict[str, T.Dict[str, T.Tuple[int, int, int]]] = load_pickle(\n",
    "    BASELINE_DATA_PATH\n",
    ")\n",
    "log(\"Loaded data.\")\n",
    "\n",
    "log(\"Loading labels...\")\n",
    "labels: T.Dict[str, T.Dict[T.Optional[str], int]] = load_pickle(\n",
    "    BASELINE_LABELS_PATH\n",
    ")\n",
    "log(\"Loaded labels.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log(\"Organising data and labels...\")\n",
    "x: T.List[T_np.NDArray[T.Any]] = []\n",
    "y_tags: T.List[str] = []\n",
    "\n",
    "log(\"Transforming...\")\n",
    "for tag, tag_counts in data.items():\n",
    "    x_tag: T.List[int] = []\n",
    "\n",
    "    for port_range in (AMPLIFICATION_SOURCE_PORTS,):\n",
    "        for port in port_range:\n",
    "            x_tag.extend(tag_counts.get(port, (0, 0, 0)))\n",
    "\n",
    "    for port_range in (\n",
    "        GENERIC_SYSTEM_SOURCE_PORTS,\n",
    "        GENERIC_USER_SOURCE_PORTS,\n",
    "        GENERIC_DYNAMIC_SOURCE_PORTS,\n",
    "    ):\n",
    "        flow_count = 0\n",
    "        byte_count = 0\n",
    "        packet_count = 0\n",
    "        for port, port_counts in tag_counts.items():\n",
    "            if port in port_range:\n",
    "                flow_count += port_counts[0]\n",
    "                byte_count += port_counts[1]\n",
    "                packet_count += port_counts[2]\n",
    "        x_tag.extend((flow_count, byte_count, packet_count))\n",
    "\n",
    "    x.append(np.array(x_tag, dtype=np.float64))\n",
    "    y_tags.append(tag)\n",
    "log(\"Transformed.\")\n",
    "\n",
    "if SAMPLE_LIMIT is not None:\n",
    "    log(f\"Warning: Using a sample limit: '{SAMPLE_LIMIT}'.\", file=sys.stderr)\n",
    "    x = x[SAMPLE_LIMIT]\n",
    "    y_tags = y_tags[SAMPLE_LIMIT]\n",
    "\n",
    "if TRIM_MAJORITY_CLASS:\n",
    "    log(\"Trimming...\")\n",
    "    x, y_tags, top_classes = trimmer(x, y_tags, labels)\n",
    "else:\n",
    "    top_classes = TOP_CLASSES\n",
    "\n",
    "binariser, y = build_binariser(y_tags, labels)\n",
    "log(\"Organised data and labels.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def baseline_classification(fold_index: int) -> Summary:\n",
    "    return training_and_testing_scores(\n",
    "        \"Baseline\", fold_index, None, build_classifier()\n",
    "    )\n",
    "\n",
    "\n",
    "baseline_fold_summaries = cross_validator(\n",
    "    x, y, binariser, baseline_classification\n",
    ")\n",
    "\n",
    "print_summaries(baseline_fold_summaries, top_classes)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Word2Vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gensim\n",
    "\n",
    "from gensim.models.callbacks import CallbackAny2Vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "WORD2VEC_CORPUS_PATH = os.path.join(\n",
    "    ARTEFACTS_PATH, \"month.word2vec.data.pickle.gz\"\n",
    ")\n",
    "\n",
    "WORD2VEC_MODEL_PATH = os.path.join(ARTEFACTS_PATH, \"month.word2vec.model\")\n",
    "\n",
    "WORD2VEC_LABELS_PATH = os.path.join(\n",
    "    ARTEFACTS_PATH, \"month.word2vec.labels.pickle.gz\"\n",
    ")\n",
    "\n",
    "VECTOR_LENGTH = 100\n",
    "CONTEXT_WINDOW_SIZE = 10\n",
    "EPOCHS = 15\n",
    "SKIPGRAM_LEARNING_MODEL = True\n",
    "NEGATIVE_SAMPLES = 5  # 5 is the Gensim default.\n",
    "NEGATIVE_SAMPLING_EXPONENT = 0.75  # 0.75 is the Gensim default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(WORD2VEC_MODEL_PATH):\n",
    "    log(\"Loading corpus...\")\n",
    "    corpus: T.Dict[str, T.List[str]] = load_pickle(WORD2VEC_CORPUS_PATH)\n",
    "    log(\"Loaded corpus.\")\n",
    "\n",
    "    class EpochsCallback(CallbackAny2Vec):\n",
    "        def __init__(self):\n",
    "            self.epoch = 1\n",
    "\n",
    "        def on_epoch_end(self, model):\n",
    "            loss = model.get_latest_training_loss()\n",
    "            if loss != 0.0:\n",
    "                log(f\"Epoch {self.epoch} completed, loss is {loss:.3f}.\")\n",
    "            else:\n",
    "                log(f\"Epoch {self.epoch} completed, loss is unavailable.\")\n",
    "            self.epoch += 1\n",
    "\n",
    "    log(\"Generating Word2Vec model...\")\n",
    "    model = gensim.models.Word2Vec(\n",
    "        sentences=iter(corpus.values()),\n",
    "        window=CONTEXT_WINDOW_SIZE,\n",
    "        epochs=EPOCHS,\n",
    "        sg=int(SKIPGRAM_LEARNING_MODEL),\n",
    "        ns_exponent=NEGATIVE_SAMPLING_EXPONENT,\n",
    "        negative=NEGATIVE_SAMPLES,\n",
    "        hs=int(NEGATIVE_SAMPLES == 0),\n",
    "        min_count=1,\n",
    "        max_vocab_size=None,\n",
    "        compute_loss=True,\n",
    "        workers=min(os.cpu_count() or 8, 16),\n",
    "        callbacks=[EpochsCallback()],\n",
    "        seed=SEED,\n",
    "    )\n",
    "\n",
    "    log(\"Generated Word2Vec model. Saving...\")\n",
    "    model.save(WORD2VEC_MODEL_PATH)\n",
    "    log(\"Saved Word2Vec model.\")\n",
    "else:\n",
    "    log(\"Loading model...\")\n",
    "    model = gensim.models.Word2Vec.load(WORD2VEC_MODEL_PATH)\n",
    "    log(\"Loaded model.\")\n",
    "\n",
    "log(\"Loading labels...\")\n",
    "labels: T.Dict[str, T.Dict[T.Optional[str], int]] = load_pickle(\n",
    "    WORD2VEC_LABELS_PATH\n",
    ")\n",
    "log(\"Loaded labels.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log(\"Organising data and labels...\")\n",
    "x: T.List[T_np.NDArray[T.Any]] = []\n",
    "y_tags: T.List[str] = []\n",
    "\n",
    "for tag, tag_labels in labels.items():\n",
    "    try:\n",
    "        x.append(model.wv[tag])\n",
    "        y_tags.append(tag)\n",
    "    except KeyError:\n",
    "        log(f\"Untrained tag (word) '{tag}' and labels: {tag_labels}\")\n",
    "\n",
    "if SAMPLE_LIMIT is not None:\n",
    "    log(f\"Warning: Using a sample limit: '{SAMPLE_LIMIT}'.\", file=sys.stderr)\n",
    "    x = x[SAMPLE_LIMIT]\n",
    "    y_tags = y_tags[SAMPLE_LIMIT]\n",
    "\n",
    "if TRIM_MAJORITY_CLASS:\n",
    "    log(\"Trimming...\")\n",
    "    x, y_tags, top_classes = trimmer(x, y_tags, labels)\n",
    "else:\n",
    "    top_classes = TOP_CLASSES\n",
    "\n",
    "binariser, y = build_binariser(y_tags, labels)\n",
    "log(\"Organised data and labels.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def word2vec_classification(fold_index: int) -> Summary:\n",
    "    return training_and_testing_scores(\n",
    "        \"Word2Vec\", fold_index, None, build_classifier()\n",
    "    )\n",
    "\n",
    "\n",
    "word2vec_fold_summaries = cross_validator(\n",
    "    x, y, binariser, word2vec_classification\n",
    ")\n",
    "\n",
    "print_summaries(word2vec_fold_summaries, top_classes)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Doc2Vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gensim\n",
    "\n",
    "from gensim.models.callbacks import CallbackAny2Vec\n",
    "from sklearn.base import BaseEstimator, TransformerMixin\n",
    "from sklearn.exceptions import NotFittedError\n",
    "from sklearn.utils.estimator_checks import check_estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DOC2VEC_CORPUS_PATH = os.path.join(\n",
    "    ARTEFACTS_PATH, \"month.standard.data.pickle.gz\"\n",
    ")\n",
    "\n",
    "DOC2VEC_LABELS_PATH = os.path.join(\n",
    "    ARTEFACTS_PATH, \"month.standard.labels.pickle.gz\"\n",
    ")\n",
    "\n",
    "VECTOR_LENGTH = 100\n",
    "CONTEXT_WINDOW_SIZE = 10\n",
    "EPOCHS = 25\n",
    "DISTRIBUTED_MEMORY_LEARNING_MODEL = True\n",
    "NEGATIVE_SAMPLES = 5  # 5 is the Gensim default.\n",
    "NEGATIVE_SAMPLING_EXPONENT = 0.75  # 0.75 is the Gensim default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log(\"Loading corpus...\")\n",
    "corpus: T.Dict[str, T.List[str]] = load_pickle(DOC2VEC_CORPUS_PATH)\n",
    "log(\"Loaded corpus.\")\n",
    "\n",
    "log(\"Loading labels...\")\n",
    "labels: T.Dict[str, T.Dict[T.Optional[str], int]] = load_pickle(\n",
    "    DOC2VEC_LABELS_PATH\n",
    ")\n",
    "log(\"Loaded labels.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log(\"Organising data and labels...\")\n",
    "x: T.List[gensim.models.doc2vec.TaggedDocument] = []\n",
    "y_tags: T.List[str] = []\n",
    "\n",
    "for tag, words in corpus.items():\n",
    "    x.append(gensim.models.doc2vec.TaggedDocument(words, [tag]))\n",
    "    y_tags.append(tag)\n",
    "\n",
    "if SAMPLE_LIMIT is not None:\n",
    "    log(f\"Warning: Using a sample limit: '{SAMPLE_LIMIT}'.\", file=sys.stderr)\n",
    "    x = x[SAMPLE_LIMIT]\n",
    "    y_tags = y_tags[SAMPLE_LIMIT]\n",
    "\n",
    "if TRIM_MAJORITY_CLASS:\n",
    "    log(\"Trimming...\")\n",
    "    x, y_tags, top_classes = trimmer(x, y_tags, labels)\n",
    "else:\n",
    "    top_classes = TOP_CLASSES\n",
    "\n",
    "binariser, y = build_binariser(y_tags, labels)\n",
    "log(\"Organised data and labels.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Doc2Vec(BaseEstimator, TransformerMixin):\n",
    "    class EpochsCallback(CallbackAny2Vec):\n",
    "        def __init__(self):\n",
    "            self.epoch = 1\n",
    "\n",
    "        def on_epoch_end(self, model):\n",
    "            loss = model.get_latest_training_loss()\n",
    "            if loss != 0.0:\n",
    "                log(f\"Epoch {self.epoch} completed, loss is {loss:.3f}.\")\n",
    "            else:\n",
    "                log(f\"Epoch {self.epoch} completed, loss is unavailable.\")\n",
    "            self.epoch += 1\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        vector_length=VECTOR_LENGTH,\n",
    "        context_window_size=CONTEXT_WINDOW_SIZE,\n",
    "        epochs=EPOCHS,\n",
    "        negative_samples=NEGATIVE_SAMPLES,\n",
    "        negative_sampling_exponent=NEGATIVE_SAMPLING_EXPONENT,\n",
    "        distributed_memory_learning_model=DISTRIBUTED_MEMORY_LEARNING_MODEL,\n",
    "        random_state: T.Optional[int] = SEED,\n",
    "    ):\n",
    "        self.vector_length = vector_length\n",
    "        self.context_window_size = context_window_size\n",
    "        self.epochs = epochs\n",
    "        self.negative_samples = negative_samples\n",
    "        self.negative_sampling_exponent = negative_sampling_exponent\n",
    "        self.distributed_memory_learning_model = (\n",
    "            distributed_memory_learning_model\n",
    "        )\n",
    "        self.random_state = random_state\n",
    "\n",
    "    def _more_tags(self):\n",
    "        return {\"X_types\": [T.List[gensim.models.doc2vec.TaggedDocument]]}\n",
    "\n",
    "    def fit(self, x, y=None):\n",
    "        log(\"Setting up Doc2Vec model...\")\n",
    "\n",
    "        self.doc2vec_model = gensim.models.Doc2Vec(\n",
    "            documents=x,\n",
    "            vector_size=self.vector_length,\n",
    "            window=self.context_window_size,\n",
    "            epochs=self.epochs,\n",
    "            dm=int(self.distributed_memory_learning_model),\n",
    "            ns_exponent=self.negative_sampling_exponent,\n",
    "            negative=self.negative_samples,\n",
    "            hs=int(self.negative_samples == 0),\n",
    "            dbow_words=0,\n",
    "            min_count=1,\n",
    "            max_vocab_size=None,\n",
    "            compute_loss=True,\n",
    "            workers=min(os.cpu_count() or 8, 16),  # Hardly helps.\n",
    "            callbacks=[self.EpochsCallback()],\n",
    "            seed=self.random_state,\n",
    "        )\n",
    "\n",
    "        self._transformation_cache: T.Dict[\n",
    "            T.Tuple[int, int], T.List[T_np.NDArray[T.Any]]\n",
    "        ] = {\n",
    "            (id(x), len(x)): [\n",
    "                self.doc2vec_model.dv[document.tags[0]] for document in x\n",
    "            ]\n",
    "        }\n",
    "\n",
    "        log(\"Fit complete.\")\n",
    "        return self\n",
    "\n",
    "    def transform(self, x, y=None, cache_transformation=True):\n",
    "        if not hasattr(self, \"doc2vec_model\"):\n",
    "            raise NotFittedError(\"Fit the transformer first.\")\n",
    "\n",
    "        if not isinstance(x[0], gensim.models.doc2vec.TaggedDocument):\n",
    "            raise ValueError(\"Expected a list of TaggedDocuments.\")\n",
    "\n",
    "        inferred = self._transformation_cache.get((id(x), len(x)), [])\n",
    "        if inferred:\n",
    "            log(\"Transformation already cached.\")\n",
    "            return inferred\n",
    "\n",
    "        log(f\"Transforming {len(x):,} documents...\")\n",
    "        progress_count = len(x) // 5\n",
    "\n",
    "        for count, document in enumerate(x, start=1):\n",
    "            inferred.append(self.doc2vec_model.infer_vector(document.words))\n",
    "            if count % progress_count == 0:\n",
    "                log(f\"{count:,} out of {len(x):,} documents transformed...\")\n",
    "\n",
    "        self._transformation_cache[(id(x), len(x))] = inferred\n",
    "        return inferred\n",
    "\n",
    "\n",
    "check_estimator(Doc2Vec())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def doc2vec_classification(fold_index: int) -> Summary:\n",
    "    return training_and_testing_scores(\n",
    "        \"Doc2Vec\", fold_index, Doc2Vec(), build_classifier()\n",
    "    )\n",
    "\n",
    "\n",
    "doc2vec_fold_summaries = cross_validator(\n",
    "    x, y, binariser, doc2vec_classification, simultaneous_folds=1\n",
    ")\n",
    "\n",
    "print_summaries(doc2vec_fold_summaries, top_classes)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Latent Semantic Analysis\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gensim\n",
    "\n",
    "from sklearn.base import BaseEstimator, TransformerMixin\n",
    "from sklearn.exceptions import NotFittedError\n",
    "from sklearn.utils.estimator_checks import check_estimator\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.feature_extraction.text import TfidfVectorizer\n",
    "from sklearn.decomposition import TruncatedSVD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LSA_CORPUS_PATH = os.path.join(ARTEFACTS_PATH, \"month.standard.data.pickle.gz\")\n",
    "\n",
    "LSA_LABELS_PATH = os.path.join(\n",
    "    ARTEFACTS_PATH, \"month.standard.labels.pickle.gz\"\n",
    ")\n",
    "\n",
    "USE_GENSIM = False  # The paper uses scikit-learn's implementation.\n",
    "\n",
    "# The untuned hyperparameters used in the paper. This was our first attempt.\n",
    "NGRAM_RANGE = (1, 3)\n",
    "SVD_COMPONENTS = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log(\"Loading corpus...\")\n",
    "corpus: T.Dict[str, T.List[str]] = load_pickle(LSA_CORPUS_PATH)\n",
    "log(\"Loaded corpus.\")\n",
    "\n",
    "log(\"Loading labels...\")\n",
    "labels: T.Dict[str, T.Dict[T.Optional[str], int]] = load_pickle(\n",
    "    LSA_LABELS_PATH\n",
    ")\n",
    "log(\"Loaded labels.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log(\"Organising data and labels...\")\n",
    "x: T.List[T.Tuple[str, ...]] = []\n",
    "y_tags: T.List[str] = []\n",
    "\n",
    "for tag, words in corpus.items():\n",
    "    x.append(tuple(words))\n",
    "    y_tags.append(tag)\n",
    "\n",
    "if SAMPLE_LIMIT is not None:\n",
    "    log(f\"Warning: Using a sample limit: '{SAMPLE_LIMIT}'.\", file=sys.stderr)\n",
    "    x = x[SAMPLE_LIMIT]\n",
    "    y_tags = y_tags[SAMPLE_LIMIT]\n",
    "\n",
    "if TRIM_MAJORITY_CLASS:\n",
    "    log(\"Trimming...\")\n",
    "    x, y_tags, top_classes = trimmer(x, y_tags, labels)\n",
    "else:\n",
    "    top_classes = TOP_CLASSES\n",
    "\n",
    "binariser, y = build_binariser(y_tags, labels)\n",
    "log(\"Organised data and labels.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LatentSemanticAnalysis(BaseEstimator, TransformerMixin):  # For Gensim.\n",
    "    def __init__(\n",
    "        self,\n",
    "        n_components=SVD_COMPONENTS,\n",
    "        ngram_range: T.Tuple[int, int] = NGRAM_RANGE,\n",
    "        ngram_joiner: str = \" \",\n",
    "        random_state: T.Optional[int] = SEED,\n",
    "        filter_extremes=False,\n",
    "    ):\n",
    "        self.n_components = n_components\n",
    "        self.ngram_range = ngram_range\n",
    "        self.ngram_joiner = ngram_joiner\n",
    "        self.random_state = random_state\n",
    "        self.filter_extremes = filter_extremes\n",
    "\n",
    "    def _more_tags(self):\n",
    "        return {\"X_types\": [T.List[T.Iterable[str]]]}\n",
    "\n",
    "    def _ngram_generator(self, x: T.Sequence[str]):\n",
    "        for n in range(self.ngram_range[0], self.ngram_range[1] + 1):\n",
    "            for gram in zip(*[x[i:] for i in range(n)]):\n",
    "                yield self.ngram_joiner.join(gram)\n",
    "\n",
    "    def fit(self, x, y=None):\n",
    "        if self.ngram_range != (1, 1):\n",
    "            log(\"Generating n-grams above unigrams...\")\n",
    "            x = [list(self._ngram_generator(document)) for document in x]\n",
    "\n",
    "        log(\"Generating dictionary...\")\n",
    "        self.dictionary = gensim.corpora.Dictionary(x, prune_at=2**64)\n",
    "\n",
    "        if self.filter_extremes:\n",
    "            log(f\"{len(self.dictionary):,} unique terms. Filtering...\")\n",
    "            self.dictionary.filter_extremes(no_below=2)\n",
    "            log(f\"{len(self.dictionary):,} unique terms after filtering.\")\n",
    "        else:\n",
    "            log(f\"{len(self.dictionary):,} unique terms.\")\n",
    "\n",
    "        log(f\"Generating corpus of {len(x):,} documents...\")\n",
    "        self.corpus = list(map(self.dictionary.doc2bow, x))\n",
    "\n",
    "        log(\"Setting up TF-IDF model...\")\n",
    "        self.tf_idf_model = gensim.models.TfidfModel(self.corpus)\n",
    "\n",
    "        log(\"Setting up LSA model...\")\n",
    "        self.lsa_model = gensim.models.LsiModel(\n",
    "            self.tf_idf_model[self.corpus],\n",
    "            onepass=False,  # TODO: Which is faster time-wise?\n",
    "            id2word=self.dictionary,\n",
    "            num_topics=self.n_components,\n",
    "            random_seed=self.random_state,\n",
    "        )\n",
    "\n",
    "        log(\"Fit complete.\")\n",
    "        return self\n",
    "\n",
    "    def transform(self, x, y=None):\n",
    "        if not hasattr(self, \"lsa_model\"):\n",
    "            raise NotFittedError(\"Fit the transformer first.\")\n",
    "\n",
    "        if isinstance(x[0], str):\n",
    "            x = [x]\n",
    "\n",
    "        log(f\"Transforming {len(x):,} documents...\")\n",
    "\n",
    "        return np.reshape(\n",
    "            np.array(\n",
    "                [\n",
    "                    gensim.matutils.sparse2full(\n",
    "                        self.lsa_model[document], self.n_components\n",
    "                    )\n",
    "                    for document in (\n",
    "                        self.tf_idf_model[\n",
    "                            self.dictionary.doc2bow(\n",
    "                                list(self._ngram_generator(raw_document))\n",
    "                            )\n",
    "                        ]\n",
    "                        for raw_document in x\n",
    "                    )\n",
    "                ]\n",
    "            ),\n",
    "            (len(x), self.n_components),\n",
    "        )\n",
    "\n",
    "\n",
    "check_estimator(LatentSemanticAnalysis())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lsa_classification(fold_index: int) -> Summary:\n",
    "    if USE_GENSIM:\n",
    "        import logging\n",
    "\n",
    "        logging.basicConfig(\n",
    "            format=\"%(asctime)s : %(levelname)s : %(message)s\",\n",
    "            level=logging.DEBUG,\n",
    "            filename=os.path.join(ARTEFACTS_PATH, \"lsa.gensim.log\"),\n",
    "            filemode=\"w\",\n",
    "        )\n",
    "\n",
    "        log(\"LSA via Gensim.\")\n",
    "        transformer = LatentSemanticAnalysis()\n",
    "    else:\n",
    "        log(\"LSA via scikit-learn.\")\n",
    "        transformer = make_pipeline(\n",
    "            TfidfVectorizer(\n",
    "                ngram_range=NGRAM_RANGE,\n",
    "                lowercase=False,\n",
    "                tokenizer=lambda x: x,\n",
    "                preprocessor=lambda x: x,\n",
    "                token_pattern=None,  # type: ignore\n",
    "                use_idf=True,\n",
    "                sublinear_tf=True,\n",
    "            ),\n",
    "            TruncatedSVD(\n",
    "                algorithm=\"arpack\",\n",
    "                n_components=SVD_COMPONENTS,\n",
    "                random_state=SEED,\n",
    "            ),\n",
    "            verbose=True,\n",
    "        )\n",
    "\n",
    "    return training_and_testing_scores(\n",
    "        \"LSA\",\n",
    "        fold_index,\n",
    "        transformer,\n",
    "        build_classifier(),\n",
    "    )\n",
    "\n",
    "\n",
    "lsa_fold_summaries = cross_validator(\n",
    "    x, y, binariser, lsa_classification, simultaneous_folds=1\n",
    ")\n",
    "\n",
    "print_summaries(lsa_fold_summaries, top_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gridsearch_fold_summaries = []\n",
    "for vector_length, minimum_frequency in it.product(\n",
    "    (100, 200, 300), (10, 1000, 10000)\n",
    "):\n",
    "    log(\n",
    "        f\"Gridsearching LSA with vector length {vector_length} \"\n",
    "        f\"and minimum frequency {minimum_frequency}...\"\n",
    "    )\n",
    "\n",
    "    def gridsearch_lsa_classification(fold_index: int) -> Summary:\n",
    "        if USE_GENSIM:\n",
    "            raise RuntimeError(\"Gensim not supported for this experiment.\")\n",
    "\n",
    "        transformer = make_pipeline(\n",
    "            TfidfVectorizer(\n",
    "                ngram_range=NGRAM_RANGE,\n",
    "                lowercase=False,\n",
    "                tokenizer=lambda x: x,\n",
    "                preprocessor=lambda x: x,\n",
    "                token_pattern=None,  # type: ignore\n",
    "                use_idf=True,\n",
    "                sublinear_tf=True,\n",
    "                min_df=minimum_frequency,\n",
    "            ),\n",
    "            TruncatedSVD(\n",
    "                algorithm=\"arpack\",\n",
    "                n_components=vector_length,\n",
    "                random_state=SEED,\n",
    "            ),\n",
    "            verbose=True,\n",
    "        )\n",
    "\n",
    "        return training_and_testing_scores(\n",
    "            f\"Gridsearch-V{vector_length}-MinDF{minimum_frequency}-LSA\",\n",
    "            fold_index,\n",
    "            transformer,\n",
    "            build_classifier(),\n",
    "        )\n",
    "\n",
    "    gridsearch_fold_summaries.append(\n",
    "        cross_validator(\n",
    "            x,\n",
    "            y,\n",
    "            binariser,\n",
    "            gridsearch_lsa_classification,\n",
    "            simultaneous_folds=1,\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ngram_fold_summaries = []\n",
    "for ngram_range in (\n",
    "    (1, 1),\n",
    "    (1, 2),\n",
    "    # (1, 3),  # This is the default we present; no need to recalculate it.\n",
    "    (2, 2),\n",
    "    (2, 3),\n",
    "    (3, 3),\n",
    "):\n",
    "    log(f\"Trying LSA with n-gram range {ngram_range}...\")\n",
    "\n",
    "    def ngram_lsa_classification(fold_index: int) -> Summary:\n",
    "        if USE_GENSIM:\n",
    "            raise RuntimeError(\"Gensim not supported for this experiment.\")\n",
    "\n",
    "        transformer = make_pipeline(\n",
    "            TfidfVectorizer(\n",
    "                ngram_range=ngram_range,\n",
    "                lowercase=False,\n",
    "                tokenizer=lambda x: x,\n",
    "                preprocessor=lambda x: x,\n",
    "                token_pattern=None,  # type: ignore\n",
    "                use_idf=True,\n",
    "                sublinear_tf=True,\n",
    "            ),\n",
    "            TruncatedSVD(\n",
    "                algorithm=\"arpack\",\n",
    "                n_components=SVD_COMPONENTS,\n",
    "                random_state=SEED,\n",
    "            ),\n",
    "            verbose=True,\n",
    "        )\n",
    "\n",
    "        return training_and_testing_scores(\n",
    "            f\"Ngram-{ngram_range[0]}-{ngram_range[1]}-LSA\",\n",
    "            fold_index,\n",
    "            transformer,\n",
    "            build_classifier(),\n",
    "        )\n",
    "\n",
    "    ngram_fold_summaries.append(\n",
    "        cross_validator(\n",
    "            x,\n",
    "            y,\n",
    "            binariser,\n",
    "            ngram_lsa_classification,\n",
    "            simultaneous_folds=1,\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Latent Semantic Analysis (without Domain Knowledge)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gensim\n",
    "\n",
    "from sklearn.base import BaseEstimator, TransformerMixin\n",
    "from sklearn.exceptions import NotFittedError\n",
    "from sklearn.utils.estimator_checks import check_estimator\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.feature_extraction.text import TfidfVectorizer\n",
    "from sklearn.decomposition import TruncatedSVD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LSA_NDK_CORPUS_PATH = os.path.join(\n",
    "    ARTEFACTS_PATH, \"month.standard-ndk.data.pickle.gz\"\n",
    ")\n",
    "\n",
    "LSA_NDK_LABELS_PATH = os.path.join(\n",
    "    ARTEFACTS_PATH, \"month.standard-ndk.labels.pickle.gz\"\n",
    ")\n",
    "\n",
    "\n",
    "USE_GENSIM = False\n",
    "\n",
    "NGRAM_RANGE = (1, 3)\n",
    "SVD_COMPONENTS = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log(\"Loading corpus...\")\n",
    "corpus: T.Dict[str, T.List[str]] = load_pickle(LSA_NDK_CORPUS_PATH)\n",
    "log(\"Loaded corpus.\")\n",
    "\n",
    "log(\"Loading labels...\")\n",
    "labels: T.Dict[str, T.Dict[T.Optional[str], int]] = load_pickle(\n",
    "    LSA_NDK_LABELS_PATH\n",
    ")\n",
    "log(\"Loaded labels.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log(\"Organising data and labels...\")\n",
    "x: T.List[T.Tuple[str, ...]] = []\n",
    "y_tags: T.List[str] = []\n",
    "\n",
    "for tag, words in corpus.items():\n",
    "    x.append(tuple(words))\n",
    "    y_tags.append(tag)\n",
    "\n",
    "if SAMPLE_LIMIT is not None:\n",
    "    log(f\"Warning: Using a sample limit: '{SAMPLE_LIMIT}'.\", file=sys.stderr)\n",
    "    x = x[SAMPLE_LIMIT]\n",
    "    y_tags = y_tags[SAMPLE_LIMIT]\n",
    "\n",
    "if TRIM_MAJORITY_CLASS:\n",
    "    log(\"Trimming...\")\n",
    "    x, y_tags, top_classes = trimmer(x, y_tags, labels)\n",
    "else:\n",
    "    top_classes = TOP_CLASSES\n",
    "\n",
    "binariser, y = build_binariser(y_tags, labels)\n",
    "log(\"Organised data and labels.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lsa_ndk_classification(fold_index: int) -> Summary:\n",
    "    if USE_GENSIM:\n",
    "        try:\n",
    "            LatentSemanticAnalysis  # type: ignore\n",
    "        except NameError:\n",
    "            raise NotImplementedError(\n",
    "                \"Run the domain knowledge version's \"\n",
    "                \"Gensim implementation block first.\"\n",
    "            )\n",
    "\n",
    "        log(\"LSA-NDK via Gensim.\")\n",
    "        transformer = LatentSemanticAnalysis(\n",
    "            ngram_range=NGRAM_RANGE, n_components=SVD_COMPONENTS\n",
    "        )\n",
    "    else:\n",
    "        log(\"LSA-NDK via scikit-learn.\")\n",
    "        transformer = make_pipeline(\n",
    "            TfidfVectorizer(\n",
    "                ngram_range=NGRAM_RANGE,\n",
    "                lowercase=False,\n",
    "                tokenizer=lambda x: x,\n",
    "                preprocessor=lambda x: x,\n",
    "                token_pattern=None,  # type: ignore\n",
    "                use_idf=True,\n",
    "                sublinear_tf=True,\n",
    "            ),\n",
    "            TruncatedSVD(\n",
    "                algorithm=\"arpack\",\n",
    "                n_components=SVD_COMPONENTS,\n",
    "                random_state=SEED,\n",
    "            ),\n",
    "            verbose=True,\n",
    "        )\n",
    "\n",
    "    return training_and_testing_scores(\n",
    "        \"LSA-NDK\",\n",
    "        fold_index,\n",
    "        transformer,\n",
    "        build_classifier(),\n",
    "    )\n",
    "\n",
    "\n",
    "lsa_ndk_fold_summaries = cross_validator(\n",
    "    x, y, binariser, lsa_ndk_classification, simultaneous_folds=1\n",
    ")\n",
    "\n",
    "print_summaries(lsa_ndk_fold_summaries, top_classes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Latent Semantic Analysis (best hyperparameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gensim\n",
    "\n",
    "from sklearn.base import BaseEstimator, TransformerMixin\n",
    "from sklearn.exceptions import NotFittedError\n",
    "from sklearn.utils.estimator_checks import check_estimator\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.feature_extraction.text import TfidfVectorizer\n",
    "from sklearn.decomposition import TruncatedSVD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LSA_NDK_CORPUS_PATH = os.path.join(\n",
    "    ARTEFACTS_PATH, \"month.standard-ndk.data.pickle.gz\"\n",
    ")\n",
    "\n",
    "LSA_NDK_LABELS_PATH = os.path.join(\n",
    "    ARTEFACTS_PATH, \"month.standard-ndk.labels.pickle.gz\"\n",
    ")\n",
    "\n",
    "\n",
    "USE_GENSIM = False\n",
    "\n",
    "NGRAM_RANGE = (1, 1)\n",
    "SVD_COMPONENTS = 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log(\"Loading corpus...\")\n",
    "corpus: T.Dict[str, T.List[str]] = load_pickle(LSA_NDK_CORPUS_PATH)\n",
    "log(\"Loaded corpus.\")\n",
    "\n",
    "log(\"Loading labels...\")\n",
    "labels: T.Dict[str, T.Dict[T.Optional[str], int]] = load_pickle(\n",
    "    LSA_NDK_LABELS_PATH\n",
    ")\n",
    "log(\"Loaded labels.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log(\"Organising data and labels...\")\n",
    "x: T.List[T.Tuple[str, ...]] = []\n",
    "y_tags: T.List[str] = []\n",
    "\n",
    "for tag, words in corpus.items():\n",
    "    x.append(tuple(words))\n",
    "    y_tags.append(tag)\n",
    "\n",
    "if SAMPLE_LIMIT is not None:\n",
    "    log(f\"Warning: Using a sample limit: '{SAMPLE_LIMIT}'.\", file=sys.stderr)\n",
    "    x = x[SAMPLE_LIMIT]\n",
    "    y_tags = y_tags[SAMPLE_LIMIT]\n",
    "\n",
    "if TRIM_MAJORITY_CLASS:\n",
    "    log(\"Trimming...\")\n",
    "    x, y_tags, top_classes = trimmer(x, y_tags, labels)\n",
    "else:\n",
    "    top_classes = TOP_CLASSES\n",
    "\n",
    "binariser, y = build_binariser(y_tags, labels)\n",
    "log(\"Organised data and labels.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lsa_best_classification(fold_index: int) -> Summary:\n",
    "    if USE_GENSIM:\n",
    "        try:\n",
    "            LatentSemanticAnalysis  # type: ignore\n",
    "        except NameError:\n",
    "            raise NotImplementedError(\n",
    "                \"Run the domain knowledge version's \"\n",
    "                \"Gensim implementation block first.\"\n",
    "            )\n",
    "\n",
    "        log(\"LSA-BEST via Gensim.\")\n",
    "        transformer = LatentSemanticAnalysis(\n",
    "            ngram_range=NGRAM_RANGE, n_components=SVD_COMPONENTS\n",
    "        )\n",
    "    else:\n",
    "        log(\"LSA-BEST via scikit-learn.\")\n",
    "        transformer = make_pipeline(\n",
    "            TfidfVectorizer(\n",
    "                ngram_range=NGRAM_RANGE,\n",
    "                lowercase=False,\n",
    "                tokenizer=lambda x: x,\n",
    "                preprocessor=lambda x: x,\n",
    "                token_pattern=None,  # type: ignore\n",
    "                use_idf=True,\n",
    "                sublinear_tf=True,\n",
    "            ),\n",
    "            TruncatedSVD(\n",
    "                algorithm=\"arpack\",\n",
    "                n_components=SVD_COMPONENTS,\n",
    "                random_state=SEED,\n",
    "            ),\n",
    "            verbose=True,\n",
    "        )\n",
    "\n",
    "    return training_and_testing_scores(\n",
    "        \"LSA-BEST\",\n",
    "        fold_index,\n",
    "        transformer,\n",
    "        build_classifier(),\n",
    "    )\n",
    "\n",
    "\n",
    "lsa_best_fold_summaries = cross_validator(\n",
    "    x, y, binariser, lsa_best_classification, simultaneous_folds=1\n",
    ")\n",
    "\n",
    "print_summaries(lsa_best_fold_summaries, top_classes)"
   ]
  }
 ],
 "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.5"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}