feat(ml): tag-eval backend — head-vs-centroid learning-curve eval (persisted)

Slice 1 of milestone #114 (tagging v2). Proves the frozen-embedding + trained- head spine on the operator's own data, reusing the SigLIP embeddings already stored on image_record — no re-embedding, no GPU. Per concept: train a logistic-regression HEAD (positives + negatives = explicit rejections + sampled unlabeled) vs the old single-CENTROID baseline; report cross-validated precision/recall/AP for both, a LEARNING CURVE (AP/F1 as tagged positives grow 10→30→100→300), and example image ids (head-would-suggest / head-doubts-positive) to eyeball. Persisted so the report SURVIVES navigation (operator-flagged): the run + full report live in a new tag_eval_run row (mirrors library_audit_run); the admin card will rehydrate from GET on mount, not transient state. - models.TagEvalRun + migration 0056; runs on the ml queue (only worker with numpy/sklearn) — numpy/sklearn lazy-imported so the API can still enqueue. - services/ml/tag_eval (compute + start helper, one-running guard), tasks.ml .tag_eval_run, api/tag-eval (POST create, GET history light / detail w/ report). - recover_stalled_tag_eval_runs sweep + retention (keep last 20) + 5-min beat (rule 89). scikit-learn added to requirements-ml. - tests: param normalization + the rehydrate read-path + create/conflict. Frontend admin card (trigger + render persisted report) follows next. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-27 22:49:10 -04:00
parent 958378312c
commit 6e3c5f697f
11 changed files with 655 additions and 0 deletions
@@ -0,0 +1,316 @@
+"""Head-vs-centroid tagging eval (#1130, milestone #114 slice 1).
+
+Proves the "frozen embedding + small trained head (with negatives)" spine on the
+operator's OWN data, reusing the SigLIP embeddings already stored on
+image_record. For each concept tag it compares:
+  - CENTROID baseline (the old approach): cosine to the mean of positive vectors.
+  - HEAD (the new approach): logistic regression trained on positives + negatives.
+and reports cross-validated precision/recall/AP for both, a LEARNING CURVE
+(accuracy as the number of tagged positives grows), and example image ids to
+eyeball.
+
+numpy + scikit-learn are imported LAZILY inside run_eval so the API worker (base
+image, no ML stack) can still import start_tag_eval_run to enqueue the ml-queue
+task — the heavy compute only runs on the ml worker.
+"""
+
+from __future__ import annotations
+
+import logging
+from datetime import UTC, datetime
+from typing import Any
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import Session
+
+from ...models import ImageRecord, Tag, TagEvalRun, TagSuggestionRejection
+from ...models.tag import image_tag
+
+log = logging.getLogger(__name__)
+
+# The operator's real concept list (mix of whole-ish + small/local cues). The
+# admin trigger can override; this is the default eval set.
+DEFAULT_CONCEPTS = [
+    "glasses", "cat", "dog", "horse", "goblin",
+    "cum", "lactation", "fellatio", "xray", "stomach bulge",
+]
+DEFAULT_CURVE_POINTS = [10, 30, 100, 300]
+DEFAULT_NEG_RATIO = 3       # negatives per positive (rejections + sampled unlabeled)
+DEFAULT_CV_FOLDS = 5
+MIN_POSITIVES = 8           # below this, a concept can't be evaluated meaningfully
+_UNLABELED_POOL = 4000      # cap on sampled unlabeled rows pulled per concept
+_EXAMPLES_K = 12
+
+
+def start_tag_eval_run(session: Session, params: dict[str, Any]) -> int:
+    """Create a TagEvalRun (status='running') and dispatch the ml-queue task.
+    Returns the new run id. Light guard: one running eval at a time."""
+    existing = session.execute(
+        select(TagEvalRun.id).where(TagEvalRun.status == "running")
+    ).scalar_one_or_none()
+    if existing is not None:
+        raise EvalAlreadyRunning(existing)
+    norm = _normalize_params(params)
+    run = TagEvalRun(params=norm, status="running", last_progress_at=datetime.now(UTC))
+    session.add(run)
+    session.flush()
+    run_id = run.id
+    # Same enqueue-by-import pattern api/suggestions.py uses for ml tasks; the
+    # commit happens in the API handler so row + dispatch are visible together.
+    from ...tasks.ml import tag_eval_run as _task
+    _task.delay(run_id)
+    return run_id
+
+
+class EvalAlreadyRunning(Exception):
+    """Raised by start_tag_eval_run when an eval is already in flight."""
+
+
+def _normalize_params(params: dict[str, Any] | None) -> dict[str, Any]:
+    params = params or {}
+    concepts = params.get("concepts") or DEFAULT_CONCEPTS
+    concepts = [str(c).strip() for c in concepts if str(c).strip()]
+    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
+    curve = params.get("curve_points") or DEFAULT_CURVE_POINTS
+    curve = sorted({int(n) for n in curve if int(n) > 0})
+    return {
+        "concepts": concepts,
+        "neg_ratio": neg_ratio,
+        "cv_folds": cv_folds,
+        "curve_points": curve,
+    }
+
+
+def _resolve_tag_id(session: Session, name: str) -> int | None:
+    """Case-insensitive tag-name match; if several share a name, take the one
+    applied to the most images (the one the operator actually uses)."""
+    rows = session.execute(
+        select(Tag.id, func.count(image_tag.c.image_record_id))
+        .outerjoin(image_tag, image_tag.c.tag_id == Tag.id)
+        .where(func.lower(Tag.name) == name.lower())
+        .group_by(Tag.id)
+        .order_by(func.count(image_tag.c.image_record_id).desc())
+    ).all()
+    return rows[0][0] if rows else None
+
+
+def _ids_with_tag(session: Session, tag_id: int) -> list[int]:
+    return [
+        r[0] for r in session.execute(
+            select(image_tag.c.image_record_id).where(image_tag.c.tag_id == tag_id)
+        ).all()
+    ]
+
+
+def _rejected_ids(session: Session, tag_id: int) -> list[int]:
+    return [
+        r[0] for r in session.execute(
+            select(TagSuggestionRejection.image_record_id)
+            .where(TagSuggestionRejection.tag_id == tag_id)
+        ).all()
+    ]
+
+
+def _sample_unlabeled(session: Session, exclude: set[int], limit: int) -> list[int]:
+    """Random image ids (with an embedding) NOT carrying the tag. Concepts are
+    sparse, so an untagged image is almost always a true negative."""
+    stmt = (
+        select(ImageRecord.id)
+        .where(ImageRecord.siglip_embedding.is_not(None))
+        .order_by(func.random())
+        .limit(limit)
+    )
+    if exclude:
+        stmt = stmt.where(ImageRecord.id.not_in(exclude))
+    return [r[0] for r in session.execute(stmt).all()]
+
+
+def _load_embeddings(session: Session, ids: list[int]) -> dict[int, Any]:
+    import numpy as np
+
+    out: dict[int, Any] = {}
+    if not ids:
+        return out
+    # Chunk the IN list to stay well under psycopg's parameter ceiling.
+    for i in range(0, len(ids), 2000):
+        chunk = ids[i:i + 2000]
+        for rid, emb in session.execute(
+            select(ImageRecord.id, ImageRecord.siglip_embedding)
+            .where(ImageRecord.id.in_(chunk))
+            .where(ImageRecord.siglip_embedding.is_not(None))
+        ).all():
+            out[rid] = np.asarray(emb, dtype=np.float32)
+    return out
+
+
+def run_eval(session: Session, params: dict[str, Any]) -> dict[str, Any]:
+    """Compute the full report. Per-concept failures are captured, not fatal."""
+    import numpy as np
+
+    cfg = _normalize_params(params)
+    concepts_out = []
+    for name in cfg["concepts"]:
+        try:
+            concepts_out.append(_eval_concept(session, name, cfg, np))
+        except Exception as exc:  # one bad concept shouldn't kill the run
+            log.exception("tag-eval concept %r failed", name)
+            concepts_out.append({"name": name, "skipped": f"error: {exc}"})
+    return {
+        "generated_at": datetime.now(UTC).isoformat(),
+        "params": cfg,
+        "concepts": concepts_out,
+    }
+
+
+def _eval_concept(session: Session, name: str, cfg: dict, np) -> dict[str, Any]:
+    tag_id = _resolve_tag_id(session, name)
+    if tag_id is None:
+        return {"name": name, "skipped": "no such tag"}
+    pos_ids = _ids_with_tag(session, tag_id)
+    if len(pos_ids) < MIN_POSITIVES:
+        return {"name": name, "tag_id": tag_id, "n_pos": len(pos_ids),
+                "skipped": f"too few positives (<{MIN_POSITIVES})"}
+
+    neg_ratio = cfg["neg_ratio"]
+    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) * neg_ratio, _EXAMPLES_K * 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 = [(i, emb[i]) for i in pos_ids if i in emb]
+    neg = [(i, emb[i]) for i in neg_ids if i in emb]
+    if len(pos) < MIN_POSITIVES or len(neg) < MIN_POSITIVES:
+        return {"name": name, "tag_id": tag_id, "n_pos": len(pos),
+                "n_neg": len(neg), "skipped": "too few embedded examples"}
+
+    ids = np.array([i for i, _ in pos] + [i for i, _ in neg])
+    X = np.vstack([v for _, v in pos] + [v for _, v in neg]).astype(np.float32)
+    y = np.array([1] * len(pos) + [0] * len(neg))
+    Xn = _l2norm(X, np)
+
+    head = _eval_head(Xn, y, cfg["cv_folds"], np)
+    centroid = _eval_centroid(Xn, y, cfg["cv_folds"], np)
+    curve = _learning_curve(Xn, y, cfg["curve_points"], neg_ratio, np)
+    examples = _examples(Xn, y, ids, np)
+
+    return {
+        "name": name, "tag_id": tag_id,
+        "n_pos": len(pos), "n_neg": len(neg),
+        "n_rejected": len(rejected),
+        "head": head, "centroid": centroid,
+        "curve": curve, "examples": examples,
+    }
+
+
+def _l2norm(X, np):
+    n = np.linalg.norm(X, axis=1, keepdims=True)
+    n[n == 0] = 1.0
+    return X / n
+
+
+def _metrics_from_scores(y, scores, np) -> dict[str, float]:
+    from sklearn.metrics import average_precision_score, precision_recall_curve
+
+    ap = float(average_precision_score(y, scores))
+    prec, rec, thr = precision_recall_curve(y, scores)
+    f1 = (2 * prec * rec) / np.clip(prec + rec, 1e-9, None)
+    best = int(np.argmax(f1))
+    # thr has len = len(prec)-1; map best index safely.
+    t = float(thr[min(best, len(thr) - 1)]) if len(thr) else 0.5
+    return {
+        "ap": round(ap, 4),
+        "precision": round(float(prec[best]), 4),
+        "recall": round(float(rec[best]), 4),
+        "f1": round(float(f1[best]), 4),
+        "threshold": round(t, 4),
+    }
+
+
+def _safe_folds(y, folds, np) -> int:
+    minority = int(min(np.bincount(y)))
+    return max(2, min(folds, minority))
+
+
+def _eval_head(Xn, y, folds, np) -> dict[str, float]:
+    from sklearn.linear_model import LogisticRegression
+    from sklearn.model_selection import StratifiedKFold, cross_val_predict
+
+    clf = LogisticRegression(max_iter=1000, class_weight="balanced")
+    cv = StratifiedKFold(n_splits=_safe_folds(y, folds, np), shuffle=True,
+                         random_state=0)
+    probs = cross_val_predict(clf, Xn, y, cv=cv, method="predict_proba")[:, 1]
+    return _metrics_from_scores(y, probs, np)
+
+
+def _eval_centroid(Xn, y, folds, np) -> dict[str, float]:
+    """Cross-validated cosine-to-positive-mean — the OLD method's quality."""
+    from sklearn.model_selection import StratifiedKFold
+
+    cv = StratifiedKFold(n_splits=_safe_folds(y, folds, np), shuffle=True,
+                         random_state=0)
+    scores = np.zeros(len(y), dtype=np.float32)
+    for train, test in cv.split(Xn, y):
+        c = Xn[train][y[train] == 1].mean(axis=0)
+        cn = c / (np.linalg.norm(c) or 1.0)
+        scores[test] = Xn[test] @ cn
+    return _metrics_from_scores(y, scores, np)
+
+
+def _learning_curve(Xn, y, points, neg_ratio, np) -> list[dict[str, float]]:
+    """Hold out a fixed test split; train the head on a growing number of
+    positives and watch AP/F1 climb — answers 'does tagging more sharpen it?'"""
+    from sklearn.linear_model import LogisticRegression
+    from sklearn.model_selection import train_test_split
+
+    rng = np.random.default_rng(0)
+    idx = np.arange(len(y))
+    try:
+        tr, te = train_test_split(idx, test_size=0.3, stratify=y, random_state=0)
+    except ValueError:
+        return []
+    tr_pos = tr[y[tr] == 1]
+    tr_neg = tr[y[tr] == 0]
+    out = []
+    for n in points:
+        if n > len(tr_pos):
+            break
+        sp = rng.choice(tr_pos, size=n, replace=False)
+        nn = min(len(tr_neg), n * neg_ratio)
+        sn = rng.choice(tr_neg, size=nn, replace=False)
+        sub = np.concatenate([sp, sn])
+        clf = LogisticRegression(max_iter=1000, class_weight="balanced")
+        clf.fit(Xn[sub], y[sub])
+        prob = clf.predict_proba(Xn[te])[:, 1]
+        m = _metrics_from_scores(y[te], prob, np)
+        out.append({"n_pos": int(n), "ap": m["ap"], "f1": m["f1"]})
+    return out
+
+
+def _examples(Xn, y, ids, np) -> dict[str, list[int]]:
+    """Train on all data, then surface: top-scoring UNLABELED-ish (highest among
+    the negative pool = what the head would newly suggest) and lowest-scoring
+    POSITIVES (where the head disagrees with the operator's tag — likely the
+    most informative to review)."""
+    from sklearn.linear_model import LogisticRegression
+
+    clf = LogisticRegression(max_iter=1000, class_weight="balanced")
+    clf.fit(Xn, y)
+    s = clf.predict_proba(Xn)[:, 1]
+    neg_idx = np.where(y == 0)[0]
+    pos_idx = np.where(y == 1)[0]
+    top_neg = neg_idx[np.argsort(s[neg_idx])[::-1][:_EXAMPLES_K]]
+    low_pos = pos_idx[np.argsort(s[pos_idx])[:_EXAMPLES_K]]
+    return {
+        "head_would_suggest": [int(ids[i]) for i in top_neg],
+        "head_doubts_positive": [int(ids[i]) for i in low_pos],
+    }