refactoring here and there

3 files changed, 78 insertions, 60 deletions

diff --git a/src/sec_certs/model/references_nlp/evaluation.py b/src/sec_certs/model/references_nlp/evaluation.py
index dc9b4b13..ed3b2700 100644
--- a/src/sec_certs/model/references_nlp/evaluation.py
+++ b/src/sec_certs/model/references_nlp/evaluation.py
@@ -15,40 +15,30 @@ logger = logging.getLogger(__name__)
 
 def evaluate_model(
     clf: DummyClassifier | CatBoostClassifier,
-    df_eval: pd.DataFrame,
+    x_eval: np.ndarray,
+    y_eval: np.ndarray,
     feature_cols: list[str],
     output_path: Path | None = None,
 ):
     logger.info("Evaluating model.")
-    x_eval = np.vstack(df_eval[feature_cols].values)
     y_pred = clf.predict(x_eval)
 
-    df_eval["y_pred"] = y_pred
-
-    df_eval.loc[df_eval.lang_matches_recertification, ["y_pred"]] = "PREVIOUS_VERSION"
-    df_eval.loc[
-        (df_eval.lang_token_set_ratio == 100)
-        & (df_eval.lang_len_difference < 5)
-        & (df_eval.y_pred != "PREVIOUS_VERSION"),
-        ["y_pred"],
-    ] = "PREVIOUS_VERSION"
-
-    print(classification_report(df_eval.label.values, df_eval.y_pred.values))
-    print(f"Balanced accuracy score: {balanced_accuracy_score(df_eval.label.values, df_eval.y_pred.values)}")
+    print(classification_report(y_eval, y_pred))
+    print(f"Balanced accuracy score: {balanced_accuracy_score(y_eval, y_pred)}")
 
     fig = ConfusionMatrixDisplay.from_predictions(
-        df_eval.label.values,
-        df_eval.y_pred.values,
+        y_eval,
+        y_pred,
         xticks_rotation=90,
     )
 
     if output_path:
-        report_dict = classification_report(df_eval.label.values, df_eval.y_pred.values, output_dict=True)
+        report_dict = classification_report(y_eval, y_pred, output_dict=True)
         report_df = pd.DataFrame(report_dict).transpose()
         report_df.to_csv(output_path / "classification_report.csv")
-        fig.figure_.savefig(output_path / "confusion_matrix.png")
+        fig.figure_.savefig(output_path / "confusion_matrix.png", bbox_inches="tight")
         with Path(output_path / "balanced_accuracy_score.txt").open("w") as handle:
-            handle.write(str(balanced_accuracy_score(df_eval.label.values, df_eval.y_pred.values)))
+            handle.write(str(balanced_accuracy_score(y_eval, y_pred)))
 
     if isinstance(clf, CatBoostClassifier):
         feature_importance = clf.get_feature_importance()
diff --git a/src/sec_certs/model/references_nlp/feature_extraction.py b/src/sec_certs/model/references_nlp/feature_extraction.py
index 9e62879e..996972d7 100644
--- a/src/sec_certs/model/references_nlp/feature_extraction.py
+++ b/src/sec_certs/model/references_nlp/feature_extraction.py
@@ -140,7 +140,7 @@ def extract_segments(
 
 def _build_transformer_embeddings(
     segments: pd.DataFrame, mode: REF_ANNOTATION_MODES, model_path: Path | None = None
-) -> pd.DataFrame:
+) -> tuple[pd.DataFrame, ReferenceAnnotator]:
     should_save_model = model_path is not None
     annotator = None
     logger.info("Building transformer embeddings.")
