# ============================================================================ # 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() ]