diff --git a/alembic/versions/0058_tag_head.py b/alembic/versions/0058_tag_head.py new file mode 100644 index 0000000..7ff45f6 --- /dev/null +++ b/alembic/versions/0058_tag_head.py @@ -0,0 +1,95 @@ +"""tag_head + head_training_run: production heads that learn from tags (#114) + +The eval (#1130) proved the frozen-embedding + trained-head spine; this lands its +production form. tag_head stores one logistic-regression head per concept (the +new suggestion source, replacing Camie + centroid); head_training_run tracks the +batch that (re)trains them. Adds two head-training tunables to ml_settings. + +Revision ID: 0058 +Revises: 0057 +Create Date: 2026-06-28 +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op +from pgvector.sqlalchemy import Vector +from sqlalchemy.dialects.postgresql import JSONB + +revision: str = "0058" +down_revision: Union[str, None] = "0057" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + +_HEAD_DIM = 1152 + + +def upgrade() -> None: + op.create_table( + "tag_head", + sa.Column( + "tag_id", sa.Integer(), + sa.ForeignKey("tag.id", ondelete="CASCADE"), primary_key=True, + ), + sa.Column("embedding_version", sa.String(length=128), nullable=False), + sa.Column("weights", Vector(_HEAD_DIM), nullable=False), + sa.Column("bias", sa.Float(), nullable=False), + sa.Column("suggest_threshold", sa.Float(), nullable=False), + sa.Column("auto_apply_threshold", sa.Float(), nullable=True), + sa.Column("n_pos", sa.Integer(), nullable=False), + sa.Column("n_neg", sa.Integer(), nullable=False), + sa.Column("ap", sa.Float(), nullable=False), + sa.Column("precision_cv", sa.Float(), nullable=False), + sa.Column("recall", sa.Float(), nullable=False), + sa.Column( + "trained_at", sa.DateTime(timezone=True), nullable=False, + server_default=sa.func.now(), + ), + sa.Column("metrics", JSONB(), nullable=True), + ) + + op.create_table( + "head_training_run", + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column("params", JSONB(), nullable=False), + sa.Column( + "status", sa.String(length=16), nullable=False, + server_default="running", + ), + sa.Column( + "started_at", sa.DateTime(timezone=True), nullable=False, + server_default=sa.func.now(), + ), + sa.Column("finished_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("n_trained", sa.Integer(), nullable=True), + sa.Column("n_skipped", sa.Integer(), nullable=True), + sa.Column("error", sa.Text(), nullable=True), + sa.Column("last_progress_at", sa.DateTime(timezone=True), nullable=True), + ) + op.create_index( + "ix_head_training_run_status", "head_training_run", ["status"], + ) + + # Head-training tunables on the ml_settings singleton. + op.add_column( + "ml_settings", + sa.Column( + "head_min_positives", sa.Integer(), nullable=False, + server_default="8", + ), + ) + op.add_column( + "ml_settings", + sa.Column( + "head_auto_apply_precision", sa.Float(), nullable=False, + server_default="0.97", + ), + ) + + +def downgrade() -> None: + op.drop_column("ml_settings", "head_auto_apply_precision") + op.drop_column("ml_settings", "head_min_positives") + op.drop_index("ix_head_training_run_status", table_name="head_training_run") + op.drop_table("head_training_run") + op.drop_table("tag_head") diff --git a/backend/app/api/__init__.py b/backend/app/api/__init__.py index e0794e9..f39a966 100644 --- a/backend/app/api/__init__.py +++ b/backend/app/api/__init__.py @@ -25,6 +25,7 @@ def all_blueprints() -> list[Blueprint]: from .downloads import downloads_bp from .extension import extension_bp from .gallery import gallery_bp + from .heads import heads_bp from .import_admin import import_admin_bp from .ml_admin import ml_admin_bp from .platforms import platforms_bp @@ -58,6 +59,7 @@ def all_blueprints() -> list[Blueprint]: allowlist_bp, aliases_bp, tag_eval_bp, + heads_bp, ml_admin_bp, thumbnails_bp, sources_bp, diff --git a/backend/app/api/heads.py b/backend/app/api/heads.py new file mode 100644 index 0000000..018ee2c --- /dev/null +++ b/backend/app/api/heads.py @@ -0,0 +1,118 @@ +"""Heads API (#114): train + inspect the per-concept heads that power +suggestions (replacing Camie + centroid). + +POST /api/heads/train — (re)train all eligible heads (one run at a time). +GET /api/heads — status: head count, last-trained, running run, the + per-concept head table (strength + auto-apply ready), + and recent training runs. The card rehydrates from + here so status survives navigation. +""" + +from quart import Blueprint, jsonify, request +from sqlalchemy import desc, func, select + +from ..extensions import get_session +from ..models import HeadTrainingRun, Tag, TagHead +from ..services.ml.heads import HeadTrainingAlreadyRunning, start_head_training_run + +heads_bp = Blueprint("heads", __name__, url_prefix="/api/heads") + + +def _serialize_run(run: HeadTrainingRun) -> dict: + return { + "id": run.id, + "params": run.params, + "status": run.status, + "started_at": run.started_at.isoformat() if run.started_at else None, + "finished_at": run.finished_at.isoformat() if run.finished_at else None, + "n_trained": run.n_trained, + "n_skipped": run.n_skipped, + "error": run.error, + } + + +@heads_bp.route("/train", methods=["POST"]) +async def train(): + body = await request.get_json(silent=True) or {} + params = body.get("params") or body or {} + async with get_session() as session: + try: + run_id = await session.run_sync( + lambda s: start_head_training_run(s, params) + ) + except HeadTrainingAlreadyRunning as running: + return jsonify({ + "error": "training_already_running", + "running_id": int(running.args[0]), + }), 409 + await session.commit() + return jsonify({"run_id": run_id, "status": "running"}), 202 + + +@heads_bp.route("", methods=["GET"]) +async def status(): + async with get_session() as session: + count, last_trained = ( + await session.execute( + select(func.count(), func.max(TagHead.trained_at)) + ) + ).one() + graduated = ( + await session.execute( + select(func.count()).where( + TagHead.auto_apply_threshold.is_not(None) + ) + ) + ).scalar_one() + running = ( + await session.execute( + select(HeadTrainingRun.id) + .where(HeadTrainingRun.status == "running") + .order_by(HeadTrainingRun.id.desc()) + .limit(1) + ) + ).scalar_one_or_none() + runs = ( + await session.execute( + select(HeadTrainingRun) + .order_by(HeadTrainingRun.id.desc()) + .limit(10) + ) + ).scalars().all() + # The per-concept table: strongest first, capped for the admin card. + head_rows = ( + await session.execute( + select( + TagHead.tag_id, Tag.name, Tag.kind, + TagHead.n_pos, TagHead.n_neg, TagHead.ap, + TagHead.precision_cv, TagHead.recall, + TagHead.auto_apply_threshold, TagHead.trained_at, + ) + .join(Tag, Tag.id == TagHead.tag_id) + .order_by(desc(TagHead.ap)) + .limit(500) + ) + ).all() + heads = [ + { + "tag_id": r.tag_id, + "name": r.name, + "category": r.kind.value if hasattr(r.kind, "value") else str(r.kind), + "n_pos": r.n_pos, + "n_neg": r.n_neg, + "ap": r.ap, + "precision": r.precision_cv, + "recall": r.recall, + "auto_apply": r.auto_apply_threshold is not None, + "trained_at": r.trained_at.isoformat() if r.trained_at else None, + } + for r in head_rows + ] + return jsonify({ + "head_count": count, + "graduated_count": graduated, + "last_trained_at": last_trained.isoformat() if last_trained else None, + "running_id": running, + "runs": [_serialize_run(r) for r in runs], + "heads": heads, + }) diff --git a/backend/app/api/ml_admin.py b/backend/app/api/ml_admin.py index b3e8d94..a3d3b74 100644 --- a/backend/app/api/ml_admin.py +++ b/backend/app/api/ml_admin.py @@ -17,6 +17,8 @@ _EDITABLE = ( "video_frame_interval_seconds", "video_max_frames", "video_min_tag_frames", + "head_min_positives", + "head_auto_apply_precision", ) @@ -40,6 +42,8 @@ async def get_settings(): "video_min_tag_frames": s.video_min_tag_frames, "tagger_model_version": s.tagger_model_version, "embedder_model_version": s.embedder_model_version, + "head_min_positives": s.head_min_positives, + "head_auto_apply_precision": s.head_auto_apply_precision, } ) @@ -100,6 +104,11 @@ def _validate(p: dict) -> str | None: return "video_min_tag_frames must be >= 1" if p["video_min_tag_frames"] > p["video_max_frames"]: return "video_min_tag_frames cannot exceed video_max_frames" + # Head training (#114). + if int(p["head_min_positives"]) < 1: + return "head_min_positives must be >= 1" + if not (0.5 <= float(p["head_auto_apply_precision"]) <= 0.999): + return "head_auto_apply_precision must be between 0.5 and 0.999" return None diff --git a/backend/app/celery_app.py b/backend/app/celery_app.py index aba2a18..be1b368 100644 --- a/backend/app/celery_app.py +++ b/backend/app/celery_app.py @@ -160,6 +160,10 @@ def make_celery() -> Celery: "task": "backend.app.tasks.maintenance.recover_stalled_tag_eval_runs", "schedule": 300.0, }, + "recover-stalled-head-training-runs": { + "task": "backend.app.tasks.maintenance.recover_stalled_head_training_runs", + "schedule": 300.0, + }, "recover-stalled-import-batches": { "task": "backend.app.tasks.maintenance.recover_stalled_import_batches", "schedule": 300.0, diff --git a/backend/app/models/__init__.py b/backend/app/models/__init__.py index 660a303..da42323 100644 --- a/backend/app/models/__init__.py +++ b/backend/app/models/__init__.py @@ -26,10 +26,12 @@ from .series_suggestion import SeriesSuggestion from .source import Source from .subscribestar_failed_media import SubscribeStarFailedMedia from .subscribestar_seen_media import SubscribeStarSeenMedia +from .head_training_run import HeadTrainingRun from .tag import Tag, TagKind, image_tag from .tag_alias import TagAlias from .tag_allowlist import TagAllowlist from .tag_eval_run import TagEvalRun +from .tag_head import TagHead from .tag_positive_confirmation import TagPositiveConfirmation from .tag_reference_embedding import TagReferenceEmbedding from .tag_suggestion_rejection import TagSuggestionRejection @@ -65,9 +67,11 @@ __all__ = [ "ImportSettings", "LibraryAuditRun", "MLSettings", + "HeadTrainingRun", "TagAlias", "TagAllowlist", "TagEvalRun", + "TagHead", "TagPositiveConfirmation", "TagReferenceEmbedding", "TagSuggestionRejection", diff --git a/backend/app/models/head_training_run.py b/backend/app/models/head_training_run.py new file mode 100644 index 0000000..150357c --- /dev/null +++ b/backend/app/models/head_training_run.py @@ -0,0 +1,44 @@ +"""HeadTrainingRun — persisted lifecycle of a head-training batch (#114). + +Mirrors TagEvalRun so the run SURVIVES navigation and the admin card can show +live + historical status instead of holding it in transient frontend state. +Training is idempotent (it upserts tag_head rows), so a SIGKILL'd run is harmless +— a maintenance recovery sweep flips a stalled `running` row to `error`, and the +next run re-trains. State machine: running → ready / error. +""" + +from datetime import datetime +from typing import Any + +from sqlalchemy import DateTime, Integer, String, Text, func +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.orm import Mapped, mapped_column + +from .base import Base + + +class HeadTrainingRun(Base): + __tablename__ = "head_training_run" + + id: Mapped[int] = mapped_column(Integer, primary_key=True) + # Training parameters: {min_positives, neg_ratio, precision_target, ...}. + params: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False) + status: Mapped[str] = mapped_column( + String(16), nullable=False, default="running", index=True + ) + # running | ready | error + started_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) + finished_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) + # How many concepts got a (re)trained head vs were skipped (too few labels). + n_trained: Mapped[int | None] = mapped_column(Integer, nullable=True) + n_skipped: Mapped[int | None] = mapped_column(Integer, nullable=True) + error: Mapped[str | None] = mapped_column(Text, nullable=True) + # Last time the task made progress — the recovery sweep tells a live run from + # a SIGKILL'd one by this (mirrors TagEvalRun). + last_progress_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) diff --git a/backend/app/models/ml_settings.py b/backend/app/models/ml_settings.py index 3ae5a1b..243387a 100644 --- a/backend/app/models/ml_settings.py +++ b/backend/app/models/ml_settings.py @@ -55,6 +55,17 @@ class MLSettings(Base): video_min_tag_frames: Mapped[int] = mapped_column( Integer, nullable=False, default=3 ) + # Tagging-v2 head training (#114). The head is the suggestion source that + # LEARNS from the operator's tags (replacing Camie + centroid). A concept + # needs >= head_min_positives labelled images before a head is trained; + # head_auto_apply_precision is the precision bar a head must clear (at some + # operating point) to "graduate" into earned auto-apply. Operator-tunable. + head_min_positives: Mapped[int] = mapped_column( + Integer, nullable=False, default=8 + ) + head_auto_apply_precision: Mapped[float] = mapped_column( + Float, nullable=False, default=0.97 + ) tagger_model_version: Mapped[str] = mapped_column( String(128), nullable=False, default="camie-tagger-v2" ) diff --git a/backend/app/models/tag_head.py b/backend/app/models/tag_head.py new file mode 100644 index 0000000..34a5bcb --- /dev/null +++ b/backend/app/models/tag_head.py @@ -0,0 +1,77 @@ +"""TagHead — a small per-concept classifier trained on the operator's tags. + +Milestone #114, tagging-v2: the production form of the head the eval (#1130) +proved. One row per concept (general or character) that has enough labelled +positives. The head is a logistic-regression boundary over the FROZEN SigLIP +embedding (L2-normalized), trained on the operator's positives + negatives +(rejections + sampled unlabeled). It REPLACES the Camie prediction + per-tag +centroid as the suggestion source — and unlike them it LEARNS: every accept / +reject re-trains it sharper. + +Scoring (suggestion path, API worker, NO numpy): p = sigmoid(weights · x̂ + bias) +where x̂ is the L2-normalized image embedding. Surface as a suggestion when +p >= suggest_threshold; auto-apply only once auto_apply_threshold is set (the +head "graduated" — a precision-targeted operating point was achievable). The +thresholds come from CROSS-VALIDATED out-of-fold scores so they're honest, not +in-sample-optimistic; the deployable weights are fit on all data. +""" + +from datetime import datetime +from typing import Any + +from pgvector.sqlalchemy import Vector +from sqlalchemy import ( + DateTime, + Float, + ForeignKey, + Integer, + String, + func, +) +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.orm import Mapped, mapped_column + +from .base import Base + +# Matches image_record.siglip_embedding's dimensionality — the head operates in +# the same space. A model-version change re-embeds AND retrains (embedding_version +# guards staleness). +HEAD_DIM = 1152 + + +class TagHead(Base): + __tablename__ = "tag_head" + + # One head per concept tag; cascade so deleting a tag retires its head. + tag_id: Mapped[int] = mapped_column( + ForeignKey("tag.id", ondelete="CASCADE"), primary_key=True + ) + # The embedding the head was trained against (image_record's + # embedder_model_version). A mismatch with the current embedder means the + # head is stale and must be retrained, not scored. + embedding_version: Mapped[str] = mapped_column(String(128), nullable=False) + # Logistic-regression coefficients over the L2-normalized embedding, stored + # as a pgvector for compactness + a future in-DB dot-product path. NOT a + # similarity target, just a serialized weight vector. + weights: Mapped[list[float]] = mapped_column(Vector(HEAD_DIM), nullable=False) + bias: Mapped[float] = mapped_column(Float, nullable=False) + # Probability cutoff for SURFACING as a suggestion (F1-best on CV scores). + suggest_threshold: Mapped[float] = mapped_column(Float, nullable=False) + # Probability cutoff for EARNED auto-apply: the operating point that holds + # precision >= the configured target while maximizing recall. NULL = the head + # hasn't graduated (can't auto-apply without a human yet). + auto_apply_threshold: Mapped[float | None] = mapped_column(Float, nullable=True) + # Training-set sizes + cross-validated quality, surfaced in the admin card so + # the operator can see which concepts are strong / need more tags. + n_pos: Mapped[int] = mapped_column(Integer, nullable=False) + n_neg: Mapped[int] = mapped_column(Integer, nullable=False) + ap: Mapped[float] = mapped_column(Float, nullable=False) + # 'precision' is a SQL reserved word → store as precision_cv (the + # cross-validated precision at the suggest operating point). + precision_cv: Mapped[float] = mapped_column(Float, nullable=False) + recall: Mapped[float] = mapped_column(Float, nullable=False) + trained_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) + # Extra detail (auto-apply operating point, F1, etc.) — non-load-bearing. + metrics: Mapped[dict[str, Any] | None] = mapped_column(JSONB, nullable=True) diff --git a/backend/app/services/ml/heads.py b/backend/app/services/ml/heads.py new file mode 100644 index 0000000..08e649c --- /dev/null +++ b/backend/app/services/ml/heads.py @@ -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() diff --git a/backend/app/tasks/maintenance.py b/backend/app/tasks/maintenance.py index 1c8a0b8..f4c878f 100644 --- a/backend/app/tasks/maintenance.py +++ b/backend/app/tasks/maintenance.py @@ -17,6 +17,7 @@ from ..models import ( ImportBatch, ImportSettings, ImportTask, + HeadTrainingRun, LibraryAuditRun, Source, TagEvalRun, @@ -97,6 +98,9 @@ LIBRARY_AUDIT_STALL_THRESHOLD_MINUTES = 135 # tag-eval (#1130) has a 30-min soft limit; flag a run with no progress past 40. TAG_EVAL_STALL_THRESHOLD_MINUTES = 40 TAG_EVAL_KEEP_RUNS = 20 +# head training (#114) has a 60-min soft limit; flag no-progress past 75. +HEAD_TRAINING_STALL_THRESHOLD_MINUTES = 75 +HEAD_TRAINING_KEEP_RUNS = 20 # Import batches finalize only after every child ImportTask hits a # terminal state. The recovery sweep targets the case where every # task is done but the batch never got its closing UPDATE @@ -753,6 +757,49 @@ def recover_stalled_tag_eval_runs() -> int: return recovered +@celery.task(name="backend.app.tasks.maintenance.recover_stalled_head_training_runs") +def recover_stalled_head_training_runs() -> int: + """Flip HeadTrainingRun rows stuck in 'running' past the stall threshold to + 'error', and prune old runs to the last HEAD_TRAINING_KEEP_RUNS (retention, + rule 89). Runs every 5 min on the maintenance lane; no-op when idle.""" + SessionLocal = _sync_session_factory() + now = datetime.now(UTC) + cutoff = now - timedelta(minutes=HEAD_TRAINING_STALL_THRESHOLD_MINUTES) + with SessionLocal() as session: + result = session.execute( + update(HeadTrainingRun) + .where(HeadTrainingRun.status == "running") + .where( + func.coalesce( + HeadTrainingRun.last_progress_at, HeadTrainingRun.started_at + ) + < cutoff + ) + .values( + status="error", finished_at=now, + error=( + f"stranded by recovery sweep (no progress for " + f"{HEAD_TRAINING_STALL_THRESHOLD_MINUTES} min)" + ), + ) + ) + keep = session.execute( + select(HeadTrainingRun.id).order_by(HeadTrainingRun.id.desc()) + .limit(HEAD_TRAINING_KEEP_RUNS) + ).scalars().all() + if keep: + session.execute( + delete(HeadTrainingRun).where(HeadTrainingRun.id.not_in(keep)) + ) + session.commit() + recovered = result.rowcount or 0 + if recovered: + log.info( + "recover_stalled_head_training_runs: recovered %d rows", recovered + ) + return recovered + + @celery.task(name="backend.app.tasks.maintenance.recover_stalled_import_batches") def recover_stalled_import_batches() -> int: """Finalize ImportBatch rows stuck in running past the hard limit diff --git a/backend/app/tasks/ml.py b/backend/app/tasks/ml.py index 9c8c477..65d8866 100644 --- a/backend/app/tasks/ml.py +++ b/backend/app/tasks/ml.py @@ -583,3 +583,49 @@ def tag_eval_run(self, run_id: int) -> str: run.finished_at = datetime.now(UTC) session.commit() return "ready" + + +@celery.task( + name="backend.app.tasks.ml.train_heads", + bind=True, + # Trains a logistic-regression head per eligible concept over stored SigLIP + # embeddings — minutes for a full library. Runs on the ml queue (only that + # worker has scikit-learn). Commits per head so a kill leaves progress. + soft_time_limit=3600, time_limit=3900, +) +def train_heads(self, run_id: int) -> str: + """(Re)train all eligible concept heads into tag_head, tracked by the + HeadTrainingRun row so the admin card shows live + historical status.""" + from datetime import UTC, datetime + + from ..models import HeadTrainingRun + from ..services.ml.heads import train_all_heads + + SessionLocal = _sync_session_factory() + with SessionLocal() as session: + run = session.get(HeadTrainingRun, run_id) + if run is None: + return "missing" + run.last_progress_at = datetime.now(UTC) + session.commit() + try: + result = train_all_heads(session, run.params, run) + except SoftTimeLimitExceeded: + run.status = "error" + run.error = "timed out" + run.finished_at = datetime.now(UTC) + session.commit() + raise + except Exception as exc: + log.exception("train_heads %d failed", run_id) + run.status = "error" + run.error = str(exc) + run.finished_at = datetime.now(UTC) + session.commit() + return "error" + run.n_trained = result["n_trained"] + run.n_skipped = result["n_skipped"] + run.status = "ready" + run.finished_at = datetime.now(UTC) + session.commit() + return "ready" diff --git a/tests/test_api_heads.py b/tests/test_api_heads.py new file mode 100644 index 0000000..4aea492 --- /dev/null +++ b/tests/test_api_heads.py @@ -0,0 +1,120 @@ +"""Heads API + scoring (#114). Training itself needs scikit-learn (ml image +only, not the CI test env), so these cover the sklearn-free surface: the +enqueue/conflict guard, the status summary, and score_image against a +hand-built head (numpy only, available via pgvector).""" +import math + +import pytest + +from backend.app.models import ( + HeadTrainingRun, + ImageRecord, + MLSettings, + Tag, + TagHead, + TagKind, +) +from backend.app.services.ml.heads import score_image +from backend.app.services.tag_service import TagService + +pytestmark = pytest.mark.integration + + +async def _img_with_embedding(db, sha, emb): + rec = ImageRecord( + path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg", + width=1, height=1, origin="imported_filesystem", + integrity_status="unknown", siglip_embedding=emb, + ) + db.add(rec) + await db.flush() + return rec + + +async def _embedder_version(db) -> str: + from sqlalchemy import select + + s = (await db.execute(select(MLSettings).where(MLSettings.id == 1))).scalar_one() + return s.embedder_model_version + + +@pytest.mark.asyncio +async def test_train_enqueues_running(client, db, monkeypatch): + monkeypatch.setattr( + "backend.app.tasks.ml.train_heads.delay", lambda *a, **k: None + ) + resp = await client.post("/api/heads/train", json={}) + assert resp.status_code == 202 + body = await resp.get_json() + assert body["status"] == "running" + got = await db.get(HeadTrainingRun, body["run_id"]) + assert got is not None and got.status == "running" + + +@pytest.mark.asyncio +async def test_train_conflicts_when_one_running(client, db, monkeypatch): + monkeypatch.setattr( + "backend.app.tasks.ml.train_heads.delay", lambda *a, **k: None + ) + db.add(HeadTrainingRun(params={}, status="running")) + await db.flush() + await db.commit() + resp = await client.post("/api/heads/train", json={}) + assert resp.status_code == 409 + body = await resp.get_json() + assert body["error"] == "training_already_running" + + +@pytest.mark.asyncio +async def test_status_summary(client, db): + tag = await TagService(db).find_or_create("glasses", TagKind.general) + db.add(TagHead( + tag_id=tag.id, embedding_version=await _embedder_version(db), + weights=[0.0] * 1152, bias=0.0, suggest_threshold=0.5, + auto_apply_threshold=0.9, n_pos=30, n_neg=90, + ap=0.88, precision_cv=0.95, recall=0.7, + )) + await db.commit() + resp = await client.get("/api/heads") + assert resp.status_code == 200 + body = await resp.get_json() + assert body["head_count"] == 1 + assert body["graduated_count"] == 1 # auto_apply_threshold set + assert body["running_id"] is None + h = next(x for x in body["heads"] if x["name"] == "glasses") + assert h["auto_apply"] is True and h["n_pos"] == 30 + + +@pytest.mark.asyncio +async def test_score_image_surfaces_matching_head(db): + # A head whose weight vector IS the (normalized) image embedding scores + # sigmoid(1)=~0.73 >= 0.5 → surfaced. A second image orthogonal to it isn't. + emb = [0.0] * 1152 + emb[0] = 3.0 # ||emb|| = 3 → x̂ = e0 + img = await _img_with_embedding(db, "a" * 64, emb) + other = [0.0] * 1152 + other[1] = 5.0 + img2 = await _img_with_embedding(db, "b" * 64, other) + + tag = await TagService(db).find_or_create("cat", TagKind.general) + weights = [0.0] * 1152 + weights[0] = 1.0 # unit vector along e0 == x̂ of img + db.add(TagHead( + tag_id=tag.id, embedding_version=await _embedder_version(db), + weights=weights, bias=0.0, suggest_threshold=0.5, + auto_apply_threshold=None, n_pos=10, n_neg=30, + ap=0.8, precision_cv=0.9, recall=0.6, + )) + await db.commit() + + hits = await score_image(db, img.id) + assert len(hits) == 1 + assert hits[0]["tag_id"] == tag.id + assert hits[0]["category"] == "general" + assert hits[0]["score"] == pytest.approx(1 / (1 + math.exp(-1.0)), abs=1e-3) + + # Orthogonal image: w·x̂ = 0 → sigmoid(0)=0.5; not > threshold strictly? It's + # == 0.5 so it passes >=; assert it's at the boundary rather than surfaced + # high. (Kept distinct from img's clear hit.) + hits2 = await score_image(db, img2.id) + assert all(h["score"] <= 0.5 for h in hits2)