feat(heads): production per-concept heads — train + score backend (#114 A)

The eval (#1130) proved the frozen-embedding + trained-head spine; this lands its production form (the first of three slices that make heads the suggestion source, replacing Camie + centroid). - tag_head: one logistic-regression head per general/character concept with enough labelled positives. Weights (pgvector), honest CV-derived suggest threshold + earned-auto-apply point, and per-concept quality metrics. - head_training_run: persisted batch lifecycle (mirrors tag_eval_run) so the admin card shows live + historical status across navigation. - services/ml/heads.py: TRAIN (sync, ml worker, reuses tag_eval's proven data loaders + metric math so production heads match measured eval numbers) and SCORE (async, API worker — numpy via pgvector, no scikit-learn): score one image's embedding against all heads → the rail's suggestions, cached on (count, max trained_at) so a retrain invalidates without per-request loads. - tasks.ml.train_heads (ml queue, commits per head so a kill leaves progress) + recover_stalled_head_training_runs sweep + retention(20) + 5-min beat (rule 89). - api/heads.py: POST /api/heads/train (one run at a time, 409 guard) + GET /api/heads (count, graduated, last-trained, running, per-concept table, recent runs). - ml_settings: head_min_positives + head_auto_apply_precision, tunable via /api/ml/settings. Scoring isn't wired into the rail yet (slice C) and the admin UI is slice B — this slice makes training + scoring exist and CI-verifiable. 'precision' column stored as precision_cv (SQL reserved word). Migration 0058. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
2026-06-28 10:36:25 -04:00
parent 179c1a9dcc
commit 22c3b54746
13 changed files with 904 additions and 0 deletions
@@ -0,0 +1,327 @@
+"""Production heads: train + score the per-concept classifiers (#114).
+
+The eval (#1130, tag_eval.py) proved the spine; this is its production form.
+- TRAIN (sync, ml worker — needs scikit-learn): for every general/character tag
+  with enough labelled positives, fit a logistic-regression head on the FROZEN
+  SigLIP embeddings (positives + negatives = rejections + sampled unlabeled),
+  derive an honest suggest threshold + earned-auto-apply point from CROSS-
+  VALIDATED scores, and upsert a TagHead row. Reuses tag_eval's proven data
+  loaders + metric helpers so production heads match the eval's measured numbers.
+- SCORE (async, API worker — numpy via pgvector, NO scikit-learn): score one
+  image's embedding against all current heads → the suggestions the rail shows,
+  REPLACING Camie predictions + per-tag centroids.
+
+scikit-learn is imported lazily inside the train path so the API worker can still
+import this module to enqueue training + to score (scoring needs only numpy).
+"""
+
+from __future__ import annotations
+
+import logging
+from datetime import UTC, datetime
+from typing import Any
+
+from sqlalchemy import delete, func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import Session
+
+from ...models import (
+    HeadTrainingRun,
+    ImageRecord,
+    MLSettings,
+    Tag,
+    TagHead,
+    TagKind,
+)
+from ...models.tag import image_tag
+# Reuse the eval's proven, identical data loaders + metric math so a production
+# head's quality matches what the eval reported for the same concept.
+from .tag_eval import (
+    _auto_apply_point,
+    _ids_with_tag,
+    _l2norm,
+    _load_embeddings,
+    _metrics_from_scores,
+    _rejected_ids,
+    _safe_folds,
+    _sample_unlabeled,
+)
+
+log = logging.getLogger(__name__)
+
+DEFAULT_NEG_RATIO = 3
+DEFAULT_CV_FOLDS = 5
+MIN_POSITIVES_FLOOR = 8     # hard floor; settings.head_min_positives can raise it
+_UNLABELED_POOL = 4000
+_EXAMPLES_MIN = 8           # need at least this many embedded +/- to fit a head
+
+# Only these tag kinds get heads (the surfaced suggestion categories).
+_HEAD_KINDS = (TagKind.general, TagKind.character)
+# tag.kind -> the suggestion category the rail groups under.
+_CATEGORY = {TagKind.general: "general", TagKind.character: "character"}
+
+
+class HeadTrainingAlreadyRunning(Exception):
+    """Raised by start_head_training_run when a run is already in flight."""
+
+
+def start_head_training_run(session: Session, params: dict[str, Any]) -> int:
+    """Create a HeadTrainingRun (status='running') + dispatch the ml-queue task.
+    Returns the run id. One training run at a time (light guard)."""
+    existing = session.execute(
+        select(HeadTrainingRun.id).where(HeadTrainingRun.status == "running")
+    ).scalar_one_or_none()
+    if existing is not None:
+        raise HeadTrainingAlreadyRunning(existing)
+    norm = _normalize_params(session, params)
+    run = HeadTrainingRun(
+        params=norm, status="running", last_progress_at=datetime.now(UTC)
+    )
+    session.add(run)
+    session.flush()
+    run_id = run.id
+    from ...tasks.ml import train_heads as _task
+    _task.delay(run_id)
+    return run_id
+
+
+def _settings(session: Session) -> MLSettings:
+    return session.execute(
+        select(MLSettings).where(MLSettings.id == 1)
+    ).scalar_one()
+
+
+def _normalize_params(session: Session, params: dict[str, Any] | None) -> dict[str, Any]:
+    params = params or {}
+    s = _settings(session)
+    try:
+        min_pos = max(MIN_POSITIVES_FLOOR, int(params.get("min_positives", s.head_min_positives)))
+    except (TypeError, ValueError):
+        min_pos = max(MIN_POSITIVES_FLOOR, s.head_min_positives)
+    try:
+        neg_ratio = max(1, int(params.get("neg_ratio", DEFAULT_NEG_RATIO)))
+    except (TypeError, ValueError):
+        neg_ratio = DEFAULT_NEG_RATIO
+    try:
+        cv_folds = max(2, int(params.get("cv_folds", DEFAULT_CV_FOLDS)))
+    except (TypeError, ValueError):
+        cv_folds = DEFAULT_CV_FOLDS
+    try:
+        precision_target = min(max(float(params.get("precision_target", s.head_auto_apply_precision)), 0.5), 0.999)
+    except (TypeError, ValueError):
+        precision_target = s.head_auto_apply_precision
+    return {
+        "min_positives": min_pos,
+        "neg_ratio": neg_ratio,
+        "cv_folds": cv_folds,
+        "precision_target": round(precision_target, 4),
+    }
+
+
+def _embedder_version(session: Session) -> str:
+    return _settings(session).embedder_model_version
+
+
+def _eligible_tag_ids(session: Session, min_pos: int) -> list[int]:
+    """Concept tags (general/character) with >= min_pos labelled images — the
+    set that gets a head. Counts all sources; source-aware filtering (#1133) is
+    a separate, optional refinement."""
+    rows = session.execute(
+        select(Tag.id)
+        .join(image_tag, image_tag.c.tag_id == Tag.id)
+        .where(Tag.kind.in_(_HEAD_KINDS))
+        .group_by(Tag.id)
+        .having(func.count(image_tag.c.image_record_id) >= min_pos)
+    ).all()
+    return [r[0] for r in rows]
+
+
+def train_all_heads(
+    session: Session, params: dict[str, Any], run: HeadTrainingRun | None = None
+) -> dict[str, int]:
+    """(Re)train a head for every eligible concept; prune heads whose tag is no
+    longer eligible. Commits per head so a SIGKILL leaves trained heads durable
+    (training is idempotent). Returns {n_trained, n_skipped}."""
+    import numpy as np
+
+    cfg = _normalize_params(session, params)
+    embedding_version = _embedder_version(session)
+    eligible = _eligible_tag_ids(session, cfg["min_positives"])
+    eligible_set = set(eligible)
+    trained = 0
+    skipped = 0
+    for i, tag_id in enumerate(eligible):
+        try:
+            ok = train_head(session, tag_id, embedding_version, cfg, np)
+        except Exception:
+            log.exception("train_head failed for tag %d", tag_id)
+            ok = False
+        session.commit()
+        trained += int(ok)
+        skipped += int(not ok)
+        if run is not None and i % 10 == 0:
+            run.last_progress_at = datetime.now(UTC)
+            session.commit()
+    # Retire heads whose concept dropped out of the eligible set (lost its
+    # positives, or the tag was re-kinded) so stale heads can't keep suggesting.
+    if eligible_set:
+        session.execute(delete(TagHead).where(TagHead.tag_id.not_in(eligible_set)))
+    else:
+        session.execute(delete(TagHead))
+    session.commit()
+    return {"n_trained": trained, "n_skipped": skipped}
+
+
+def train_head(
+    session: Session, tag_id: int, embedding_version: str, cfg: dict, np
+) -> bool:
+    """Fit + upsert one head. Returns True if a head was written, False if the
+    concept had too few usable examples to train (the row is then removed)."""
+    from sklearn.linear_model import LogisticRegression
+    from sklearn.model_selection import StratifiedKFold, cross_val_predict
+
+    pos_ids = _ids_with_tag(session, tag_id)
+    if len(pos_ids) < cfg["min_positives"]:
+        session.execute(delete(TagHead).where(TagHead.tag_id == tag_id))
+        return False
+
+    pos_set = set(pos_ids)
+    rejected = [i for i in _rejected_ids(session, tag_id) if i not in pos_set]
+    want_neg = max(len(pos_ids) * cfg["neg_ratio"], _EXAMPLES_MIN * 4)
+    sampled = _sample_unlabeled(
+        session, pos_set | set(rejected), min(_UNLABELED_POOL, want_neg)
+    )
+    neg_ids = rejected + [i for i in sampled if i not in pos_set]
+
+    emb = _load_embeddings(session, pos_ids + neg_ids)
+    pos = [emb[i] for i in pos_ids if i in emb]
+    neg = [emb[i] for i in neg_ids if i in emb]
+    if len(pos) < _EXAMPLES_MIN or len(neg) < _EXAMPLES_MIN:
+        session.execute(delete(TagHead).where(TagHead.tag_id == tag_id))
+        return False
+
+    X = np.vstack(pos + neg).astype(np.float32)
+    y = np.array([1] * len(pos) + [0] * len(neg))
+    Xn = _l2norm(X, np)
+
+    clf = LogisticRegression(max_iter=1000, class_weight="balanced")
+    cv = StratifiedKFold(
+        n_splits=_safe_folds(y, cfg["cv_folds"], np), shuffle=True, random_state=0
+    )
+    # Honest thresholds from out-of-fold scores; deployable weights from a final
+    # fit on ALL the data.
+    cv_probs = cross_val_predict(clf, Xn, y, cv=cv, method="predict_proba")[:, 1]
+    metrics = _metrics_from_scores(y, cv_probs, np)
+    auto = _auto_apply_point(y, cv_probs, cfg["precision_target"], np)
+    clf.fit(Xn, y)
+
+    head = session.get(TagHead, tag_id)
+    if head is None:
+        head = TagHead(tag_id=tag_id)
+        session.add(head)
+    head.embedding_version = embedding_version
+    head.weights = clf.coef_[0].astype(np.float32).tolist()
+    head.bias = float(clf.intercept_[0])
+    head.suggest_threshold = float(metrics["threshold"])
+    head.auto_apply_threshold = float(auto["threshold"]) if auto else None
+    head.n_pos = len(pos)
+    head.n_neg = len(neg)
+    head.ap = float(metrics["ap"])
+    head.precision_cv = float(metrics["precision"])
+    head.recall = float(metrics["recall"])
+    head.trained_at = datetime.now(UTC)
+    head.metrics = {"f1": metrics["f1"], "auto_apply": auto}
+    return True
+
+
+# --- Scoring (async, API worker) -----------------------------------------
+# Score one image against every current head to produce the rail's suggestions.
+# A tiny in-process cache holds the stacked weight matrix keyed on (count,
+# max(trained_at)) so a retrain invalidates it without per-request weight loads.
+_HEADS_CACHE: dict[str, Any] = {"key": None, "heads": None}
+
+
+async def _current_heads(session: AsyncSession, embedding_version: str):
+    """Stacked (W, b, thresholds, tag_id/name/category) for heads matching the
+    current embedding, cached until the next retrain."""
+    import numpy as np
+
+    sig = (
+        await session.execute(
+            select(func.count(), func.max(TagHead.trained_at)).where(
+                TagHead.embedding_version == embedding_version
+            )
+        )
+    ).one()
+    key = f"{embedding_version}:{sig[0]}:{sig[1].isoformat() if sig[1] else '-'}"
+    cached = _HEADS_CACHE.get("heads")
+    if cached is not None and _HEADS_CACHE.get("key") == key:
+        return cached
+    rows = (
+        await session.execute(
+            select(
+                TagHead.tag_id, Tag.name, Tag.kind,
+                TagHead.weights, TagHead.bias,
+                TagHead.suggest_threshold, TagHead.auto_apply_threshold,
+            )
+            .join(Tag, Tag.id == TagHead.tag_id)
+            .where(TagHead.embedding_version == embedding_version)
+        )
+    ).all()
+    if not rows:
+        loaded = {"W": None, "rows": []}
+    else:
+        W = np.vstack([np.asarray(r.weights, dtype=np.float32) for r in rows])
+        b = np.asarray([r.bias for r in rows], dtype=np.float32)
+        thr = np.asarray([r.suggest_threshold for r in rows], dtype=np.float32)
+        meta = [
+            {
+                "tag_id": r.tag_id,
+                "name": r.name,
+                "category": _CATEGORY.get(r.kind, "general"),
+                "auto_apply_threshold": r.auto_apply_threshold,
+            }
+            for r in rows
+        ]
+        loaded = {"W": W, "b": b, "thr": thr, "meta": meta}
+    _HEADS_CACHE["key"] = key
+    _HEADS_CACHE["heads"] = loaded
+    return loaded
+
+
+async def score_image(session: AsyncSession, image_id: int) -> list[dict]:
+    """Suggestions for one image from the trained heads: [{tag_id, name,
+    category, score}], score >= each head's suggest_threshold, ranked. Empty if
+    the image has no embedding or no heads exist yet."""
+    import numpy as np
+
+    img = await session.get(ImageRecord, image_id)
+    if img is None or img.siglip_embedding is None:
+        return []
+    settings = await _settings_async(session)
+    heads = await _current_heads(session, settings.embedder_model_version)
+    if heads["W"] is None:
+        return []
+    x = np.asarray(img.siglip_embedding, dtype=np.float32)
+    n = float(np.linalg.norm(x)) or 1.0
+    xn = x / n
+    z = heads["W"] @ xn + heads["b"]
+    probs = 1.0 / (1.0 + np.exp(-z))
+    out = []
+    for i, p in enumerate(probs):
+        if p >= heads["thr"][i]:
+            m = heads["meta"][i]
+            out.append({
+                "tag_id": m["tag_id"],
+                "name": m["name"],
+                "category": m["category"],
+                "score": float(p),
+            })
+    out.sort(key=lambda d: d["score"], reverse=True)
+    return out
+
+
+async def _settings_async(session: AsyncSession) -> MLSettings:
+    return (
+        await session.execute(select(MLSettings).where(MLSettings.id == 1))
+    ).scalar_one()