workbench_classifiers.py

# ============================================================================
# workbench_classifiers.py — Reference-sentence cosine-similarity classifier
# ============================================================================
"""Reference-based cosine classification — no training, no train/test split.

Purpose
-------
Researcher provides a small set of REFERENCE sentences with class tags.
Each test sentence is then classified by cosine similarity in MiniLM
embedding space, against either:
  - the nearest individual reference sentence (1-NN by cosine), or
  - the nearest class prototype (mean of references per class).

This is the workbench's "supervised" path: vector-only, no scikit-learn
fit step, fully interpretable (every prediction has a "nearest reference"
the researcher can read).

A secondary block — six standard classifiers on an 80/20 train/test
split — is exposed via train_six_classifiers() for users who explicitly
want that comparison.

Inputs
------
    refs       : list[str] — reference sentences
    labels     : list[str] — class label per reference (parallel to refs)
    test_sents : list[str] — corpus to classify

Outputs
-------
    PredictionResult dataclass — per-sentence prediction, similarity, and
    nearest-reference text.

Side effects
------------
    None. Pure functions.

Contract
--------
    classify_by_reference(refs, ref_labels, test_sents, mode, embedding_*)
    train_six_classifiers(X_train, y_train, X_test, y_test) — optional.
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Any

import numpy as np

import providers


# ----------------------------------------------------------------
# Mode dispatch table — replaces an if/elif ladder. Every mode key
# names a function that takes (test_norm, ref_norm, ref_labels) and
# returns (predictions, top_sims, nearest_indices_for_each_test).
# ----------------------------------------------------------------
def _predict_nearest_reference(
    test_norm: np.ndarray, ref_norm: np.ndarray, ref_labels: list[str],
) -> tuple[list[str], list[float], list[int]]:
    """1-NN by cosine: each test sentence inherits its nearest reference's label."""
    sims = test_norm @ ref_norm.T          # (n_test, n_refs)
    idx = np.argmax(sims, axis=1)           # nearest reference index per test
    top_sims = sims[np.arange(sims.shape[0]), idx].tolist()
    preds = [ref_labels[i] for i in idx]
    return preds, top_sims, idx.tolist()


def _predict_class_prototype(
    test_norm: np.ndarray, ref_norm: np.ndarray, ref_labels: list[str],
) -> tuple[list[str], list[float], list[int]]:
    """Mean-embedding-per-class then nearest prototype.

    The "nearest reference" we report for transparency is the reference
    inside the predicted class that is closest to the test sentence, so
    the researcher can still read a concrete exemplar.
    """
    classes = sorted(set(ref_labels))
    # Prototype per class — mean of normalised reference vectors, re-normalised.
    protos = np.stack([
        ref_norm[[i for i, lab in enumerate(ref_labels) if lab == c]].mean(axis=0)
        for c in classes
    ])
    protos = protos / (np.linalg.norm(protos, axis=1, keepdims=True) + 1e-12)

    sims_to_proto = test_norm @ protos.T   # (n_test, n_classes)
    proto_idx = np.argmax(sims_to_proto, axis=1)
    preds = [classes[i] for i in proto_idx]
    top_sims = sims_to_proto[np.arange(sims_to_proto.shape[0]), proto_idx].tolist()

    # For each test, find the nearest reference INSIDE its predicted class
    # so the per-row "nearest reference" remains readable.
    sims_to_refs = test_norm @ ref_norm.T  # (n_test, n_refs)
    nearest = []
    for t_idx, pred_class in enumerate(preds):
        in_class = [i for i, lab in enumerate(ref_labels) if lab == pred_class]
        sub = sims_to_refs[t_idx, in_class]
        nearest.append(in_class[int(np.argmax(sub))])
    return preds, top_sims, nearest


_MODE_TABLE = {
    "nearest_reference": _predict_nearest_reference,
    "class_prototype":   _predict_class_prototype,
}


# ----------------------------------------------------------------
# Result dataclass
# ----------------------------------------------------------------
@dataclass
class PredictionRow:
    """One test sentence's classification result, ready for the UI."""

    sentence: str
    predicted_class: str
    similarity: float
    nearest_reference: str


