feat(heads): production per-concept heads — train + score backend (#114 A)
CI / lint (push) Failing after 3s
CI / frontend-build (push) Successful in 19s
CI / backend-lint-and-test (push) Successful in 26s
CI / integration (push) Failing after 3m26s

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
This commit is contained in:
2026-06-28 10:36:25 -04:00
parent 179c1a9dcc
commit 22c3b54746
13 changed files with 904 additions and 0 deletions
+95
View File
@@ -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")
+2
View File
@@ -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,
+118
View File
@@ -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,
})
+9
View File
@@ -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
+4
View File
@@ -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,
+4
View File
@@ -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",
+44
View File
@@ -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
)
+11
View File
@@ -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"
)
+77
View File
@@ -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)
+327
View File
@@ -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()
+47
View File
@@ -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
+46
View File
@@ -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"
+120
View File
@@ -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)