Upload src/pubguard/train.py with huggingface_hub

src/pubguard/train.py

@@ -0,0 +1,280 @@
+"""
+Training pipeline for PubGuard classification heads.
+
+Trains lightweight linear classifiers on frozen model2vec embeddings.
+This follows the same paradigm as the openalex-topic-classifier:
+the expensive embedding is pre-computed once, and the classifier
+itself is a single matrix multiply — fast to train, fast to infer.
+
+Training strategy:
+    1. Load + cache model2vec embeddings for all training data
+    2. For each head, fit a logistic regression (sklearn) with
+       class-balanced weights and L2 regularisation
+    3. Export weights as .npz for the numpy-only inference path
+    4. Report per-class precision / recall / F1 on held-out split
+
+The entire pipeline trains in <5 minutes on CPU for ~50K samples,
+consistent with your existing toolchain.
+"""
+
+import json
+import logging
+import time
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import classification_report
+from sklearn.model_selection import train_test_split
+
+from .config import PubGuardConfig, DOC_TYPE_LABELS, AI_DETECT_LABELS, TOXICITY_LABELS
+from .classifier import LinearHead
+from .text import clean_text, extract_structural_features, N_STRUCTURAL_FEATURES
+
+logger = logging.getLogger(__name__)
+
+
+def load_ndjson(path: Path) -> Tuple[List[str], List[str]]:
+    """Load NDJSON file → (texts, labels)."""
+    texts, labels = [], []
+    with open(path) as f:
+        for line in f:
+            if line.strip():
+                row = json.loads(line)
+                texts.append(row["text"])
+                labels.append(row["label"])
+    return texts, labels
+
+
+def embed_texts(
+    texts: List[str],
+    config: PubGuardConfig,
+    cache_path: Optional[Path] = None,
+) -> np.ndarray:
+    """
+    Encode texts with model2vec, L2-normalise, return (N, D) float32.
+
+    Optionally caches to disk to avoid re-embedding on repeat runs.
+    """
+    if cache_path and cache_path.exists():
+        logger.info(f"Loading cached embeddings from {cache_path}")
+        return np.load(cache_path)
+
+    from model2vec import StaticModel
+
+    model_path = config.distilled_model_path
+    if model_path.exists():
+        model = StaticModel.from_pretrained(str(model_path))
+    else:
+        model = StaticModel.from_pretrained(config.model_name)
+        model_path.parent.mkdir(parents=True, exist_ok=True)
+        model.save_pretrained(str(model_path))
+
+    logger.info(f"Embedding {len(texts)} texts...")
+    cleaned = [clean_text(t, config.max_text_chars) for t in texts]
+    embeddings = model.encode(cleaned, show_progress_bar=True)
+
+    # L2-normalise
+    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
+    norms = np.where(norms == 0, 1, norms)
+    embeddings = (embeddings / norms).astype("float32")
+
+    if cache_path:
+        cache_path.parent.mkdir(parents=True, exist_ok=True)
+        np.save(cache_path, embeddings)
+        logger.info(f"Cached embeddings to {cache_path}")
+
+    return embeddings
+
+
+def compute_structural_features(texts: List[str]) -> np.ndarray:
+    """Compute structural features for all texts."""
+    feats = []
+    for t in texts:
+        cleaned = clean_text(t)
+        feat_dict = extract_structural_features(cleaned)
+        feats.append(list(feat_dict.values()))
+    return np.array(feats, dtype="float32")
+
+
+def train_head(
+    X_train: np.ndarray,
+    y_train: np.ndarray,
+    X_test: np.ndarray,
+    y_test: np.ndarray,
+    labels: List[str],
+    head_name: str,
+    C: float = 1.0,
+    max_iter: int = 1000,
+) -> LinearHead:
+    """
+    Train a single linear classification head.
+
+    Uses sklearn LogisticRegression with:
+        - L2 regularisation (C parameter)
+        - class_weight='balanced' for imbalanced data
+        - lbfgs solver (good for moderate feature counts)
+        - multinomial objective even for binary (consistent API)
+
+    Extracts W and b into a LinearHead for numpy-only inference.
+    """
+    logger.info(f"\n{'='*60}")
+    logger.info(f"Training {head_name} head")
+    logger.info(f"{'='*60}")
+    logger.info(f"  Train: {X_train.shape[0]:,} | Test: {X_test.shape[0]:,}")
+    logger.info(f"  Features: {X_train.shape[1]} | Classes: {len(labels)}")
+
+    # Class distribution
+    unique, counts = np.unique(y_train, return_counts=True)
+    for u, c in zip(unique, counts):
+        logger.info(f"    {u}: {c:,}")
+
+    start = time.time()
+
+    clf = LogisticRegression(
+        C=C,
+        max_iter=max_iter,
+        class_weight="balanced",
+        solver="lbfgs",
+        n_jobs=-1,
+        random_state=42,
+    )
+    clf.fit(X_train, y_train)
+
+    elapsed = time.time() - start
+    logger.info(f"  Trained in {elapsed:.1f}s")
+
+    # Evaluate
+    y_pred = clf.predict(X_test)
+    report = classification_report(y_test, y_pred, target_names=labels, digits=4)
+    logger.info(f"\n{report}")
+
+    # Extract weights into LinearHead
+    head = LinearHead(labels)
+    # sklearn stores coef_ as (n_classes, n_features) for multinomial
+    # We want W as (n_features, n_classes) for X @ W + b
+    if clf.coef_.shape[0] == 1:
+        # Binary case: sklearn only stores one row
+        # Expand to full 2-class format
+        head.W = np.vstack([-clf.coef_[0], clf.coef_[0]]).T.astype("float32")
+        head.b = np.array([-clf.intercept_[0], clf.intercept_[0]], dtype="float32")
+    else:
+        head.W = clf.coef_.T.astype("float32")  # (features, classes)
+        head.b = clf.intercept_.astype("float32")
+
+    # Sanity check: reproduce sklearn predictions
+    logits = X_test[:5] @ head.W + head.b
+    e = np.exp(logits - logits.max(axis=-1, keepdims=True))
+    probs = e / e.sum(axis=-1, keepdims=True)
+    np_pred_idx = np.argmax(probs, axis=1)
+    sk_pred_idx = clf.predict(X_test[:5])  # returns integer class indices
+    assert list(np_pred_idx) == list(int(x) for x in sk_pred_idx), \
+        f"Mismatch: {list(np_pred_idx)} vs {list(sk_pred_idx)}"
+    logger.info("  ✓ Numpy inference matches sklearn predictions")
+
+    return head
+
+
+def train_all(
+    data_dir: Path,
+    config: Optional[PubGuardConfig] = None,
+    test_size: float = 0.15,
+):
+    """
+    Train all three classification heads.
+
+    Args:
+        data_dir: Directory containing the prepared NDJSON files
+        config: PubGuard configuration
+        test_size: Fraction of data held out for evaluation
+    """
+    config = config or PubGuardConfig()
+    data_dir = Path(data_dir)
+    cache_dir = data_dir / "embeddings_cache"
+
+    logger.info("=" * 60)
+    logger.info("PubGuard Training Pipeline")
+    logger.info("=" * 60)
+    logger.info(f"Data dir: {data_dir}")
+    logger.info(f"Models dir: {config.models_dir}")
+    start_total = time.time()
+
+    # ── HEAD 1: doc_type ────────────────────────────────────────
+    doc_type_path = data_dir / "doc_type_train.ndjson"
+    if doc_type_path.exists():
+        texts, labels = load_ndjson(doc_type_path)
+        label_to_idx = {l: i for i, l in enumerate(DOC_TYPE_LABELS)}
+
+        # Embed
+        embeddings = embed_texts(
+            texts, config,
+            cache_path=cache_dir / "doc_type_emb.npy",
+        )
+
+        # Add structural features
+        logger.info("Computing structural features...")
+        struct = compute_structural_features(texts)
+        X = np.concatenate([embeddings, struct], axis=1)
+
+        y = np.array([label_to_idx.get(l, 0) for l in labels])
+
+        X_tr, X_te, y_tr, y_te = train_test_split(
+            X, y, test_size=test_size, stratify=y, random_state=42
+        )
+
+        head = train_head(X_tr, y_tr, X_te, y_te, DOC_TYPE_LABELS, "doc_type")
+        head.save(config.doc_type_head_path)
+        logger.info(f"Saved → {config.doc_type_head_path}")
+    else:
+        logger.warning(f"doc_type data not found: {doc_type_path}")
+
+    # ── HEAD 2: ai_detect ───────────────────────────────────────
+    ai_path = data_dir / "ai_detect_train.ndjson"
+    if ai_path.exists():
+        texts, labels = load_ndjson(ai_path)
+        label_to_idx = {l: i for i, l in enumerate(AI_DETECT_LABELS)}
+
+        embeddings = embed_texts(
+            texts, config,
+            cache_path=cache_dir / "ai_detect_emb.npy",
+        )
+
+        y = np.array([label_to_idx.get(l, 0) for l in labels])
+
+        X_tr, X_te, y_tr, y_te = train_test_split(
+            embeddings, y, test_size=test_size, stratify=y, random_state=42
+        )
+
+        head = train_head(X_tr, y_tr, X_te, y_te, AI_DETECT_LABELS, "ai_detect")
+        head.save(config.ai_detect_head_path)
+        logger.info(f"Saved → {config.ai_detect_head_path}")
+    else:
+        logger.warning(f"ai_detect data not found: {ai_path}")
+
+    # ── HEAD 3: toxicity ────────────────────────────────────────
+    tox_path = data_dir / "toxicity_train.ndjson"
+    if tox_path.exists():
+        texts, labels = load_ndjson(tox_path)
+        label_to_idx = {l: i for i, l in enumerate(TOXICITY_LABELS)}
+
+        embeddings = embed_texts(
+            texts, config,
+            cache_path=cache_dir / "toxicity_emb.npy",
+        )
+
+        y = np.array([label_to_idx.get(l, 0) for l in labels])
+
+        X_tr, X_te, y_tr, y_te = train_test_split(
+            embeddings, y, test_size=test_size, stratify=y, random_state=42
+        )
+
+        head = train_head(X_tr, y_tr, X_te, y_te, TOXICITY_LABELS, "toxicity")
+        head.save(config.toxicity_head_path)
+        logger.info(f"Saved → {config.toxicity_head_path}")
+    else:
+        logger.warning(f"toxicity data not found: {tox_path}")
+
+    elapsed = time.time() - start_total
+    logger.info(f"\nTotal training time: {elapsed/60:.1f} minutes")
+    logger.info("All heads saved to: " + str(config.models_dir))