@@ -204,12 +204,12 @@ def build_embeddings(
     mode: REF_ANNOTATION_MODES,
     method: REF_EMBEDDING_METHOD,
     model_path: Path | None = None,
-) -> pd.DataFrame:
-    return (
-        _build_transformer_embeddings(segments, mode, model_path)
-        if method == "transformer"
-        else _build_tf_idf_embeddings(segments, mode)
-    )
+) -> tuple[pd.DataFrame, ReferenceAnnotator | None]:
+    if method == "transformer":
+        return _build_transformer_embeddings(segments, mode, model_path)
+    if method == "tf_idf":
+        return _build_tf_idf_embeddings(segments, mode), None
+    raise ValueError(f"Unknown embedding method {method}")
 
 
 def extract_language_features(df: pd.DataFrame, cc_dset: CCDataset) -> pd.DataFrame:
@@ -521,14 +521,9 @@ def extract_geometrical_features(df: pd.DataFrame) -> pd.DataFrame:
     return pd.concat([df, df_features_pca, df_features_umap], axis=1)
 
 
-def get_data_for_clf(
-    df: pd.DataFrame,
-    mode: REF_ANNOTATION_MODES,
-    use_pca: bool = True,
-    use_umap: bool = True,
-    use_lang: bool = True,
-    use_pred: bool = True,
-) -> tuple[np.ndarray, np.ndarray, pd.DataFrame | None, list[str]]:
+def _choose_feature_columns(
+    df: pd.DataFrame, use_pca: bool = True, use_umap: bool = True, use_lang: bool = True, use_pred: bool = True
+) -> list[str]:
     feature_columns = []
     if not use_pca and not use_umap and not use_lang and not use_pred:
         raise ValueError("At least one of PCA, UMAP or language features must be used.")
