source/train_embedding.py

from __future__ import annotations

import json

import joblib
import numpy as np
import pandas as pd
import trackio
from sentence_transformers import SentenceTransformer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
)
from train import ARTIFACT_DIR, DATA_DIR

BASE_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
OUTPUT_DIR = ARTIFACT_DIR.parent / "protocol-guardian-minilm-linear"


def load(name: str) -> tuple[list[str], np.ndarray]:
    frame = pd.read_parquet(DATA_DIR / f"{name}.parquet")
    return (
        frame["text"].tolist(),
        frame["label"].to_numpy(dtype=np.int64, copy=True),
    )


def metrics(labels: np.ndarray, predictions: np.ndarray) -> dict:
    return {
        "accuracy": float(accuracy_score(labels, predictions)),
        "precision": float(precision_score(labels, predictions, zero_division=0)),
        "recall": float(recall_score(labels, predictions, zero_division=0)),
        "f1": float(f1_score(labels, predictions, zero_division=0)),
        "confusion_matrix": confusion_matrix(labels, predictions).tolist(),
        "examples": len(labels),
    }


def main() -> None:
    train_text, train_labels = load("train")
    validation_text, validation_labels = load("validation")
    test_text, test_labels = load("test")
    encoder = SentenceTransformer(BASE_MODEL, device="cpu")
    embeddings = {
        "train": encoder.encode(
            train_text,
            batch_size=128,
            normalize_embeddings=True,
            show_progress_bar=True,
        ),
        "validation": encoder.encode(
            validation_text,
            batch_size=128,
            normalize_embeddings=True,
            show_progress_bar=True,
        ),
        "test": encoder.encode(
            test_text,
            batch_size=128,
            normalize_embeddings=True,
            show_progress_bar=True,
        ),
    }
    trackio.init(
        project="protocol-guardian",
        name="minilm-embedding-linear-v1",
        config={
            "base_model": BASE_MODEL,
            "embedding_dimensions": embeddings["train"].shape[1],
            "train_examples": len(train_labels),
            "validation_template_family": "held-out from train",
            "test_template_family": "held-out from train and validation",
        },
    )
    candidates = []
    for regularization in [0.03, 0.1, 0.3, 1.0, 3.0, 10.0]:
        classifier = LogisticRegression(
            C=regularization,
            class_weight="balanced",
            max_iter=2000,
            random_state=2026,
        )
        classifier.fit(embeddings["train"], train_labels)
        predictions = classifier.predict(embeddings["validation"])
        score = f1_score(validation_labels, predictions)
        candidates.append((score, regularization, classifier))
        trackio.log(
            {
                "regularization_c": regularization,
                "validation_f1": score,
                "validation_accuracy": accuracy_score(
                    validation_labels,
                    predictions,
                ),
            }
        )
    _, best_c, classifier = max(candidates, key=lambda item: item[0])
    validation_predictions = classifier.predict(embeddings["validation"])
    test_predictions = classifier.predict(embeddings["test"])
    results = {
        "model": "Protocol Guardian MiniLM Linear",
        "base_model": BASE_MODEL,
        "embedding_dimensions": int(embeddings["train"].shape[1]),
        "linear_head_parameters": int(classifier.coef_.size + classifier.intercept_.size),
        "selected_c": best_c,
        "validation": metrics(validation_labels, validation_predictions),
        "test": metrics(test_labels, test_predictions),
        "test_split": "1200 examples from a third, unseen template family",
    }
    trackio.log(
        {
            "test_accuracy": results["test"]["accuracy"],
            "test_f1": results["test"]["f1"],
            "test_recall": results["test"]["recall"],
        }
    )
    trackio.finish()
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    joblib.dump(
        {
            "classifier": classifier,
            "base_model": BASE_MODEL,
            "labels": {0: "ROUTINE", 1: "HAZARDOUS"},
        },
        OUTPUT_DIR / "classifier.joblib",
        compress=3,
    )
    (OUTPUT_DIR / "training_summary.json").write_text(
        json.dumps(results, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()