bvandeusen
2d44a26bdf
Makes auto-apply truly "soft" for heads: _ids_with_tag (head positives) and _eligible_tag_ids (graduation count) now count human-applied + operator-confirmed tags only, via a shared _AUTO_SOURCES (head_auto/ccip_auto/ml_auto) exclusion. Unconfirmed auto-applied tags no longer train the head that judges them, so a misfire can't reinforce itself and the retraction sweep can actually drop it. Confirming a tag (TagPositiveConfirmation) promotes it to a positive AND protects it from retraction. sklearn-free tests. CCIP reference exclusion is the companion piece, next. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CDgx8bQS5YrGRK76v8HUnM
796 lines
32 KiB
Python
796 lines
32 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, exists, func, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy.orm import Session
|
|
|
|
from ...models import (
|
|
HeadAutoApplyRun,
|
|
HeadTrainingRun,
|
|
ImageRecord,
|
|
ImageRegion,
|
|
MLSettings,
|
|
Tag,
|
|
TagHead,
|
|
TagKind,
|
|
TagPositiveConfirmation,
|
|
TagSuggestionRejection,
|
|
)
|
|
from ...models.tag import image_tag
|
|
from .training_data import (
|
|
_AUTO_SOURCES,
|
|
_auto_apply_point,
|
|
_hygiene_excluded_ids,
|
|
_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"}
|
|
|
|
# System-tag (wip/banner/editor screenshot) heads surface as suggestions at
|
|
# this FLAT confidence floor instead of their auto-derived (precision-tuned)
|
|
# suggest threshold. The auto threshold is high, so it hides the borderline /
|
|
# false-positive guesses — which are exactly the cases the operator wants to
|
|
# SEE and REJECT to sharpen these heads (hard-negative mining: "negatively
|
|
# reinforce what isn't a system tag"). Operator-set 0.65 (2026-07-03): high
|
|
# enough not to spam near-zero scores, low enough to surface real mistakes.
|
|
# Content-tag heads keep their own thresholds; the typed-dropdown's
|
|
# threshold_override still overrides everything (show-all mode).
|
|
_SYSTEM_TAG_SUGGEST_FLOOR = 0.65
|
|
|
|
|
|
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 POSITIVE images — the set
|
|
that gets a head. Counts human-applied + operator-confirmed tags only;
|
|
unconfirmed auto-applied predictions do NOT count toward eligibility (they
|
|
don't train the head — milestone 139), so a concept can't graduate on its own
|
|
guesses."""
|
|
confirmed = exists().where(
|
|
TagPositiveConfirmation.image_record_id == image_tag.c.image_record_id,
|
|
TagPositiveConfirmation.tag_id == image_tag.c.tag_id,
|
|
)
|
|
rows = session.execute(
|
|
select(Tag.id)
|
|
.join(image_tag, image_tag.c.tag_id == Tag.id)
|
|
.where(Tag.kind.in_(_HEAD_KINDS))
|
|
.where(image_tag.c.source.not_in(_AUTO_SOURCES) | confirmed)
|
|
.group_by(Tag.id)
|
|
.having(func.count(image_tag.c.image_record_id) >= min_pos)
|
|
).all()
|
|
return [r[0] for r in rows]
|
|
|
|
|
|
def _head_fingerprints(session: Session, tag_ids: list[int]) -> dict[int, str]:
|
|
"""Per-tag training-data fingerprint: (positive count, latest positive
|
|
created_at) + (rejection count, latest rejected_at). It moves whenever a tag
|
|
gains/loses a positive or a rejection — the incremental-retrain change
|
|
detector (#1317 p2). A newly-added positive/rejection always has the latest
|
|
timestamp, so even a remove-one-add-one (unchanged count) is caught. The
|
|
sampled-unlabeled negative pool + the hygiene set drift GLOBALLY and are
|
|
reconciled by the nightly full run, not captured here."""
|
|
if not tag_ids:
|
|
return {}
|
|
pos = session.execute(
|
|
select(
|
|
image_tag.c.tag_id,
|
|
func.count(image_tag.c.image_record_id),
|
|
func.max(image_tag.c.created_at),
|
|
)
|
|
.where(image_tag.c.tag_id.in_(tag_ids))
|
|
.group_by(image_tag.c.tag_id)
|
|
).all()
|
|
pos_map = {t: (c, m) for t, c, m in pos}
|
|
rej = session.execute(
|
|
select(
|
|
TagSuggestionRejection.tag_id,
|
|
func.count(),
|
|
func.max(TagSuggestionRejection.rejected_at),
|
|
)
|
|
.where(TagSuggestionRejection.tag_id.in_(tag_ids))
|
|
.group_by(TagSuggestionRejection.tag_id)
|
|
).all()
|
|
rej_map = {t: (c, m) for t, c, m in rej}
|
|
out = {}
|
|
for t in tag_ids:
|
|
pc, pm = pos_map.get(t, (0, None))
|
|
rc, rm = rej_map.get(t, (0, None))
|
|
out[t] = f"{pc}:{pm}:{rc}:{rm}"
|
|
return out
|
|
|
|
|
|
def _heads_needing_retrain(
|
|
session: Session, eligible: list[int], embedding_version: str,
|
|
fps: dict[int, str], full: bool,
|
|
) -> list[int]:
|
|
"""The eligible tag_ids to (re)fit: no head yet, a head trained in a DIFFERENT
|
|
embedding space (a model swap), or a changed training-data fingerprint.
|
|
full=True forces every eligible tag. sklearn-free (train_head itself needs
|
|
scikit-learn) so the incremental decision is unit-testable on its own."""
|
|
if full:
|
|
return list(eligible)
|
|
existing = {
|
|
tag_id: (fp, ev)
|
|
for tag_id, fp, ev in session.execute(
|
|
select(
|
|
TagHead.tag_id, TagHead.train_fingerprint,
|
|
TagHead.embedding_version,
|
|
)
|
|
).all()
|
|
}
|
|
out = []
|
|
for tag_id in eligible:
|
|
prev = existing.get(tag_id)
|
|
if (
|
|
prev is None
|
|
or prev[1] != embedding_version
|
|
or prev[0] != fps.get(tag_id)
|
|
):
|
|
out.append(tag_id)
|
|
return out
|
|
|
|
|
|
def train_all_heads(
|
|
session: Session, params: dict[str, Any], run: HeadTrainingRun | None = None
|
|
) -> dict[str, int]:
|
|
"""(Re)train eligible concept heads, INCREMENTALLY by default (#1317 p2):
|
|
refit only the tags whose training data changed since last fit, so a manual
|
|
Retrain click is fast. `params["full"]=True` (the nightly run) refits every
|
|
head to reconcile sampled-negative + hygiene drift. Prunes heads whose tag is
|
|
no longer eligible. Commits per head so a SIGKILL leaves trained heads durable.
|
|
Returns {n_trained, n_skipped} (n_skipped = unchanged + too-few-examples)."""
|
|
import numpy as np
|
|
|
|
cfg = _normalize_params(session, params)
|
|
embedding_version = _embedder_version(session)
|
|
full = bool((params or {}).get("full"))
|
|
eligible = _eligible_tag_ids(session, cfg["min_positives"])
|
|
eligible_set = set(eligible)
|
|
# Computed once per run, not per head — the hygiene set is identical for
|
|
# every non-system concept.
|
|
hygiene = _hygiene_excluded_ids(session)
|
|
fps = _head_fingerprints(session, eligible)
|
|
to_train = set(
|
|
_heads_needing_retrain(session, eligible, embedding_version, fps, full)
|
|
)
|
|
trained = 0
|
|
failed = 0
|
|
for i, tag_id in enumerate(eligible):
|
|
if tag_id not in to_train:
|
|
continue
|
|
try:
|
|
ok = train_head(
|
|
session, tag_id, embedding_version, cfg, np, hygiene=hygiene
|
|
)
|
|
except Exception:
|
|
log.exception("train_head failed for tag %d", tag_id)
|
|
ok = False
|
|
if ok:
|
|
# Stamp the fingerprint we trained against so an unchanged tag is
|
|
# skipped on the next incremental run.
|
|
head = session.get(TagHead, tag_id)
|
|
if head is not None:
|
|
head.train_fingerprint = fps.get(tag_id)
|
|
session.commit()
|
|
trained += int(ok)
|
|
failed += 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()
|
|
# n_skipped = unchanged (not attempted) + failed-to-fit (too few examples).
|
|
return {
|
|
"n_trained": trained,
|
|
"n_skipped": (len(eligible) - len(to_train)) + failed,
|
|
}
|
|
|
|
|
|
def head_training_ids(
|
|
session: Session, tag_id: int, cfg: dict, hygiene: set[int] | None = None,
|
|
) -> tuple[list[int], list[int]] | None:
|
|
"""Select (pos_ids, neg_ids) for one head. Split out of train_head and
|
|
kept sklearn-free so the hygiene exclusion is testable in the CI env
|
|
(sklearn only exists in the ml image). Returns None when the concept has
|
|
too few usable positives.
|
|
|
|
Training hygiene (#128): images carrying a system tag are ABSENT from
|
|
every other concept's training — dropped as positives AND kept out of
|
|
the rejection/sampled negative pool (see _hygiene_excluded_ids). A system
|
|
tag's own head trains on them unfiltered: its positives ARE the hygiene
|
|
images."""
|
|
tag = session.get(Tag, tag_id)
|
|
if tag is not None and tag.is_system:
|
|
hygiene = set()
|
|
elif hygiene is None:
|
|
hygiene = _hygiene_excluded_ids(session)
|
|
|
|
pos_ids = [i for i in _ids_with_tag(session, tag_id) if i not in hygiene]
|
|
if len(pos_ids) < cfg["min_positives"]:
|
|
return None
|
|
|
|
pos_set = set(pos_ids)
|
|
rejected = [
|
|
i for i in _rejected_ids(session, tag_id)
|
|
if i not in pos_set and i not in hygiene
|
|
]
|
|
want_neg = max(len(pos_ids) * cfg["neg_ratio"], _EXAMPLES_MIN * 4)
|
|
sampled = _sample_unlabeled(
|
|
session, pos_set | set(rejected) | hygiene,
|
|
min(_UNLABELED_POOL, want_neg),
|
|
)
|
|
return pos_ids, rejected + [i for i in sampled if i not in pos_set]
|
|
|
|
|
|
def train_head(
|
|
session: Session, tag_id: int, embedding_version: str, cfg: dict, np,
|
|
hygiene: set[int] | None = None,
|
|
) -> 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
|
|
|
|
ids = head_training_ids(session, tag_id, cfg, hygiene)
|
|
if ids is None:
|
|
session.execute(delete(TagHead).where(TagHead.tag_id == tag_id))
|
|
return False
|
|
pos_ids, neg_ids = ids
|
|
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, Tag.is_system,
|
|
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,
|
|
# System tags (wip/banner/editor) are kind=general but group under
|
|
# their OWN "system" suggestion category so the operator reviews
|
|
# them apart from content tags (they still train as general heads).
|
|
"category": "system" if r.is_system else _CATEGORY.get(r.kind, "general"),
|
|
"auto_apply_threshold": r.auto_apply_threshold,
|
|
"is_system": bool(r.is_system),
|
|
}
|
|
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 _image_bag(
|
|
session: AsyncSession, image_id: int, cur_version: str,
|
|
) -> tuple[list, list[dict | None]]:
|
|
"""The max-over-bag inputs for one image: the whole-image SigLIP vector (when
|
|
it's in the current model's space) PLUS every concept-region crop embedded in
|
|
that space. Returns (bag, bag_meta) as PARALLEL lists — bag_meta[i] is None for
|
|
the whole-image row, else the region's {bbox, kind, detector} so a surfaced tag
|
|
can point back at the crop that produced it (#1206 grounding).
|
|
|
|
Only current-version embeddings enter the bag: mid model-swap (#1190) an image
|
|
still carrying an 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). Shared by live scoring (score_image)
|
|
and on-demand applied-tag grounding (ground_applied_tag, #1206 Step 4)."""
|
|
import numpy as np
|
|
|
|
img = await session.get(ImageRecord, image_id)
|
|
bag: list = []
|
|
bag_meta: list[dict | None] = []
|
|
if img is None:
|
|
return bag, bag_meta
|
|
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))
|
|
bag_meta.append(None)
|
|
region_rows = (
|
|
await session.execute(
|
|
select(
|
|
ImageRegion.siglip_embedding,
|
|
ImageRegion.rx, ImageRegion.ry, ImageRegion.rw, ImageRegion.rh,
|
|
ImageRegion.kind, ImageRegion.detector_version,
|
|
)
|
|
.where(ImageRegion.image_record_id == image_id)
|
|
.where(ImageRegion.siglip_embedding.is_not(None))
|
|
.where(ImageRegion.embedding_version == cur_version)
|
|
)
|
|
).all()
|
|
for vec, rx, ry, rw, rh, kind, detector in region_rows:
|
|
if vec is not None:
|
|
bag.append(np.asarray(vec, dtype=np.float32))
|
|
bag_meta.append(
|
|
{"bbox": [rx, ry, rw, rh], "kind": kind, "detector": detector}
|
|
)
|
|
return bag, bag_meta
|
|
|
|
|
|
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). System-tag heads (wip/banner/editor) instead use a flat
|
|
_SYSTEM_TAG_SUGGEST_FLOOR so their false positives surface for rejection
|
|
(still overridden by threshold_override). 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
|
|
|
|
settings = await _settings_async(session)
|
|
cur_version = settings.embedder_model_version
|
|
heads = await _current_heads(session, cur_version)
|
|
if heads["W"] is None:
|
|
return []
|
|
bag, bag_meta = await _image_bag(session, image_id, cur_version)
|
|
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_bag = 1.0 / (1.0 + np.exp(-Z)) # (B, H)
|
|
probs = probs_bag.max(axis=0) # (H,) best over the bag
|
|
# ARGMAX beside the max: WHICH bag row won each head → the region that grounds
|
|
# the tag (bag_meta[win]); None when the whole-image vector won (#1206).
|
|
winners = probs_bag.argmax(axis=0) # (H,)
|
|
out = []
|
|
for i, p in enumerate(probs):
|
|
if threshold_override is not None:
|
|
cut = threshold_override
|
|
elif heads["meta"][i]["is_system"]:
|
|
# System tags surface at the flat floor (see _SYSTEM_TAG_SUGGEST_FLOOR)
|
|
# so their false positives show up for the operator to reject.
|
|
cut = _SYSTEM_TAG_SUGGEST_FLOOR
|
|
else:
|
|
cut = 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),
|
|
"grounding": bag_meta[int(winners[i])],
|
|
})
|
|
out.sort(key=lambda d: d["score"], reverse=True)
|
|
return out
|
|
|
|
|
|
async def ground_applied_tag(
|
|
session: AsyncSession, image_id: int, tag_id: int,
|
|
) -> tuple[dict | None, bool]:
|
|
"""On-demand grounding for an ALREADY-APPLIED tag (#1206 Step 4). Applied tags
|
|
aren't scored live, so recompute the max-over-bag argmax for just this tag's
|
|
head — which crop region best explains the tag on this image — mirroring what
|
|
score_image records for live suggestions. Returns (grounding, has_head):
|
|
|
|
- has_head False → the tag has no head in the current embedding space (manual/
|
|
artist/meta tags, or a concept below the head floor). Nothing to localize
|
|
with, so the UI shows no overlay (distinct from "the whole image won").
|
|
- grounding None (has_head True) → the whole-image vector best explains it,
|
|
not any crop; the UI shows the subtle whole-image frame.
|
|
- grounding {bbox, kind, detector} → the winning region.
|
|
|
|
Character heads are covered too (character is a head kind); this deliberately
|
|
reuses the SigLIP head bag rather than the CCIP figure path so every applied
|
|
concept grounds through one consistent mechanism."""
|
|
import numpy as np
|
|
|
|
cur_version = (await _settings_async(session)).embedder_model_version
|
|
row = (
|
|
await session.execute(
|
|
select(TagHead.weights, TagHead.bias).where(
|
|
TagHead.tag_id == tag_id,
|
|
TagHead.embedding_version == cur_version,
|
|
)
|
|
)
|
|
).one_or_none()
|
|
if row is None:
|
|
return None, False
|
|
bag, bag_meta = await _image_bag(session, image_id, cur_version)
|
|
if not bag:
|
|
return None, True
|
|
|
|
X = np.vstack(bag)
|
|
norms = np.linalg.norm(X, axis=1, keepdims=True)
|
|
norms[norms == 0] = 1.0
|
|
Xn = X / norms
|
|
# The sigmoid is monotonic in the logit, so the highest-probability bag row is
|
|
# just argmax of the raw score — no need to exponentiate to pick the winner.
|
|
z = Xn @ np.asarray(row.weights, dtype=np.float32) + float(row.bias) # (B,)
|
|
return bag_meta[int(z.argmax())], True
|
|
|
|
|
|
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}
|
|
|
|
|
|
def retract_auto_applied_heads(session: Session) -> int:
|
|
"""Soft auto-apply (milestone 139): re-score every standing source='head_auto'
|
|
tag against its CURRENT head and REMOVE the ones now BELOW the head's
|
|
auto_apply_threshold — i.e. the head sharpened (or the operator raised the bar)
|
|
and no longer supports them. Skips operator-confirmed tags
|
|
(TagPositiveConfirmation). SILENT: a low score isn't proof the tag was wrong,
|
|
so no hard negative is recorded — that's reserved for an operator removal.
|
|
No-op unless head_auto_apply_enabled. Only re-scores the images that ALREADY
|
|
carry the auto-tag (bounded), never the whole library. Returns n_retracted."""
|
|
import numpy as np
|
|
|
|
settings = _settings(session)
|
|
if not settings.head_auto_apply_enabled:
|
|
return 0
|
|
heads = session.execute(
|
|
select(
|
|
TagHead.tag_id, TagHead.weights, TagHead.bias,
|
|
TagHead.auto_apply_threshold,
|
|
)
|
|
.where(TagHead.embedding_version == settings.embedder_model_version)
|
|
.where(TagHead.auto_apply_threshold.is_not(None))
|
|
).all()
|
|
retracted = 0
|
|
for tag_id, weights, bias, thr in heads:
|
|
auto_ids = [
|
|
iid for (iid,) in session.execute(
|
|
select(image_tag.c.image_record_id)
|
|
.where(image_tag.c.tag_id == tag_id)
|
|
.where(image_tag.c.source == "head_auto")
|
|
)
|
|
]
|
|
if not auto_ids:
|
|
continue
|
|
confirmed = {
|
|
iid for (iid,) in session.execute(
|
|
select(TagPositiveConfirmation.image_record_id)
|
|
.where(TagPositiveConfirmation.tag_id == tag_id)
|
|
.where(TagPositiveConfirmation.image_record_id.in_(auto_ids))
|
|
)
|
|
}
|
|
candidates = [i for i in auto_ids if i not in confirmed]
|
|
emb = _load_embeddings(session, candidates)
|
|
cids = [i for i in candidates if i in emb]
|
|
if not cids:
|
|
continue
|
|
Xn = _l2norm(np.vstack([emb[i] for i in cids]).astype(np.float32), np)
|
|
w = np.asarray(weights, dtype=np.float32)
|
|
probs = 1.0 / (1.0 + np.exp(-(Xn @ w + float(bias))))
|
|
below = [cids[k] for k in np.where(probs < float(thr))[0]]
|
|
for iid in below:
|
|
session.execute(
|
|
image_tag.delete()
|
|
.where(image_tag.c.image_record_id == iid)
|
|
.where(image_tag.c.tag_id == tag_id)
|
|
.where(image_tag.c.source == "head_auto")
|
|
)
|
|
retracted += 1
|
|
session.commit()
|
|
return retracted
|