@@ -540,7 +535,10 @@ def get_data_for_clf(
         feature_columns.extend([x for x in df.columns if x.startswith("lang_")])
     if use_pred:
         feature_columns.extend([x for x in df.columns if x.startswith("pred_")])
+    return feature_columns
 
+
+def _split_df(df: pd.DataFrame, mode: REF_ANNOTATION_MODES) -> tuple[pd.DataFrame, pd.DataFrame | None]:
     if mode == "training":
         train_df = df.loc[df.split == "train"].copy()
         eval_df = df.loc[df.split == "valid"].copy()
@@ -555,10 +553,36 @@ def get_data_for_clf(
         eval_df = None
     else:
         raise ValueError(f"Unknown mode {mode}")
+    return train_df, eval_df
 
-    return (
+
+def dataframe_to_training_arrays(
+    df: pd.DataFrame,
+    mode: REF_ANNOTATION_MODES,
+    use_pca: bool = True,
+    use_umap: bool = True,
+    use_lang: bool = True,
+    use_pred: bool = True,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray | None, np.ndarray | None, list[str]]:
+    feature_columns = _choose_feature_columns(df, use_pca, use_umap, use_lang, use_pred)
+    train_df, eval_df = _split_df(df.loc[df.label.notnull()].copy(), mode)
+
+    x_train, y_train = (
         np.vstack(train_df[feature_columns].values),
         train_df.label.values,
-        eval_df,
+    )
+    if eval_df is not None:
+        x_valid, y_valid = (
+            np.vstack(eval_df[feature_columns].values),
+            eval_df.label.values,
+        )
+    else:
+        x_valid, y_valid = None, None
+
+    return (
+        x_train,
+        y_train,
+        x_valid,
+        y_valid,
         feature_columns,
     )
diff --git a/src/sec_certs/model/references_nlp/training.py b/src/sec_certs/model/references_nlp/training.py
index d65bae59..e0abe94f 100644
--- a/src/sec_certs/model/references_nlp/training.py
+++ b/src/sec_certs/model/references_nlp/training.py
@@ -9,20 +9,23 @@ from sklearn.metrics import balanced_accuracy_score
 from sklearn.model_selection import KFold
 
 from sec_certs.constants import RANDOM_STATE, REF_ANNOTATION_MODES
-from sec_certs.model.references_nlp.feature_extraction import get_data_for_clf
+from sec_certs.model.references_nlp.feature_extraction import dataframe_to_training_arrays
 
 logger = logging.getLogger(__name__)
 
 
 def _train_model(
-    x_train,
-    y_train,
-    x_eval,
-    y_eval,
+    mode: REF_ANNOTATION_MODES,
+    x_train: np.ndarray,
+    y_train: np.ndarray,
+    x_eval: np.ndarray | None = None,
+    y_eval: np.ndarray | None = None,
     learning_rate: float = 0.03,
     depth: int = 6,
     l2_leaf_reg: float = 3,
 ):
+    # In production mode, we don't have early stopping on validation set. Hence we use number of iterations that worked during evaluation.
+    n_iters = 20 if mode == "production" else 1000
     clf = CatBoostClassifier(
         learning_rate=learning_rate,
         depth=depth,
@@ -30,10 +33,11 @@ def _train_model(
         task_type="GPU",
         devices=os.environ["CUDA_VISIBLE_DEVICES"],
         random_seed=RANDOM_STATE,
+        iterations=n_iters,
     )
 
     train_pool = Pool(x_train, y_train)
-    eval_pool = Pool(x_eval, y_eval)
+    eval_pool = Pool(x_eval, y_eval) if x_eval is not None else None
     clf.fit(
         train_pool,
         eval_set=eval_pool,
@@ -46,46 +50,46 @@ def _train_model(
 
 
 def train_model(
-    df: pd.DataFrame,
     mode: REF_ANNOTATION_MODES,
+    x_train: np.ndarray,
+    y_train: np.ndarray,
+    x_eval: np.ndarray | None = None,
+    y_eval: np.ndarray | None = None,
     train_baseline: bool = False,
-    use_pca: bool = True,
-    use_umap: bool = True,
-    use_lang: bool = True,
-    use_pred: bool = True,
     learning_rate: float = 0.079573,
     depth: int = 10,
     l2_leaf_reg: float = 7.303517,
-) -> tuple[DummyClassifier | CatBoostClassifier, pd.DataFrame, list[str]]:
-    logger.info(f"Training model for mode {mode}")
-    X_train, y_train, eval_df, feature_cols = get_data_for_clf(df, mode, use_pca, use_umap, use_lang, use_pred)
+) -> DummyClassifier | CatBoostClassifier:
+    logger.info(f"Training model with baselne={train_baseline}")
+
     if train_baseline:
         clf = DummyClassifier(random_state=RANDOM_STATE)
-        clf.fit(X_train, y_train)
+        clf.fit(x_train, y_train)
     else:
-        assert eval_df is not None
         clf = _train_model(
-            X_train,
+            mode,
+            x_train,
             y_train,
-            eval_df[feature_cols],
-            eval_df.label,
+            x_eval,
+            y_eval,
             learning_rate,
             depth,
             l2_leaf_reg,
         )
-
-    return clf, eval_df, feature_cols
+    return clf
 
 
-def cross_validate_model(df: pd.DataFrame, learning_rate: float = 0.03, depth: int = 6, l2_leaf_reg: int = 3) -> float:
+def cross_validate_model(
+    mode: REF_ANNOTATION_MODES, df: pd.DataFrame, learning_rate: float = 0.03, depth: int = 6, l2_leaf_reg: int = 3
+) -> float:
     logger.info("Cross-validating model")
-    X_train, y_train, _, _ = get_data_for_clf(df, "cross-validation", True, True, True, True)
+    X_train, y_train, _, _, _ = dataframe_to_training_arrays(df, "cross-validation", True, True, True, True)
     kf = KFold(n_splits=5, shuffle=True, random_state=RANDOM_STATE)
     scores = []
     for train_index, test_index in kf.split(X_train):
         X_train_, X_test_ = X_train[train_index], X_train[test_index]
         y_train_, y_test_ = y_train[train_index], y_train[test_index]
-        clf = _train_model(X_train_, y_train_, X_test_, y_test_, learning_rate, depth, l2_leaf_reg)
+        clf = _train_model(mode, X_train_, y_train_, X_test_, y_test_, learning_rate, depth, l2_leaf_reg)
         scores.append(balanced_accuracy_score(y_test_, clf.predict(X_test_)))
 
     return np.mean(scores)