FabledCurator/backend/app/services/ml/heads.py

"""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()