@dataclass
class PredictionResult:
    """End-to-end classification output."""

    rows: list[PredictionRow] = field(default_factory=list)
    mode: str = ""
    n_classes: int = 0
    n_references: int = 0
    n_test: int = 0


# ----------------------------------------------------------------
# Embedding helper
# ----------------------------------------------------------------
def _embed(
    texts: list[str], embedding_provider: str, embedding_key: str,
) -> np.ndarray:
    """Vectorise via providers.embed_texts; raise with clear framing on error."""
    try:
        vecs = providers.embed_texts(
            texts, provider_name=embedding_provider, api_key=embedding_key,
        )
    except Exception as exc:
        raise RuntimeError(
            f"Embedding step failed (provider={embedding_provider}): {exc}"
        ) from exc
    return np.asarray(vecs, dtype=np.float32)


def _l2_normalise(vecs: np.ndarray) -> np.ndarray:
    """Row-wise L2 normalisation so dot product == cosine similarity."""
    return vecs / (np.linalg.norm(vecs, axis=1, keepdims=True) + 1e-12)


# ----------------------------------------------------------------
# PUBLIC: reference-sentence cosine classification
# ----------------------------------------------------------------
def classify_by_reference(
    refs: list[str],
    ref_labels: list[str],
    test_sents: list[str],
    mode: str = "nearest_reference",
    embedding_provider: str = "MiniLM (local)",
    embedding_key: str = "",
) -> PredictionResult:
    """Cosine-similarity classification, no training step.

    Args:
        refs:               Labelled exemplar sentences.
        ref_labels:         Class label for each reference (parallel list).
        test_sents:         Sentences to classify.
        mode:               "nearest_reference" or "class_prototype".
        embedding_provider: Workbench embedding provider key.
        embedding_key:      Optional API key (empty for local MiniLM).

    Returns:
        PredictionResult with one PredictionRow per test sentence.

    Raises:
        ValueError:    on shape mismatches between refs and ref_labels,
                       on unknown mode, or on empty references.
        RuntimeError:  bubbled up from the embedding step.

    Example:
        >>> r = classify_by_reference(
        ...     refs=["card stolen overseas", "standard grocery purchase"],
        ...     ref_labels=["fraud", "normal"],
        ...     test_sents=["unauthorised charge in Madrid"],
        ... )
        >>> r.rows[0].predicted_class
        'fraud'
    """
    if len(refs) != len(ref_labels):
        raise ValueError(
            f"refs and ref_labels length mismatch: {len(refs)} vs {len(ref_labels)}"
        )
    if not refs:
        raise ValueError("No reference sentences provided.")
    if mode not in _MODE_TABLE:
        raise ValueError(
            f"Unknown classification mode {mode!r}. "
            f"Valid: {list(_MODE_TABLE)}"
        )
    if not test_sents:
        return PredictionResult(
            mode=mode,
            n_classes=len(set(ref_labels)),
            n_references=len(refs),
            n_test=0,
        )

    ref_vecs = _embed(refs, embedding_provider, embedding_key)
    test_vecs = _embed(test_sents, embedding_provider, embedding_key)
    ref_norm = _l2_normalise(ref_vecs)
    test_norm = _l2_normalise(test_vecs)

    predict_fn = _MODE_TABLE[mode]
    preds, sims, nearest_idx = predict_fn(test_norm, ref_norm, ref_labels)

    rows = [
        PredictionRow(
            sentence=test_sents[i],
            predicted_class=preds[i],
            similarity=float(sims[i]),
            nearest_reference=refs[nearest_idx[i]],
        )
        for i in range(len(test_sents))
    ]

    return PredictionResult(
        rows=rows,
        mode=mode,
        n_classes=len(set(ref_labels)),
        n_references=len(refs),
        n_test=len(test_sents),
    )


# ============================================================================
# OPTIONAL: six standard classifiers on an 80/20 split
# ============================================================================
# Provided only as a comparison block. Off by default in the UI.
# Each entry is a name + a zero-argument factory that returns a fresh,
# unfit estimator. The factory pattern keeps the table declarative and
# avoids importing scikit-learn at module load time when the comparison
# block is never used.
# ============================================================================

