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/api/suggestions.py b/backend/app/api/suggestions.py
index 46fa2e5..7a5ca1d 100644
--- a/backend/app/api/suggestions.py
+++ b/backend/app/api/suggestions.py
@@ -61,6 +61,10 @@ async def get_suggestions(image_id: int):
# modal's "Treat as alias"/"Remove alias" affordances.
"raw_name": s.raw_name,
"via_alias": s.via_alias,
+ # operator dismissed this tag for this image — surfaced
+ # (not dropped) so the rail can show it rejected + offer
+ # one-click un-reject.
+ "rejected": s.rejected,
}
for s in items
]
@@ -131,6 +135,21 @@ async def dismiss_suggestion(image_id: int):
return "", 204
+@suggestions_bp.route(
+ "/images//suggestions/undismiss", methods=["POST"]
+)
+async def undismiss_suggestion(image_id: int):
+ """Reverse a per-image dismissal (reject-recovery). Idempotent — undoing a
+ tag that isn't rejected is a no-op delete."""
+ body = await request.get_json()
+ if not body or "tag_id" not in body:
+ return jsonify({"error": "tag_id required"}), 400
+ async with get_session() as session:
+ await AllowlistService(session).undismiss(image_id, body["tag_id"])
+ await session.commit()
+ return "", 204
+
+
@suggestions_bp.route("/suggestions/bulk", methods=["POST"])
async def bulk_suggestions():
body = await request.get_json()
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..04025ed 100644
--- a/backend/app/models/__init__.py
+++ b/backend/app/models/__init__.py
@@ -8,6 +8,7 @@ from .base import Base
from .credential import Credential
from .download_event import DownloadEvent
from .external_link import ExternalLink
+from .head_training_run import HeadTrainingRun
from .image_prediction import ImagePrediction
from .image_provenance import ImageProvenance
from .image_record import ImageRecord
@@ -30,6 +31,7 @@ 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/allowlist.py b/backend/app/services/ml/allowlist.py
index 8f7d14e..857d665 100644
--- a/backend/app/services/ml/allowlist.py
+++ b/backend/app/services/ml/allowlist.py
@@ -89,6 +89,12 @@ class AllowlistService:
)
await self.session.execute(stmt)
+ async def undismiss(self, image_id: int, tag_id: int) -> None:
+ """Undo a per-image dismissal — drop the TagSuggestionRejection so the
+ suggestion reverts to a live (un-rejected) state. Backs the rail's
+ one-click reject-recovery (operator-asked 2026-06-27)."""
+ await self._clear_rejection(image_id, tag_id)
+
async def reject_applied_tag(self, image_id: int, tag_id: int) -> None:
"""Operator removed an applied tag from an image. Remove the
image_tag row AND record a rejection so the allowlist won't
diff --git a/backend/app/services/ml/heads.py b/backend/app/services/ml/heads.py
new file mode 100644
index 0000000..aed4e69
--- /dev/null
+++ b/backend/app/services/ml/heads.py
@@ -0,0 +1,330 @@
+"""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
+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, 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."""
+ 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):
+ 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()
diff --git a/backend/app/services/ml/suggestions.py b/backend/app/services/ml/suggestions.py
index 7b0f996..1cb306f 100644
--- a/backend/app/services/ml/suggestions.py
+++ b/backend/app/services/ml/suggestions.py
@@ -1,24 +1,22 @@
-"""The suggestion read-path: raw predictions + centroids -> alias-resolved,
-threshold-filtered, category-grouped, ranked suggestions for one image.
+"""The suggestion read-path: trained HEADS score one image's frozen embedding
+into alias-resolved, category-grouped, ranked suggestions.
+
+Tagging-v2 (#114): suggestions now come from the per-concept heads that LEARN
+from the operator's tags (services/ml/heads.py) — the Camie prediction source
+and the per-tag SigLIP centroid have been REMOVED. A head exists only for an
+existing concept tag, so every suggestion is a canonical tag (no raw model key,
+no alias remap, no creates-new). Rejected tags stay in the list FLAGGED (not
+dropped) so the rail can show + reverse a dismissal.
"""
from dataclasses import dataclass, field
-from sqlalchemy import func, select
+from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
-from ...models import (
- ImagePrediction,
- ImageRecord,
- MLSettings,
- Tag,
- TagSuggestionRejection,
-)
+from ...models import ImageRecord, TagSuggestionRejection
from ...models.tag import image_tag
-from .aliases import AliasService
-from .centroids import CentroidService
-from .tag_name import normalize as normalize_tag_name
-from .tagger import SURFACED_CATEGORIES
+from .heads import score_image
@dataclass(frozen=True)
@@ -29,7 +27,7 @@ class Suggestion:
display_name: str
category: str
score: float
- source: str # 'tagger' | 'centroid' | 'both'
+ source: str # 'head' (Camie 'tagger'/'centroid' sources removed in v2)
creates_new_tag: bool
# raw_name = the booru model vocab key behind this suggestion. It's the key
# an alias MUST be stored under (resolution looks up the raw key), so the
@@ -39,6 +37,11 @@ class Suggestion:
# via_alias = this suggestion was surfaced because an operator alias remapped
# the raw prediction to this canonical tag. Lets the UI mark it + offer undo.
via_alias: bool = False
+ # rejected = the operator dismissed this tag for this image (a stored
+ # TagSuggestionRejection). It stays in the list — flagged, not dropped — so
+ # the rejection is VISIBLE and REVERSIBLE in the rail (misclick recovery,
+ # operator-asked 2026-06-27) instead of silently vanishing or re-suggesting.
+ rejected: bool = False
@dataclass
@@ -49,67 +52,24 @@ class SuggestionList:
class SuggestionService:
def __init__(self, session: AsyncSession):
self.session = session
- self.aliases = AliasService(session)
- self.centroids = CentroidService(session)
-
- async def _settings(self) -> MLSettings:
- return (
- await self.session.execute(select(MLSettings).where(MLSettings.id == 1))
- ).scalar_one()
-
- async def _load_predictions(self, image_id: int) -> dict:
- """Predictions for one image from the normalized image_prediction
- table (#768), in the {raw_name: {category, confidence}} shape the rest
- of this service consumed from the old JSON column — so all downstream
- threshold/alias/merge logic is unchanged."""
- rows = (
- await self.session.execute(
- select(
- ImagePrediction.raw_name,
- ImagePrediction.category,
- ImagePrediction.score,
- ).where(ImagePrediction.image_record_id == image_id)
- )
- ).all()
- return {
- r.raw_name: {"category": r.category, "confidence": r.score}
- for r in rows
- }
-
- def _threshold_for(
- self, s: MLSettings, category: str, override: float | None = None,
- ) -> float:
- # 'artist' (FC-2d-vii-c) and 'copyright' (2026-06-01) retired;
- # both fall through to the 1.01 "never surfaces" default like any
- # unsurfaced category.
- # override (the typed-dropdown "show everything the model saw" mode)
- # applies to the surfaced categories only — unsurfaced ones are already
- # skipped before the threshold check, so they can't leak in.
- if override is not None:
- return override
- return {
- "character": s.suggestion_threshold_character,
- "general": s.suggestion_threshold_general,
- }.get(category, 1.01)
async def for_image(
- self, image_id: int, *, threshold_override: float | None = None,
+ self, image_id: int, threshold_override: float | None = None,
) -> SuggestionList:
- """Ranked suggestions for one image.
+ """Head-scored suggestions for one image, grouped by category and ranked.
- threshold_override surfaces EVERY stored tagger prediction (down to the
- ingest STORE_FLOOR) regardless of the configured per-category suggestion
- thresholds — backs the tag-input dropdown's "search all of the model's
- predictions, including low-confidence ones, in the canonical formatting"
- mode (operator-asked 2026-06-09). The Suggestions panel still calls with
- no override so it stays the curated above-threshold list."""
+ Each trained head scores the image's frozen embedding; a concept surfaces
+ when its score clears the head's own suggest threshold. threshold_override
+ (used by the typed tag-input dropdown's "show everything" mode) replaces
+ that per-head cut with a flat floor (0 → every head), so a low-scoring
+ concept can still be typed + picked in canonical formatting.
+
+ Already-applied tags are dropped; rejected tags stay FLAGGED and sink to
+ the bottom of their category so a dismissal is visible + reversible."""
img = await self.session.get(ImageRecord, image_id)
if img is None:
return SuggestionList()
- settings = await self._settings()
- predictions: dict = await self._load_predictions(image_id)
-
applied = set(
(
await self.session.execute(
@@ -129,148 +89,30 @@ class SuggestionService:
).scalars().all()
)
- # --- Camie predictions ---
- # candidates carry (raw_name, display_name, category, confidence).
- # raw_name = the booru-formatted vocab key, kept for alias_map
- # lookup since alias rows are hand-curated against raw keys.
- # display_name = normalize_tag_name(raw_name) — what the operator
- # sees AND what gets written to tag.name on Accept.
- candidates: list[tuple[str, str, str, float]] = []
- for name, p in predictions.items():
- category = p.get("category", "general")
- if category not in SURFACED_CATEGORIES:
- continue
- conf = float(p.get("confidence", 0.0))
- if conf < self._threshold_for(settings, category, threshold_override):
- continue
- display = normalize_tag_name(name)
- if display is None:
- # emoticon / pure-punctuation vocab entry — drop entirely
- continue
- candidates.append((name, display, category, conf))
-
- alias_map = await self.aliases.resolve_many(
- [(raw, c) for raw, _disp, c, _conf in candidates]
+ hits = await score_image(
+ self.session, image_id, threshold_override=threshold_override
)
- merged: dict[object, Suggestion] = {}
-
- def _merge(key, sug: Suggestion):
- existing = merged.get(key)
- if existing is None:
- merged[key] = sug
- elif sug.score > existing.score:
- merged[key] = Suggestion(
- canonical_tag_id=existing.canonical_tag_id,
- display_name=existing.display_name,
- category=existing.category,
- score=sug.score,
- source="both"
- if existing.source != sug.source
- else existing.source,
- creates_new_tag=existing.creates_new_tag,
- # Keep the alias identity from `existing`: the tagger pass
- # (which carries raw_name / via_alias) runs before centroid
- # augmentation, so it's always the first writer for a key.
- raw_name=existing.raw_name,
- via_alias=existing.via_alias,
- )
-
- for raw, display, category, conf in candidates:
- canonical = alias_map.get((raw, category))
- if canonical is not None:
- if canonical.id in applied or canonical.id in rejected:
- continue
- _merge(
- canonical.id,
- Suggestion(
- canonical_tag_id=canonical.id,
- display_name=canonical.name,
- category=category,
- score=conf,
- source="tagger",
- creates_new_tag=False,
- raw_name=raw,
- via_alias=True,
- ),
- )
- else:
- # Case-insensitive match on BOTH the raw camie key AND
- # the normalized form — covers legacy underscore-named
- # Tag rows accepted before normalization shipped, AND
- # any tag the operator created with the human form.
- existing_tag = (
- await self.session.execute(
- select(Tag).where(
- func.lower(Tag.name).in_(
- [raw.lower(), display.lower()]
- )
- )
- )
- ).scalars().first()
- if existing_tag is not None:
- if (
- existing_tag.id in applied
- or existing_tag.id in rejected
- ):
- continue
- _merge(
- existing_tag.id,
- Suggestion(
- canonical_tag_id=existing_tag.id,
- display_name=existing_tag.name,
- category=category,
- score=conf,
- source="tagger",
- creates_new_tag=False,
- raw_name=raw,
- via_alias=False,
- ),
- )
- else:
- _merge(
- f"raw:{display}:{category}",
- Suggestion(
- canonical_tag_id=None,
- display_name=display,
- category=category,
- score=conf,
- source="tagger",
- creates_new_tag=True,
- raw_name=raw,
- via_alias=False,
- ),
- )
-
- # --- Centroid augmentation ---
- hits = await self.centroids.find_similar_tags(image_id, limit=30)
- for hit in hits:
- if hit.similarity < settings.centroid_similarity_threshold:
- continue
- if hit.tag_id in applied or hit.tag_id in rejected:
- continue
- tag = await self.session.get(Tag, hit.tag_id)
- if tag is None:
- continue
- cat = tag.kind.value if hasattr(tag.kind, "value") else str(tag.kind)
- display_cat = cat if cat in SURFACED_CATEGORIES else "general"
- _merge(
- tag.id,
- Suggestion(
- canonical_tag_id=tag.id,
- display_name=tag.name,
- category=display_cat,
- score=hit.similarity,
- source="centroid",
- creates_new_tag=False,
- ),
- )
-
result = SuggestionList()
- for sug in merged.values():
- result.by_category.setdefault(sug.category, []).append(sug)
+ for h in hits:
+ tag_id = h["tag_id"]
+ if tag_id in applied:
+ continue
+ result.by_category.setdefault(h["category"], []).append(
+ Suggestion(
+ canonical_tag_id=tag_id,
+ display_name=h["name"],
+ category=h["category"],
+ score=h["score"],
+ source="head",
+ creates_new_tag=False,
+ rejected=tag_id in rejected,
+ )
+ )
for cat in result.by_category:
- result.by_category[cat].sort(key=lambda s: s.score, reverse=True)
+ # Live suggestions first (by score), rejected ones sink to the
+ # bottom of the category — visible for recovery, out of the way.
+ result.by_category[cat].sort(key=lambda s: (s.rejected, -s.score))
return result
async def for_selection(
@@ -297,6 +139,11 @@ class SuggestionService:
for s in items:
if s.canonical_tag_id is None or s.creates_new_tag:
continue
+ # for_image keeps rejected tags (flagged) for the rail;
+ # bulk consensus must still ignore them — a tag dismissed on
+ # an image isn't a suggestion for that image.
+ if s.rejected:
+ continue
st = stats.get(s.canonical_tag_id)
if st is None:
st = {
diff --git a/backend/app/tasks/maintenance.py b/backend/app/tasks/maintenance.py
index 1c8a0b8..fa0397b 100644
--- a/backend/app/tasks/maintenance.py
+++ b/backend/app/tasks/maintenance.py
@@ -13,6 +13,7 @@ from ..celery_app import celery
from ..models import (
BackupRun,
DownloadEvent,
+ HeadTrainingRun,
ImageRecord,
ImportBatch,
ImportSettings,
@@ -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/frontend/src/components/modal/SuggestionItem.vue b/frontend/src/components/modal/SuggestionItem.vue
index 1b79848..b9df764 100644
--- a/frontend/src/components/modal/SuggestionItem.vue
+++ b/frontend/src/components/modal/SuggestionItem.vue
@@ -1,29 +1,54 @@
-
+ 2026-06-01). The row itself is informational; the green ✓ / red ✗
+ verdict pair + 3-dot alias menu are the action affordances. -->
+
{{ suggestion.display_name }}
- rejected
+ + new
alias
{{ scorePct }}
-
- Accept
-
+
+
+ mdi-check
+ mdi-undo-variant
+ mdi-close
+
+ teleported image modal — #711). Only rendered when an alias action
+ applies — dismiss now lives on the red ✗, so a centroid hit with no
+ alias option has no menu. -->
@@ -54,9 +76,15 @@ import { computed } from 'vue'
import KebabMenu from '../common/KebabMenu.vue'
const props = defineProps({ suggestion: { type: Object, required: true } })
-defineEmits(['accept', 'alias', 'remove-alias', 'dismiss'])
+defineEmits(['accept', 'alias', 'remove-alias', 'dismiss', 'undismiss'])
const scorePct = computed(() => `${Math.round(props.suggestion.score * 100)}%`)
+// Kebab now only carries alias actions: show it when this suggestion can be
+// aliased (raw model key, not yet aliased) or is already aliased (so it can be
+// un-aliased). Centroid hits (no raw_name, no alias) have an empty menu → hide.
+const hasMenu = computed(() =>
+ Boolean(props.suggestion.raw_name) || Boolean(props.suggestion.via_alias)
+)
diff --git a/frontend/src/components/modal/SuggestionsCategoryGroup.vue b/frontend/src/components/modal/SuggestionsCategoryGroup.vue
index c2b35bc..4a317f6 100644
--- a/frontend/src/components/modal/SuggestionsCategoryGroup.vue
+++ b/frontend/src/components/modal/SuggestionsCategoryGroup.vue
@@ -18,6 +18,7 @@
@alias="$emit('alias', $event)"
@remove-alias="$emit('remove-alias', $event)"
@dismiss="$emit('dismiss', $event)"
+ @undismiss="$emit('undismiss', $event)"
/>
@@ -33,7 +34,7 @@ const props = defineProps({
collapsible: { type: Boolean, default: false },
defaultOpen: { type: Boolean, default: true }
})
-defineEmits(['accept', 'alias', 'remove-alias', 'dismiss'])
+defineEmits(['accept', 'alias', 'remove-alias', 'dismiss', 'undismiss'])
const open = ref(props.collapsible ? props.defaultOpen : true)
diff --git a/frontend/src/components/modal/SuggestionsPanel.vue b/frontend/src/components/modal/SuggestionsPanel.vue
index ffa2804..e62fc8b 100644
--- a/frontend/src/components/modal/SuggestionsPanel.vue
+++ b/frontend/src/components/modal/SuggestionsPanel.vue
@@ -21,14 +21,14 @@
v-show="store.byCategory[cat] && store.byCategory[cat].length"
:label="labelFor(cat)" :items="store.byCategory[cat] || []"
@accept="onAccept" @alias="onAlias" @remove-alias="onRemoveAlias"
- @dismiss="store.dismiss"
+ @dismiss="store.dismiss" @undismiss="store.undismiss"
/>
diff --git a/frontend/src/components/modal/TagAutocomplete.vue b/frontend/src/components/modal/TagAutocomplete.vue
index ba6fe25..43ab0e1 100644
--- a/frontend/src/components/modal/TagAutocomplete.vue
+++ b/frontend/src/components/modal/TagAutocomplete.vue
@@ -202,6 +202,10 @@ const suggestionHits = computed(() => {
const out = []
for (const list of Object.values(suggestions.allByCategory)) {
for (const s of list || []) {
+ // Rejected suggestions now stay in allByCategory (flagged) so the panel
+ // can show + un-reject them; keep them OUT of the type-to-add dropdown,
+ // whose job is finding a tag to ADD (un-reject lives in the panel).
+ if (s.rejected) continue
const key = `${s.category}:${s.display_name.toLowerCase()}`
if (!s.display_name.toLowerCase().includes(q)) continue
if (seen.has(key)) continue
diff --git a/frontend/src/components/settings/HeadsCard.vue b/frontend/src/components/settings/HeadsCard.vue
new file mode 100644
index 0000000..673ff1f
--- /dev/null
+++ b/frontend/src/components/settings/HeadsCard.vue
@@ -0,0 +1,241 @@
+
+
+
+ A head is a tiny classifier trained on the SigLIP
+ embeddings already stored on your images — your positives plus your
+ negatives (rejections). One is built per general/character concept with at
+ least {{ minPositives }} tagged images. Retrain after a
+ tagging session to fold in your latest accepts/rejects; scoring is live, so
+ the rail reflects a retrain on the next image you open.
+
+
+
+
+
+
{{ headCount }}
+
heads
+
+
+
{{ graduatedCount }}
+
auto-apply ready
+
+
+
{{ lastTrained }}
+
last trained
+
+
+
+ {{ headCount > 0 ? 'Retrain heads' : 'Train heads' }}
+
+
+
+
Training… (started {{ startedAgo }})
+
+
+ Training failed: {{ lastError }}
+
+
+
+
mdi-brain
+
+ No heads yet. Tag a handful of images for the concepts you care about,
+ then train — each concept with ≥ {{ minPositives }} tags becomes a head.
+
+
+
+
+
+
+ {{ heads.length }} concept{{ heads.length === 1 ? '' : 's' }}, strongest first
+ (AP = average precision; auto-apply ⚡ = precise enough to fire without review)
+
+
+
+
+
+ Concept
+ Cat
+ +tags
+ AP
+ P
+ R
+ ⚡
+
+
+
+
+ {{ h.name }}
+ {{ h.category }}
+ {{ h.n_pos }}
+ {{ pct(h.ap) }}
+ {{ pct(h.precision) }}
+ {{ pct(h.recall) }}
+
+ mdi-lightning-bolt
+ —
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/frontend/src/components/settings/MaintenancePanel.vue b/frontend/src/components/settings/MaintenancePanel.vue
index 68c30d1..44dafe6 100644
--- a/frontend/src/components/settings/MaintenancePanel.vue
+++ b/frontend/src/components/settings/MaintenancePanel.vue
@@ -26,6 +26,7 @@
+
@@ -52,6 +53,7 @@ import ArchiveReextractCard from './ArchiveReextractCard.vue'
import MissingFileRepairCard from './MissingFileRepairCard.vue'
import DbMaintenanceCard from './DbMaintenanceCard.vue'
import MLThresholdSliders from './MLThresholdSliders.vue'
+import HeadsCard from './HeadsCard.vue'
import AllowlistTable from './AllowlistTable.vue'
import AliasTable from './AliasTable.vue'
import TagEvalCard from './TagEvalCard.vue'
diff --git a/frontend/src/stores/heads.js b/frontend/src/stores/heads.js
new file mode 100644
index 0000000..c1f481e
--- /dev/null
+++ b/frontend/src/stores/heads.js
@@ -0,0 +1,24 @@
+import { defineStore } from 'pinia'
+
+import { useApi } from '../composables/useApi.js'
+
+// Heads (#114): the per-concept classifiers that LEARN from your tags and power
+// suggestions (replacing Camie + centroid). Training runs as a background task;
+// the card rehydrates status from GET /api/heads on mount so it survives
+// navigation (the run lives in head_training_run server-side).
+export const useHeadsStore = defineStore('heads', () => {
+ const api = useApi()
+
+ // Summary: head_count, graduated_count, last_trained_at, running_id, the
+ // per-concept head table, and recent training runs.
+ async function status() {
+ return await api.get('/api/heads')
+ }
+
+ // (Re)train all eligible heads. One run at a time (409 if already running).
+ async function train(params = {}) {
+ return await api.post('/api/heads/train', { body: { params } })
+ }
+
+ return { status, train }
+})
diff --git a/frontend/src/stores/suggestions.js b/frontend/src/stores/suggestions.js
index 36c926a..b7dd6c5 100644
--- a/frontend/src/stores/suggestions.js
+++ b/frontend/src/stores/suggestions.js
@@ -84,6 +84,19 @@ export const useSuggestionsStore = defineStore('suggestions', () => {
}
}
+ // Flip the `rejected` flag on the matching suggestion in BOTH lists in place,
+ // so a reject/un-reject shows immediately without dropping the row (visible,
+ // reversible rejection — misclick recovery, operator-asked 2026-06-27).
+ function _setRejectedEverywhere(suggestion, value) {
+ const key = _keyOf(suggestion)
+ const cat = suggestion.category
+ for (const map of [byCategory.value, allByCategory.value]) {
+ for (const s of map[cat] || []) {
+ if (_keyOf(s) === key) s.rejected = value
+ }
+ }
+ }
+
async function accept(suggestion) {
// Capture imageId so a mid-flight prev/next can't reroute the
// accept POST to a different image AND push the tag to that
@@ -168,21 +181,36 @@ export const useSuggestionsStore = defineStore('suggestions', () => {
async function dismiss(suggestion) {
const imageId = currentImageId
if (imageId == null) return
- // Dismiss needs a tag_id; raw tags have none, so dismissing a raw
- // suggestion just hides it client-side (nothing to persist a rejection
- // against until the tag exists).
- if (suggestion.canonical_tag_id != null) {
- await api.post(`/api/images/${imageId}/suggestions/dismiss`, {
- body: { tag_id: suggestion.canonical_tag_id }
- })
+ // Dismiss needs a tag_id. Raw tags (creates_new_tag) have none, so there's
+ // nothing to persist a rejection against — drop them client-side as before.
+ // Canonical tags persist a rejection and STAY in the list flagged rejected,
+ // so the operator can see it and one-click un-reject (misclick recovery).
+ if (suggestion.canonical_tag_id == null) {
+ if (currentImageId === imageId) _dropEverywhere(suggestion)
+ return
}
+ await api.post(`/api/images/${imageId}/suggestions/dismiss`, {
+ body: { tag_id: suggestion.canonical_tag_id }
+ })
if (currentImageId === imageId) {
- _dropEverywhere(suggestion)
+ _setRejectedEverywhere(suggestion, true)
+ }
+ }
+
+ // Undo a per-image dismissal — the suggestion reverts to a live row.
+ async function undismiss(suggestion) {
+ const imageId = currentImageId
+ if (imageId == null || suggestion.canonical_tag_id == null) return
+ await api.post(`/api/images/${imageId}/suggestions/undismiss`, {
+ body: { tag_id: suggestion.canonical_tag_id }
+ })
+ if (currentImageId === imageId) {
+ _setRejectedEverywhere(suggestion, false)
}
}
return {
byCategory, allByCategory, loading, error,
- load, loadAll, accept, aliasAccept, removeAlias, dismiss
+ load, loadAll, accept, aliasAccept, removeAlias, dismiss, undismiss
}
})
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)
diff --git a/tests/test_api_suggestions.py b/tests/test_api_suggestions.py
index 2b1b425..ac02374 100644
--- a/tests/test_api_suggestions.py
+++ b/tests/test_api_suggestions.py
@@ -1,7 +1,8 @@
import pytest
+from sqlalchemy import select
from backend.app.celery_app import celery
-from backend.app.models import ImageRecord, TagKind
+from backend.app.models import ImageRecord, MLSettings, TagHead, TagKind
from backend.app.services.tag_service import TagService
pytestmark = pytest.mark.integration
@@ -31,13 +32,30 @@ async def _img(db, preds, sha="s" * 64):
@pytest.mark.asyncio
async def test_get_suggestions(client, db):
- img = await _img(
- db, {"sword": {"category": "general", "confidence": 0.97}}
+ # Suggestions come from a trained head now (Camie/centroid removed): an image
+ # whose embedding aligns with the head surfaces that concept.
+ s = (await db.execute(select(MLSettings).where(MLSettings.id == 1))).scalar_one()
+ img = ImageRecord(
+ path="/images/headsug.jpg", sha256="h" * 64, size_bytes=1,
+ mime="image/jpeg", width=1, height=1, origin="imported_filesystem",
+ integrity_status="unknown", siglip_embedding=[3.0] + [0.0] * 1151,
)
+ db.add(img)
+ await db.flush()
+ tag = await TagService(db).find_or_create("sword", TagKind.general)
+ db.add(TagHead(
+ tag_id=tag.id, embedding_version=s.embedder_model_version,
+ weights=[1.0] + [0.0] * 1151, 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()
resp = await client.get(f"/api/images/{img.id}/suggestions")
assert resp.status_code == 200
body = await resp.get_json()
- assert "general" in body["by_category"]
+ general = body["by_category"].get("general", [])
+ s2 = next(x for x in general if x["canonical_tag_id"] == tag.id)
+ assert s2["source"] == "head"
@pytest.mark.asyncio
@@ -95,6 +113,25 @@ async def test_dismiss(client, db):
assert resp.status_code == 204
+@pytest.mark.asyncio
+async def test_undismiss_reverses_rejection(client, db):
+ img = await _img(db, {})
+ tag = await TagService(db).find_or_create("UndismissMe", TagKind.general)
+ await db.commit()
+ await client.post(
+ f"/api/images/{img.id}/suggestions/dismiss", json={"tag_id": tag.id}
+ )
+ resp = await client.post(
+ f"/api/images/{img.id}/suggestions/undismiss", json={"tag_id": tag.id}
+ )
+ assert resp.status_code == 204
+ # Idempotent: un-rejecting again (nothing to clear) is still a 204.
+ resp2 = await client.post(
+ f"/api/images/{img.id}/suggestions/undismiss", json={"tag_id": tag.id}
+ )
+ assert resp2.status_code == 204
+
+
@pytest.mark.asyncio
async def test_alias_requires_fields(client, db):
img = await _img(db, {})
@@ -102,67 +139,3 @@ async def test_alias_requires_fields(client, db):
f"/api/images/{img.id}/suggestions/alias", json={"alias_string": "x"}
)
assert resp.status_code == 400
-
-
-async def _img_at(db, path, sha, preds):
- from tests._prediction_helpers import seed_predictions
-
- img = ImageRecord(
- path=path, sha256=sha, size_bytes=1, mime="image/jpeg",
- width=1, height=1, origin="imported_filesystem",
- integrity_status="unknown",
- )
- db.add(img)
- await db.commit()
- await seed_predictions(db, img.id, preds)
- await db.commit()
- return img
-
-
-@pytest.mark.asyncio
-async def test_alias_roundtrip_resolves_by_raw_key(client, db):
- """Locks the modal-alias contract: the suggestion exposes the RAW model key,
- an alias authored with that key resolves on a later image, and the resolved
- suggestion is flagged via_alias. (Pre-fix the modal stored the normalized
- display name, which never resolved.)"""
- canonical = await TagService(db).find_or_create(
- "Sasuke Uchiha", TagKind.character
- )
- await db.commit()
- preds = {"uchiha_sasuke": {"category": "character", "confidence": 0.99}}
- img_a = await _img_at(db, "/images/alias_a.jpg", "a" * 64, preds)
-
- # (a) raw_name is exposed so the modal can author the alias with it; the
- # raw prediction doesn't textually match the tag, so it'd otherwise be +new.
- body = await (
- await client.get(f"/api/images/{img_a.id}/suggestions")
- ).get_json()
- sug = body["by_category"]["character"][0]
- assert sug["raw_name"] == "uchiha_sasuke"
- assert sug["via_alias"] is False
- assert sug["creates_new_tag"] is True
-
- # Author the alias keyed by the RAW key (what the frontend now sends).
- resp = await client.post(
- f"/api/images/{img_a.id}/suggestions/alias",
- json={
- "alias_string": sug["raw_name"],
- "alias_category": "character",
- "canonical_tag_id": canonical.id,
- },
- )
- assert resp.status_code == 200
- assert (await resp.get_json())["allowlisted"] is True
-
- # (b) A DIFFERENT image with the same prediction now resolves via the alias
- # (image A's tag is applied, so it's filtered there). Had the alias been
- # stored under the display name, this would NOT resolve.
- img_b = await _img_at(db, "/images/alias_b.jpg", "b" * 64, preds)
- body_b = await (
- await client.get(f"/api/images/{img_b.id}/suggestions")
- ).get_json()
- sug_b = body_b["by_category"]["character"][0]
- assert sug_b["canonical_tag_id"] == canonical.id
- assert sug_b["via_alias"] is True
- assert sug_b["creates_new_tag"] is False
- assert sug_b["raw_name"] == "uchiha_sasuke"
diff --git a/tests/test_ml_artist_retired.py b/tests/test_ml_artist_retired.py
index 955583a..0a3e208 100644
--- a/tests/test_ml_artist_retired.py
+++ b/tests/test_ml_artist_retired.py
@@ -14,14 +14,12 @@ def test_artist_not_centroid_eligible():
assert TagKind.artist not in ELIGIBLE_KINDS
-def test_threshold_for_artist_is_unsurfaced():
- from backend.app.services.ml.suggestions import SuggestionService
+def test_artist_not_head_eligible():
+ # Tagging-v2: suggestions come from heads, and heads are only trained for
+ # general/character concepts — so 'artist' (and any other kind) can't surface.
+ from backend.app.models import TagKind
+ from backend.app.services.ml.heads import _HEAD_KINDS
- class _S:
- suggestion_threshold_character = 0.5
- suggestion_threshold_general = 0.5
-
- svc = SuggestionService.__new__(SuggestionService)
- # 'artist' and 'copyright' both retired — fall through to 1.01
- assert svc._threshold_for(_S(), "artist") == 1.01
- assert svc._threshold_for(_S(), "copyright") == 1.01
+ assert TagKind.general in _HEAD_KINDS
+ assert TagKind.character in _HEAD_KINDS
+ assert TagKind.artist not in _HEAD_KINDS
diff --git a/tests/test_ml_suggestions.py b/tests/test_ml_suggestions.py
index 3f532bd..7a427e2 100644
--- a/tests/test_ml_suggestions.py
+++ b/tests/test_ml_suggestions.py
@@ -1,149 +1,133 @@
+"""Suggestion read-path (tagging-v2): suggestions come from trained HEADS, not
+Camie predictions or centroids. Heads are inserted directly (training needs
+scikit-learn, ml image only); scoring is numpy-only (available via pgvector)."""
import pytest
+from sqlalchemy import select
-from backend.app.models import ImageRecord, TagKind
+from backend.app.models import ImageRecord, MLSettings, TagHead, TagKind
from backend.app.models.tag import image_tag
-from backend.app.services.ml.aliases import AliasService
+from backend.app.services.ml.allowlist import AllowlistService
from backend.app.services.ml.suggestions import SuggestionService
from backend.app.services.tag_service import TagService
pytestmark = pytest.mark.integration
-def _img(sha: str) -> ImageRecord:
- return ImageRecord(
- path=f"/images/{sha}.jpg",
- sha256=sha,
- size_bytes=1,
- mime="image/jpeg",
- width=1,
- height=1,
- origin="imported_filesystem",
- integrity_status="unknown",
+def _emb(slot: int, val: float = 3.0) -> list[float]:
+ """An embedding pointing along axis `slot` (so its L2-normalized form is the
+ unit vector e_slot — a head with weights e_slot scores it sigmoid(1)≈0.73)."""
+ v = [0.0] * 1152
+ v[slot] = val
+ return v
+
+
+async def _img(db, sha: str, emb=None) -> ImageRecord:
+ img = 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,
)
-
-
-async def _seed_img(db, sha: str, predictions: dict) -> ImageRecord:
- """#768: create an image + seed its predictions into image_prediction
- (the read path's source), returning the flushed record."""
- from tests._prediction_helpers import seed_predictions
-
- img = _img(sha)
db.add(img)
await db.flush()
- await seed_predictions(db, img.id, predictions)
return img
+async def _embver(db) -> str:
+ s = (await db.execute(select(MLSettings).where(MLSettings.id == 1))).scalar_one()
+ return s.embedder_model_version
+
+
+async def _head(db, tag_id: int, slot: int, suggest_threshold: float = 0.5):
+ weights = [0.0] * 1152
+ weights[slot] = 1.0
+ db.add(TagHead(
+ tag_id=tag_id, embedding_version=await _embver(db),
+ weights=weights, bias=0.0, suggest_threshold=suggest_threshold,
+ auto_apply_threshold=None, n_pos=10, n_neg=30,
+ ap=0.8, precision_cv=0.9, recall=0.6,
+ ))
+
+
@pytest.mark.asyncio
-async def test_threshold_filters_low_confidence_general(db):
- # Default general threshold is 0.50 (alembic 0029 lowered it from
- # 0.95). Use 0.30/0.60 to keep the test asserting threshold behavior
- # rather than the exact cutoff number.
- img = await _seed_img(
- db,
- "a" * 64,
- {
- "lowconf": {"category": "general", "confidence": 0.30},
- "sword": {"category": "general", "confidence": 0.97},
- },
- )
+async def test_head_suggestion_surfaces_for_matching_image(db):
+ tag = await TagService(db).find_or_create("glasses", TagKind.general)
+ img = await _img(db, "a" * 64, _emb(0))
+ await _head(db, tag.id, slot=0)
+ await db.commit()
+
sl = await SuggestionService(db).for_image(img.id)
- names = [s.display_name for s in sl.by_category.get("general", [])]
- # display_name is normalized (tag_name.normalize) before surfacing.
- assert "Sword" in names
- assert "Lowconf" not in names
+ general = sl.by_category["general"]
+ assert len(general) == 1
+ s = general[0]
+ assert s.canonical_tag_id == tag.id
+ assert s.source == "head"
+ assert s.creates_new_tag is False
+ assert s.via_alias is False and s.raw_name is None
+ assert s.score > 0.5
@pytest.mark.asyncio
-async def test_threshold_override_surfaces_low_confidence(db):
- # The typed-dropdown "show everything the model saw" mode: threshold_override
- # surfaces stored predictions below the configured threshold (in canonical
- # formatting) so they can be picked instead of hand-typed (2026-06-09).
- img = await _seed_img(
- db,
- "d" * 64,
- {
- "lowconf": {"category": "general", "confidence": 0.30},
- "sword": {"category": "general", "confidence": 0.97},
- },
- )
- sl = await SuggestionService(db).for_image(img.id, threshold_override=0.0)
- names = [s.display_name for s in sl.by_category.get("general", [])]
- assert "Sword" in names
- assert "Lowconf" in names # below the configured threshold, surfaced anyway
-
- # Unsurfaced categories are still excluded even with the override.
- img2 = await _seed_img(
- db, "e" * 64, {"safe": {"category": "rating", "confidence": 0.99}}
- )
- sl2 = await SuggestionService(db).for_image(img2.id, threshold_override=0.0)
- assert "rating" not in sl2.by_category
+async def test_no_embedding_means_no_suggestions(db):
+ img = await _img(db, "b" * 64, None)
+ tag = await TagService(db).find_or_create("cat", TagKind.general)
+ await _head(db, tag.id, slot=0)
+ await db.commit()
+ assert (await SuggestionService(db).for_image(img.id)).by_category == {}
@pytest.mark.asyncio
-async def test_unsurfaced_category_dropped(db):
- img = await _seed_img(
- db,
- "b" * 64,
- {"safe": {"category": "rating", "confidence": 0.99}},
- )
- sl = await SuggestionService(db).for_image(img.id)
- assert "rating" not in sl.by_category
-
-
-@pytest.mark.asyncio
-async def test_alias_resolution(db):
- tags = TagService(db)
- canonical = await tags.find_or_create("Sasuke Uchiha", TagKind.character)
- await AliasService(db).create("uchiha_sasuke", "character", canonical.id)
- img = await _seed_img(
- db,
- "c" * 64,
- {"uchiha_sasuke": {"category": "character", "confidence": 0.96}},
- )
- sl = await SuggestionService(db).for_image(img.id)
- chars = sl.by_category["character"]
- assert len(chars) == 1
- assert chars[0].display_name == "Sasuke Uchiha"
- assert chars[0].canonical_tag_id == canonical.id
- assert chars[0].creates_new_tag is False
- # Surfaced via an alias on the raw model key — the UI marks it + offers undo.
- assert chars[0].via_alias is True
- assert chars[0].raw_name == "uchiha_sasuke"
-
-
-@pytest.mark.asyncio
-async def test_raw_tag_creates_new(db):
- img = await _seed_img(
- db,
- "d" * 64,
- {"brand_new_tag": {"category": "character", "confidence": 0.96}},
- )
- sl = await SuggestionService(db).for_image(img.id)
- chars = sl.by_category["character"]
- # display_name is the normalized Camie name (underscores -> spaces,
- # title-cased), not the raw vocab key.
- assert chars[0].display_name == "Brand New Tag"
- assert chars[0].creates_new_tag is True
- # Not aliased, but the raw key is carried so the modal can author one.
- assert chars[0].via_alias is False
- assert chars[0].raw_name == "brand_new_tag"
- assert chars[0].canonical_tag_id is None
+async def test_no_heads_means_no_suggestions(db):
+ img = await _img(db, "c" * 64, _emb(0))
+ await db.commit() # no heads trained yet
+ assert (await SuggestionService(db).for_image(img.id)).by_category == {}
@pytest.mark.asyncio
async def test_applied_tag_not_suggested(db):
- tags = TagService(db)
- tag = await tags.find_or_create("alreadyhere", TagKind.character)
- img = await _seed_img(
- db,
- "e" * 64,
- {"alreadyhere": {"category": "character", "confidence": 0.96}},
- )
+ tag = await TagService(db).find_or_create("dog", TagKind.general)
+ img = await _img(db, "d" * 64, _emb(0))
+ await _head(db, tag.id, slot=0)
await db.execute(
image_tag.insert().values(
image_record_id=img.id, tag_id=tag.id, source="manual"
)
)
+ await db.commit()
sl = await SuggestionService(db).for_image(img.id)
- assert "character" not in sl.by_category or not sl.by_category["character"]
+ assert "general" not in sl.by_category or not sl.by_category["general"]
+
+
+@pytest.mark.asyncio
+async def test_threshold_override_surfaces_below_cut(db):
+ # A head with a high suggest_threshold won't surface on a so-so score, but
+ # the dropdown's override=0 floor surfaces every head regardless.
+ tag = await TagService(db).find_or_create("horse", TagKind.general)
+ img = await _img(db, "e" * 64, _emb(1)) # orthogonal to the head → score 0.5
+ await _head(db, tag.id, slot=0, suggest_threshold=0.6)
+ await db.commit()
+ svc = SuggestionService(db)
+ assert svc and not (await svc.for_image(img.id)).by_category.get("general")
+ flooded = await svc.for_image(img.id, threshold_override=0.0)
+ assert any(s.canonical_tag_id == tag.id for s in flooded.by_category["general"])
+
+
+@pytest.mark.asyncio
+async def test_rejected_tag_surfaced_flagged_then_reversible(db):
+ # A dismissed suggestion is NOT dropped: it stays flagged rejected so the
+ # rail can show it + offer one-click un-reject (operator-asked 2026-06-27).
+ tag = await TagService(db).find_or_create("goblin", TagKind.general)
+ img = await _img(db, "f" * 64, _emb(0))
+ await _head(db, tag.id, slot=0)
+ await db.commit()
+ await AllowlistService(db).dismiss(img.id, tag.id)
+ await db.commit()
+
+ sl = await SuggestionService(db).for_image(img.id)
+ s = next(x for x in sl.by_category["general"] if x.canonical_tag_id == tag.id)
+ assert s.rejected is True
+
+ await AllowlistService(db).undismiss(img.id, tag.id)
+ await db.commit()
+ sl2 = await SuggestionService(db).for_image(img.id)
+ s2 = next(x for x in sl2.by_category["general"] if x.canonical_tag_id == tag.id)
+ assert s2.rejected is False
diff --git a/tests/test_suggestions_bulk.py b/tests/test_suggestions_bulk.py
index d0d101a..5df13f2 100644
--- a/tests/test_suggestions_bulk.py
+++ b/tests/test_suggestions_bulk.py
@@ -1,88 +1,84 @@
+"""Consensus (for_selection) over the tagging-v2 HEAD suggestion source."""
import pytest
+from sqlalchemy import select
from backend.app import create_app
-from backend.app.models import ImageRecord, TagKind
+from backend.app.models import ImageRecord, MLSettings, TagHead, TagKind
from backend.app.models.tag import image_tag
from backend.app.services.ml.suggestions import SuggestionService
from backend.app.services.tag_service import TagService
-from tests._prediction_helpers import seed_predictions
pytestmark = pytest.mark.integration
-def _img(sha: str) -> ImageRecord:
- return ImageRecord(
- path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1,
- mime="image/jpeg", width=1, height=1,
- origin="imported_filesystem", integrity_status="unknown",
+def _emb(slot: int) -> list[float]:
+ v = [0.0] * 1152
+ v[slot] = 3.0
+ return v
+
+
+async def _img(db, sha: str, emb=None) -> ImageRecord:
+ img = 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(img)
+ await db.flush()
+ return img
+
+
+async def _head(db, tag_id: int, slot: int = 0):
+ s = (await db.execute(select(MLSettings).where(MLSettings.id == 1))).scalar_one()
+ weights = [0.0] * 1152
+ weights[slot] = 1.0
+ db.add(TagHead(
+ tag_id=tag_id, embedding_version=s.embedder_model_version,
+ 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,
+ ))
@pytest.mark.asyncio
async def test_consensus_includes_tag_over_threshold(db):
- tags = TagService(db)
- t = await tags.find_or_create("sword", TagKind.general)
- a = _img("a" * 64)
- b = _img("b" * 64)
- db.add_all([a, b])
- await db.flush()
- await seed_predictions(db, a.id, {"sword": {"category": "general", "confidence": 0.97}})
- await seed_predictions(db, b.id, {"sword": {"category": "general", "confidence": 0.95}})
+ t = await TagService(db).find_or_create("sword", TagKind.general)
+ a = await _img(db, "a" * 64, _emb(0))
+ b = await _img(db, "b" * 64, _emb(0))
+ await _head(db, t.id, slot=0)
+ await db.commit()
res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8)
- gen = res["general"]
- assert any(s["canonical_tag_id"] == t.id for s in gen)
- s = next(s for s in gen if s["canonical_tag_id"] == t.id)
- assert s["coverage"] == 1.0
- assert 0.95 <= s["confidence"] <= 0.97
+ s = next(s for s in res["general"] if s["canonical_tag_id"] == t.id)
+ assert s["coverage"] == 1.0 # suggested on both
+ assert s["confidence"] > 0.5
@pytest.mark.asyncio
async def test_consensus_counts_already_applied_for_coverage(db):
- tags = TagService(db)
- t = await tags.find_or_create("sky", TagKind.general)
- a = _img("c" * 64)
- b = _img("d" * 64) # no prediction
- db.add_all([a, b])
- await db.flush()
- await seed_predictions(db, a.id, {"sky": {"category": "general", "confidence": 0.96}})
- # b already has the tag applied -> counts toward coverage, not confidence
+ t = await TagService(db).find_or_create("sky", TagKind.general)
+ a = await _img(db, "c" * 64, _emb(0)) # head suggests it
+ b = await _img(db, "d" * 64, None) # no embedding; tag applied instead
+ await _head(db, t.id, slot=0)
await db.execute(
image_tag.insert().values(
image_record_id=b.id, tag_id=t.id, source="manual"
)
)
+ await db.commit()
res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8)
s = next(s for s in res["general"] if s["canonical_tag_id"] == t.id)
assert s["coverage"] == 1.0 # 1 suggested + 1 applied / 2
- assert s["confidence"] == pytest.approx(0.96, abs=1e-4)
@pytest.mark.asyncio
async def test_consensus_excludes_below_threshold(db):
- tags = TagService(db)
- await tags.find_or_create("rare", TagKind.general)
- a = _img("e" * 64)
- b = _img("f" * 64)
- db.add_all([a, b])
- await db.flush()
- await seed_predictions(db, a.id, {"rare": {"category": "general", "confidence": 0.96}})
+ t = await TagService(db).find_or_create("rare", TagKind.general)
+ a = await _img(db, "e" * 64, _emb(0)) # suggested here
+ b = await _img(db, "f" * 64, None) # not here → coverage 0.5 < 0.8
+ await _head(db, t.id, slot=0)
+ await db.commit()
res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8)
- assert all(
- s["name"] != "rare" for s in res.get("general", [])
- ) # coverage 0.5 < 0.8
-
-
-@pytest.mark.asyncio
-async def test_consensus_skips_creates_new_tag(db):
- a = _img("g" * 64)
- b = _img("h" * 64)
- db.add_all([a, b])
- await db.flush()
- await seed_predictions(db, a.id, {"neverseen": {"category": "general", "confidence": 0.99}})
- await seed_predictions(db, b.id, {"neverseen": {"category": "general", "confidence": 0.99}})
- res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8)
- # 'neverseen' has no Tag row -> creates_new_tag -> excluded from consensus
- assert all(s["name"] != "neverseen" for s in res.get("general", []))
+ assert all(s["name"] != "rare" for s in res.get("general", []))
@pytest.mark.asyncio
@@ -93,13 +89,9 @@ async def test_consensus_threshold_clamped_and_empty_for_no_ids(db):
@pytest.mark.asyncio
async def test_bulk_suggestions_route(db):
-
- tags = TagService(db)
- await tags.find_or_create("sword", TagKind.general)
- a = _img("i" * 64)
- db.add(a)
- await db.commit()
- await seed_predictions(db, a.id, {"sword": {"category": "general", "confidence": 0.97}})
+ t = await TagService(db).find_or_create("sword", TagKind.general)
+ a = await _img(db, "i" * 64, _emb(0))
+ await _head(db, t.id, slot=0)
await db.commit()
app = create_app()
async with app.test_client() as c:
@@ -115,7 +107,6 @@ async def test_bulk_suggestions_route(db):
@pytest.mark.asyncio
async def test_bulk_suggestions_requires_ids(db):
-
app = create_app()
async with app.test_client() as c:
resp = await c.post("/api/suggestions/bulk", json={})