reference annotater improvements and hyperparam search poc

3 files changed, 72 insertions, 13 deletions

diff --git a/src/sec_certs/model/references/annotator.py b/src/sec_certs/model/references/annotator.py
index 0b971a8b..42615528 100644
--- a/src/sec_certs/model/references/annotator.py
+++ b/src/sec_certs/model/references/annotator.py
@@ -24,7 +24,10 @@ class ReferenceAnnotator:
 
     _model: Any
     _label_mapping: dict[int, str]
+    _soft_voting_power: int = 2
+    _use_analytical_rule_name_similarity: bool = True
 
+    # TODO: This does not load hyperparameters, only the model and label mapping
     @classmethod
     def from_pretrained(cls, model_dir: str | Path) -> ReferenceAnnotator:
         """
@@ -43,6 +46,7 @@ class ReferenceAnnotator:
 
         return cls(model, label_mapping)
 
+    # TODO: This does not save hyperparameters, only the model and label mapping
     def save_pretrained(self, model_dir: str | Path):
         """
         Will dump _model and _label_mapping into a directory.
@@ -74,7 +78,7 @@ class ReferenceAnnotator:
         3. Sum probabilities for each label
         4. softmax
         """
-        return softmax(np.power(self._model.predict_proba(sample, as_numpy=True), 2).sum(axis=0))
+        return softmax(np.power(self._model.predict_proba(sample, as_numpy=True), self._soft_voting_power).sum(axis=0))
 
     def predict_df(self, df: pd.DataFrame) -> pd.DataFrame:
         """
@@ -84,6 +88,15 @@ class ReferenceAnnotator:
         y_proba = self.predict_proba(df.segments)
         df_new["y_proba"] = y_proba
         df_new["y_pred"] = df_new.y_proba.map(lambda x: self._label_mapping[int(np.argmax(x))])
+
+        if self._use_analytical_rule_name_similarity:
+            df_new.loc[
+                (df_new.name_similarity_stripped_version == 100)
+                & (df_new.name_len_diff < 5)
+                & ((df_new.y_pred != "RECERTIFICATION") & (df_new.y_pred != "PREVIOUS_VERSION")),
+                ["y_pred"],
+            ] = "PREVIOUS_VERSION"
+
         df_new["correct"] = df_new.apply(
             lambda row: row["y_pred"] == row["label"] if not pd.isnull(row["label"]) else np.NaN, axis=1
         )
diff --git a/src/sec_certs/model/references/annotator_trainer.py b/src/sec_certs/model/references/annotator_trainer.py
index e0ae01e5..d2c53636 100644
--- a/src/sec_certs/model/references/annotator_trainer.py
+++ b/src/sec_certs/model/references/annotator_trainer.py
@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 import logging
+from functools import partial
 from typing import Callable, Literal
 
 import pandas as pd
 from datasets import ClassLabel, Dataset, Features, NamedSplit, Value
 from sentence_transformers.losses import CosineSimilarityLoss
 from setfit import SetFitModel, SetFitTrainer
+from sklearn.metrics import f1_score
 
 from sec_certs.model.references.annotator import ReferenceAnnotator
 from sec_certs.utils.nlp import prepare_reference_annotations_df
@@ -20,12 +22,28 @@ class ReferenceAnnotatorTrainer:
         train_dataset: pd.DataFrame,
         eval_dataset: pd.DataFrame,
         metric: Callable,
+        use_analytical_rule_name_similarity: bool = True,
+        n_iterations: int = 20,
+        n_epochs: int = 1,
+        batch_size: int = 16,
+        segmenter_metric: Literal["accuracy", "f1"] = "accuracy",
+        ensemble_soft_voting_power: int = 2,
     ):
         self._train_dataset = train_dataset
         self._eval_dataset = eval_dataset
         self._metric = metric
+        self.use_analytical_rule_name_similarity = use_analytical_rule_name_similarity
+        self.n_iterations = n_iterations
+        self.n_epochs = n_epochs
+        self.batch_size = batch_size
+        self.segmenter_metric = segmenter_metric
+        self.ensemble_soft_voting_power = ensemble_soft_voting_power
+
         self._model, self._trainer, self.label_mapping = self._init_trainer()
