bvandeusen
eaea4308fc
The head-vs-centroid eval (#1130) existed to prove the 'frozen embedding + trained head' spine; the operator accepted the tagging system and dropped the harness. Removed per rule 22: TagEvalCard + store, /api/tag_eval blueprint, tag_eval_run ml task, recover-stalled-tag-eval-runs sweep + beat entry, TagEvalRun model + table (migration 0073), and its tests. The eval's data loaders + metric helpers were NOT eval-specific — the nightly heads trainer runs on them — so they moved verbatim to services/ml/training_data.py (heads.py import updated; behavior unchanged). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CDgx8bQS5YrGRK76v8HUnM
503 lines
19 KiB
Python
503 lines
19 KiB
Python
"""Production heads: train + score the per-concept classifiers (#114).
|
|
|
|
The eval harness (#1130) proved the spine, then retired 2026-07-02 once the
|
|
tagging system was accepted; this is the 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. Uses the eval-proven data loaders
|
|
+ metric helpers (training_data.py) so heads match the 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 (
|
|
HeadAutoApplyRun,
|
|
HeadTrainingRun,
|
|
ImageRecord,
|
|
ImageRegion,
|
|
MLSettings,
|
|
Tag,
|
|
TagHead,
|
|
TagKind,
|
|
TagSuggestionRejection,
|
|
)
|
|
from ...models.tag import image_tag
|
|
from .training_data 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, threshold_override: float | None = None,
|
|
) -> list[dict]:
|
|
"""Suggestions for one image from the trained heads: [{tag_id, name,
|
|
category, score}], ranked. A concept surfaces when its score clears the
|
|
head's own suggest_threshold — or, when threshold_override is given (the
|
|
typed-dropdown "show everything" mode), that flat floor instead (0 → every
|
|
head). Empty if the image has no embedding or no heads exist yet.
|
|
|
|
MAX-OVER-BAG: the image is scored as a BAG of embeddings — the whole-image
|
|
vector PLUS every concept-region crop the agent embedded (same model
|
|
version) — and each head takes its MAX score across the bag. A small/local
|
|
concept (glasses, a stomach bulge) that the whole-image vector washes out
|
|
can still surface from the crop where it dominates. The whole-image vector is
|
|
always in the bag, so this can never score lower than whole-image alone."""
|
|
import numpy as np
|
|
|
|
img = await session.get(ImageRecord, image_id)
|
|
if img is None:
|
|
return []
|
|
settings = await _settings_async(session)
|
|
cur_version = settings.embedder_model_version
|
|
heads = await _current_heads(session, cur_version)
|
|
if heads["W"] is None:
|
|
return []
|
|
|
|
# Only embeddings in the CURRENT model's space enter the bag. Mid model-swap
|
|
# (#1190), an image still carrying the OLD-version whole-image vector is
|
|
# skipped rather than scored by heads trained in a different space; a legacy
|
|
# NULL version is treated as current (those predate per-row stamping).
|
|
bag = []
|
|
if img.siglip_embedding is not None and img.siglip_model_version in (
|
|
cur_version, None,
|
|
):
|
|
bag.append(np.asarray(img.siglip_embedding, dtype=np.float32))
|
|
region_vecs = (
|
|
await session.execute(
|
|
select(ImageRegion.siglip_embedding)
|
|
.where(ImageRegion.image_record_id == image_id)
|
|
.where(ImageRegion.siglip_embedding.is_not(None))
|
|
.where(ImageRegion.embedding_version == cur_version)
|
|
)
|
|
).all()
|
|
for (vec,) in region_vecs:
|
|
if vec is not None:
|
|
bag.append(np.asarray(vec, dtype=np.float32))
|
|
if not bag:
|
|
return []
|
|
|
|
X = np.vstack(bag) # (B, D)
|
|
norms = np.linalg.norm(X, axis=1, keepdims=True)
|
|
norms[norms == 0] = 1.0
|
|
Xn = X / norms
|
|
Z = Xn @ heads["W"].T + heads["b"] # (B, H)
|
|
probs = (1.0 / (1.0 + np.exp(-Z))).max(axis=0) # (H,) best over the bag
|
|
out = []
|
|
for i, p in enumerate(probs):
|
|
cut = threshold_override if threshold_override is not None else heads["thr"][i]
|
|
if p >= cut:
|
|
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()
|
|
|
|
|
|
# --- Earned auto-apply (sync, ml worker) ---------------------------------
|
|
# A graduated head can apply its tag to images it scores above the head's
|
|
# auto_apply_threshold, without a human. Gated by a master switch + a support
|
|
# floor so a precise-looking but under-supported head can't spray tags.
|
|
|
|
_AUTO_APPLY_CHUNK = 5000
|
|
|
|
|
|
class HeadAutoApplyAlreadyRunning(Exception):
|
|
"""Raised when an auto-apply sweep is already in flight."""
|
|
|
|
|
|
class HeadAutoApplyDisabled(Exception):
|
|
"""Raised when a real (non-dry-run) sweep is requested but the master
|
|
switch (head_auto_apply_enabled) is off."""
|
|
|
|
|
|
def start_head_auto_apply_run(session: Session, params: dict[str, Any]) -> int:
|
|
"""Create a HeadAutoApplyRun + dispatch the ml-queue sweep. dry_run previews
|
|
(writes nothing); a real sweep needs the master switch on. One run at a time."""
|
|
dry_run = bool((params or {}).get("dry_run", False))
|
|
existing = session.execute(
|
|
select(HeadAutoApplyRun.id).where(HeadAutoApplyRun.status == "running")
|
|
).scalar_one_or_none()
|
|
if existing is not None:
|
|
raise HeadAutoApplyAlreadyRunning(existing)
|
|
if not dry_run and not _settings(session).head_auto_apply_enabled:
|
|
raise HeadAutoApplyDisabled()
|
|
run = HeadAutoApplyRun(
|
|
dry_run=dry_run, params={"dry_run": dry_run}, status="running",
|
|
last_progress_at=datetime.now(UTC),
|
|
)
|
|
session.add(run)
|
|
session.flush()
|
|
run_id = run.id
|
|
from ...tasks.ml import apply_head_tags as _task
|
|
_task.delay(run_id)
|
|
return run_id
|
|
|
|
|
|
def _auto_apply_heads(session: Session, embedding_version: str, min_pos: int):
|
|
"""Eligible heads to fire: graduated (auto_apply_threshold set), enough
|
|
support, current embedding. Returns the row list (tag_id/name/weights/...)."""
|
|
return session.execute(
|
|
select(
|
|
TagHead.tag_id, Tag.name, TagHead.weights, TagHead.bias,
|
|
TagHead.auto_apply_threshold,
|
|
)
|
|
.join(Tag, Tag.id == TagHead.tag_id)
|
|
.where(TagHead.embedding_version == embedding_version)
|
|
.where(TagHead.auto_apply_threshold.is_not(None))
|
|
.where(TagHead.n_pos >= min_pos)
|
|
).all()
|
|
|
|
|
|
def auto_apply_sweep(
|
|
session: Session, run: HeadAutoApplyRun, dry_run: bool
|
|
) -> dict[str, Any]:
|
|
"""Score every embedded image against the eligible heads and apply (or, for
|
|
dry_run, just count) each head's tag where score >= its auto_apply_threshold
|
|
and the tag isn't already applied or rejected on that image. Streams
|
|
embeddings in chunks; commits per chunk on a real run. Returns
|
|
{n_applied, concepts:[{tag_id,name,applied,scanned,threshold}]}."""
|
|
import numpy as np
|
|
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
|
|
|
settings = _settings(session)
|
|
rows = _auto_apply_heads(
|
|
session, settings.embedder_model_version,
|
|
settings.head_auto_apply_min_positives,
|
|
)
|
|
if not rows:
|
|
return {"n_applied": 0, "concepts": []}
|
|
|
|
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.auto_apply_threshold for r in rows], dtype=np.float32)
|
|
tag_ids = [r.tag_id for r in rows]
|
|
names = [r.name for r in rows]
|
|
|
|
# Skip images that already carry, or have rejected, each tag.
|
|
skip = {tid: set() for tid in tag_ids}
|
|
for tid in tag_ids:
|
|
for (iid,) in session.execute(
|
|
select(image_tag.c.image_record_id).where(image_tag.c.tag_id == tid)
|
|
):
|
|
skip[tid].add(iid)
|
|
for (iid,) in session.execute(
|
|
select(TagSuggestionRejection.image_record_id).where(
|
|
TagSuggestionRejection.tag_id == tid
|
|
)
|
|
):
|
|
skip[tid].add(iid)
|
|
|
|
applied = [0] * len(rows)
|
|
scanned = 0
|
|
all_ids = list(session.execute(
|
|
select(ImageRecord.id).where(ImageRecord.siglip_embedding.is_not(None))
|
|
).scalars())
|
|
for start in range(0, len(all_ids), _AUTO_APPLY_CHUNK):
|
|
chunk = all_ids[start:start + _AUTO_APPLY_CHUNK]
|
|
emb = _load_embeddings(session, chunk)
|
|
cids = [i for i in chunk if i in emb]
|
|
if not cids:
|
|
continue
|
|
Xn = _l2norm(np.vstack([emb[i] for i in cids]).astype(np.float32), np)
|
|
probs = 1.0 / (1.0 + np.exp(-(Xn @ W.T + b))) # (N, H)
|
|
scanned += len(cids)
|
|
for h in range(len(rows)):
|
|
tid = tag_ids[h]
|
|
for idx in np.where(probs[:, h] >= thr[h])[0]:
|
|
iid = cids[int(idx)]
|
|
if iid in skip[tid]:
|
|
continue
|
|
skip[tid].add(iid)
|
|
applied[h] += 1
|
|
if not dry_run:
|
|
session.execute(
|
|
pg_insert(image_tag)
|
|
.values(image_record_id=iid, tag_id=tid, source="head_auto")
|
|
.on_conflict_do_nothing()
|
|
)
|
|
if not dry_run:
|
|
session.commit()
|
|
run.last_progress_at = datetime.now(UTC)
|
|
session.commit()
|
|
|
|
concepts = [
|
|
{"tag_id": tag_ids[h], "name": names[h], "applied": applied[h],
|
|
"scanned": scanned, "threshold": float(thr[h])}
|
|
for h in range(len(rows))
|
|
]
|
|
return {"n_applied": sum(applied), "concepts": concepts}
|