_CLASSIFIER_FACTORIES: dict[str, Any] = {
    "LogisticRegression": lambda: __import__(
        "sklearn.linear_model", fromlist=["LogisticRegression"]
    ).LogisticRegression(max_iter=1000, n_jobs=-1),
    "LinearSVC": lambda: __import__(
        "sklearn.svm", fromlist=["LinearSVC"]
    ).LinearSVC(),
    "KNN_cosine_k5": lambda: __import__(
        "sklearn.neighbors", fromlist=["KNeighborsClassifier"]
    ).KNeighborsClassifier(n_neighbors=5, metric="cosine"),
    "RandomForest": lambda: __import__(
        "sklearn.ensemble", fromlist=["RandomForestClassifier"]
    ).RandomForestClassifier(n_estimators=200, n_jobs=-1, random_state=42),
    "GradientBoosting": lambda: __import__(
        "sklearn.ensemble", fromlist=["GradientBoostingClassifier"]
    ).GradientBoostingClassifier(random_state=42),
    "MLPClassifier": lambda: __import__(
        "sklearn.neural_network", fromlist=["MLPClassifier"]
    ).MLPClassifier(hidden_layer_sizes=(128,), max_iter=300, random_state=42),
}


@dataclass
class ClassifierScore:
    """One classifier's score row for the comparison block."""

    name: str
    accuracy: float
    macro_f1: float
    fit_seconds: float


def _fit_one(
    name: str, factory: Any,
    X_train: np.ndarray, y_train: np.ndarray,
    X_test: np.ndarray, y_test: np.ndarray,
) -> ClassifierScore:
    """Fit one classifier, score it, return a single ClassifierScore row.

    Errors during fit are caught and returned as a row with NaN scores so
    the comparison table can still render the other classifiers.
    """
    from sklearn.metrics import accuracy_score, f1_score

    try:
        clf = factory()
        t0 = time.perf_counter()
        clf.fit(X_train, y_train)
        fit_s = time.perf_counter() - t0
        y_pred = clf.predict(X_test)
        return ClassifierScore(
            name=name,
            accuracy=float(accuracy_score(y_test, y_pred)),
            macro_f1=float(f1_score(y_test, y_pred, average="macro", zero_division=0)),
            fit_seconds=float(fit_s),
        )
    except Exception as exc:
        return ClassifierScore(
            name=f"{name} (failed: {exc})",
            accuracy=float("nan"),
            macro_f1=float("nan"),
            fit_seconds=float("nan"),
        )


def train_six_classifiers(
    sentences: list[str],
    labels: list[str],
    train_fraction: float = 0.8,
    embedding_provider: str = "MiniLM (local)",
    embedding_key: str = "",
    random_state: int = 42,
) -> list[ClassifierScore]:
    """Optional comparison block — train six standard classifiers and score them.

    Args:
        sentences:          Labelled corpus.
        labels:             Class label per sentence.
        train_fraction:     0.8 means 80/20 train/test split.
        embedding_provider: Workbench embedding provider key.
        embedding_key:      Optional API key (empty for local MiniLM).
        random_state:       For deterministic split.

    Returns:
        A list of ClassifierScore — one per classifier in
        _CLASSIFIER_FACTORIES, in dispatch-table order.

    Raises:
        ValueError:   when sentences/labels lengths don't match.
        RuntimeError: bubbled up from the embedding step.
    """
    if len(sentences) != len(labels):
        raise ValueError(
            f"sentences and labels length mismatch: {len(sentences)} vs {len(labels)}"
        )

    from sklearn.model_selection import train_test_split

    X = _embed(sentences, embedding_provider, embedding_key)
    y = np.asarray(labels)

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, train_size=train_fraction, random_state=random_state, stratify=y,
    )

    return [
        _fit_one(name, factory, X_train, y_train, X_test, y_test)
        for name, factory in _CLASSIFIER_FACTORIES.items()
    ]