-        self.clf = ReferenceAnnotator(self._model, self.label_mapping)
+
+        self.clf = ReferenceAnnotator(
+            self._model, self.label_mapping, self.ensemble_soft_voting_power, self.use_analytical_rule_name_similarity
+        )
 
     @classmethod
     def from_df(
@@ -33,6 +51,12 @@ class ReferenceAnnotatorTrainer:
         df: pd.DataFrame,
         metric: Callable,
         mode: Literal["training", "production"] = "training",
+        use_analytical_rule_name_similarity: bool = True,
+        n_iterations: int = 20,
+        n_epochs: int = 1,
+        batch_size: int = 16,
+        segmenter_metric: Literal["accuracy", "f1"] = "accuracy",
+        ensemble_soft_voting_power: int = 2,
     ):
         df = prepare_reference_annotations_df(df)
         dataset_generation_method = {
@@ -41,7 +65,17 @@ class ReferenceAnnotatorTrainer:
         }
 
         train_dataset, eval_dataset = dataset_generation_method[mode](df)
-        return cls(train_dataset, eval_dataset, metric)
+        return cls(
+            train_dataset,
+            eval_dataset,
+            metric,
+            use_analytical_rule_name_similarity,
+            n_iterations,
+            n_epochs,
+            batch_size,
+            segmenter_metric,
+            ensemble_soft_voting_power,
+        )
 
     @staticmethod
     def split_df_for_training(df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
@@ -58,25 +92,37 @@ class ReferenceAnnotatorTrainer:
         model = SetFitModel.from_pretrained("paraphrase-multilingual-mpnet-base-v2")
         # model = SetFitModel.from_pretrained("all-mpnet-base-v2")
 
-        internal_train_dataset = self._get_hugging_face_datasets_from_df(self._train_dataset, "train")
-        internal_validation_dataset = self._get_hugging_face_datasets_from_df(self._eval_dataset, "validation")
+        train_dataset_relevant_cols = self._train_dataset[["dgst", "referenced_cert_id", "segments", "label"]]
+        eval_dataset_relevant_cols = self._eval_dataset[["dgst", "referenced_cert_id", "segments", "label"]]
+        internal_train_dataset = self._get_hugging_face_datasets_from_df(train_dataset_relevant_cols, "train")
+        internal_validation_dataset = self._get_hugging_face_datasets_from_df(eval_dataset_relevant_cols, "validation")
+
+        # Align labels alphabetically
+        labels_alphabetically = sorted(internal_train_dataset.features["label"].names)
+        label2id = {label: index for index, label in enumerate(labels_alphabetically)}
+        internal_train_dataset = internal_train_dataset.align_labels_with_mapping(label2id, "label")
+        internal_validation_dataset = internal_validation_dataset.align_labels_with_mapping(label2id, "label")
+
+        if self.segmenter_metric == "accuracy":
+            metric_to_use = "accuracy"
+        else:
+            metric_to_use = partial(f1_score, average="weighted", zero_division=0)
 
         trainer = SetFitTrainer(
             model=model,
             train_dataset=internal_train_dataset,
             eval_dataset=internal_validation_dataset,
             loss_class=CosineSimilarityLoss,
-            metric=self._metric,
-            batch_size=16,
-            num_iterations=40,  # The number of text pairs to generate for contrastive learning
-            num_epochs=1,  # The number of epochs to use for contrastive learning
+            metric=metric_to_use,
+            batch_size=self.batch_size,
+            num_iterations=self.n_iterations,  # The number of text pairs to generate for contrastive learning
+            num_epochs=self.n_epochs,  # The number of epochs to use for contrastive learning
             column_mapping={
                 "segment": "text",
                 "label": "label",
             },  # Map dataset columns to text/label expected by trainer
         )
-        label_mapping = {index: x for index, x in enumerate(internal_train_dataset.features["label"].names)}
-        return model, trainer, label_mapping
+        return model, trainer, {index: label for label, index in label2id.items()}
 
     @staticmethod
     def _get_hugging_face_datasets_from_df(df: pd.DataFrame, split: NamedSplit) -> Dataset:
diff --git a/src/sec_certs/utils/nlp.py b/src/sec_certs/utils/nlp.py
index 4c136ce9..a48e6d73 100644
--- a/src/sec_certs/utils/nlp.py
+++ b/src/sec_certs/utils/nlp.py
@@ -9,8 +9,8 @@ from sklearn.metrics import precision_score, recall_score
 
 def prec_recall_metric(y_pred, y_true):
     return {
-        "precision": precision_score(y_true, y_pred, zero_division="warn", average="micro"),
-        "recall": recall_score(y_true, y_pred, zero_division="warn", average="micro"),
+        "precision": precision_score(y_true, y_pred, zero_division="warn", average="weighted"),
+        "recall": recall_score(y_true, y_pred, zero_division="warn", average="weighted"),
     }