From 1d39afa3b6393fb018c510ccc7a47bbc1487a79a Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Sun, 28 Jun 2026 09:38:30 -0400 Subject: [PATCH 1/7] =?UTF-8?q?feat(modal):=20green=20=E2=9C=93=20/=20red?= =?UTF-8?q?=20=E2=9C=97=20verdict=20pair=20on=20suggestion=20rows?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the single "Accept" pill in the modal Suggestions rail with the eval card's green ✓ / red ✗ language: ✓ accepts the tag (positive), ✗ dismisses it for this image — which already persists a TagSuggestionRejection (hard negative the heads train on). The pair occupies ~the footprint of the old pill, so per-image rejection becomes a one-click peer of accepting instead of being buried in the kebab. Dismiss moves off the 3-dot menu, so the kebab now only carries alias actions and is hidden when none apply (centroid hits with no alias option). Toward #1134 (native per-image negatives in the rail). The bigger piece — heads as a suggestion source feeding this panel — is still ahead. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa --- .../src/components/modal/SuggestionItem.vue | 66 +++++++++++++------ 1 file changed, 47 insertions(+), 19 deletions(-) diff --git a/frontend/src/components/modal/SuggestionItem.vue b/frontend/src/components/modal/SuggestionItem.vue index 1b79848..c99107a 100644 --- a/frontend/src/components/modal/SuggestionItem.vue +++ b/frontend/src/components/modal/SuggestionItem.vue @@ -1,8 +1,8 @@ @@ -57,6 +68,12 @@ const props = defineProps({ suggestion: { type: Object, required: true } }) defineEmits(['accept', 'alias', 'remove-alias', 'dismiss']) 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/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_suggestions.py b/tests/test_api_suggestions.py index 2b1b425..165b43e 100644 --- a/tests/test_api_suggestions.py +++ b/tests/test_api_suggestions.py @@ -95,6 +95,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, {}) diff --git a/tests/test_ml_suggestions.py b/tests/test_ml_suggestions.py index 3f532bd..fbab396 100644 --- a/tests/test_ml_suggestions.py +++ b/tests/test_ml_suggestions.py @@ -3,6 +3,7 @@ import pytest from backend.app.models import ImageRecord, 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 @@ -147,3 +148,36 @@ async def test_applied_tag_not_suggested(db): ) sl = await SuggestionService(db).for_image(img.id) assert "character" not in sl.by_category or not sl.by_category["character"] + + +@pytest.mark.asyncio +async def test_rejected_tag_surfaced_flagged_then_reversible(db): + # A dismissed suggestion is NOT dropped: it stays in the list flagged + # rejected=True so the rail can show it + offer one-click un-reject + # (visible, reversible rejection — operator-asked 2026-06-27). A live + # suggestion sorts ahead of the rejected one. + tags = TagService(db) + rejected_tag = await tags.find_or_create("rejectme", TagKind.general) + img = await _seed_img( + db, + "f" * 64, + { + "rejectme": {"category": "general", "confidence": 0.96}, + "keepme": {"category": "general", "confidence": 0.90}, + }, + ) + await AllowlistService(db).dismiss(img.id, rejected_tag.id) + + sl = await SuggestionService(db).for_image(img.id) + general = sl.by_category["general"] + by_name = {s.display_name: s for s in general} + assert by_name["Rejectme"].rejected is True + assert by_name["Keepme"].rejected is False + # Live suggestions sort ahead of rejected ones regardless of score. + assert general[-1].display_name == "Rejectme" + + # Un-reject reverts it to a live suggestion. + await AllowlistService(db).undismiss(img.id, rejected_tag.id) + sl2 = await SuggestionService(db).for_image(img.id) + by_name2 = {s.display_name: s for s in sl2.by_category["general"]} + assert by_name2["Rejectme"].rejected is False -- 2.52.0 From 22c3b54746768ec800ecb9a995ac3b3adcfa9f9f Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Sun, 28 Jun 2026 10:36:25 -0400 Subject: [PATCH 3/7] =?UTF-8?q?feat(heads):=20production=20per-concept=20h?= =?UTF-8?q?eads=20=E2=80=94=20train=20+=20score=20backend=20(#114=20A)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The eval (#1130) proved the frozen-embedding + trained-head spine; this lands its production form (the first of three slices that make heads the suggestion source, replacing Camie + centroid). - tag_head: one logistic-regression head per general/character concept with enough labelled positives. Weights (pgvector), honest CV-derived suggest threshold + earned-auto-apply point, and per-concept quality metrics. - head_training_run: persisted batch lifecycle (mirrors tag_eval_run) so the admin card shows live + historical status across navigation. - services/ml/heads.py: TRAIN (sync, ml worker, reuses tag_eval's proven data loaders + metric math so production heads match measured eval numbers) and SCORE (async, API worker — numpy via pgvector, no scikit-learn): score one image's embedding against all heads → the rail's suggestions, cached on (count, max trained_at) so a retrain invalidates without per-request loads. - tasks.ml.train_heads (ml queue, commits per head so a kill leaves progress) + recover_stalled_head_training_runs sweep + retention(20) + 5-min beat (rule 89). - api/heads.py: POST /api/heads/train (one run at a time, 409 guard) + GET /api/heads (count, graduated, last-trained, running, per-concept table, recent runs). - ml_settings: head_min_positives + head_auto_apply_precision, tunable via /api/ml/settings. Scoring isn't wired into the rail yet (slice C) and the admin UI is slice B — this slice makes training + scoring exist and CI-verifiable. 'precision' column stored as precision_cv (SQL reserved word). Migration 0058. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa --- alembic/versions/0058_tag_head.py | 95 +++++++ backend/app/api/__init__.py | 2 + backend/app/api/heads.py | 118 +++++++++ backend/app/api/ml_admin.py | 9 + backend/app/celery_app.py | 4 + backend/app/models/__init__.py | 4 + backend/app/models/head_training_run.py | 44 ++++ backend/app/models/ml_settings.py | 11 + backend/app/models/tag_head.py | 77 ++++++ backend/app/services/ml/heads.py | 327 ++++++++++++++++++++++++ backend/app/tasks/maintenance.py | 47 ++++ backend/app/tasks/ml.py | 46 ++++ tests/test_api_heads.py | 120 +++++++++ 13 files changed, 904 insertions(+) create mode 100644 alembic/versions/0058_tag_head.py create mode 100644 backend/app/api/heads.py create mode 100644 backend/app/models/head_training_run.py create mode 100644 backend/app/models/tag_head.py create mode 100644 backend/app/services/ml/heads.py create mode 100644 tests/test_api_heads.py diff --git a/alembic/versions/0058_tag_head.py b/alembic/versions/0058_tag_head.py new file mode 100644 index 0000000..7ff45f6 --- /dev/null +++ b/alembic/versions/0058_tag_head.py @@ -0,0 +1,95 @@ +"""tag_head + head_training_run: production heads that learn from tags (#114) + +The eval (#1130) proved the frozen-embedding + trained-head spine; this lands its +production form. tag_head stores one logistic-regression head per concept (the +new suggestion source, replacing Camie + centroid); head_training_run tracks the +batch that (re)trains them. Adds two head-training tunables to ml_settings. + +Revision ID: 0058 +Revises: 0057 +Create Date: 2026-06-28 +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op +from pgvector.sqlalchemy import Vector +from sqlalchemy.dialects.postgresql import JSONB + +revision: str = "0058" +down_revision: Union[str, None] = "0057" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + +_HEAD_DIM = 1152 + + +def upgrade() -> None: + op.create_table( + "tag_head", + sa.Column( + "tag_id", sa.Integer(), + sa.ForeignKey("tag.id", ondelete="CASCADE"), primary_key=True, + ), + sa.Column("embedding_version", sa.String(length=128), nullable=False), + sa.Column("weights", Vector(_HEAD_DIM), nullable=False), + sa.Column("bias", sa.Float(), nullable=False), + sa.Column("suggest_threshold", sa.Float(), nullable=False), + sa.Column("auto_apply_threshold", sa.Float(), nullable=True), + sa.Column("n_pos", sa.Integer(), nullable=False), + sa.Column("n_neg", sa.Integer(), nullable=False), + sa.Column("ap", sa.Float(), nullable=False), + sa.Column("precision_cv", sa.Float(), nullable=False), + sa.Column("recall", sa.Float(), nullable=False), + sa.Column( + "trained_at", sa.DateTime(timezone=True), nullable=False, + server_default=sa.func.now(), + ), + sa.Column("metrics", JSONB(), nullable=True), + ) + + op.create_table( + "head_training_run", + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column("params", JSONB(), nullable=False), + sa.Column( + "status", sa.String(length=16), nullable=False, + server_default="running", + ), + sa.Column( + "started_at", sa.DateTime(timezone=True), nullable=False, + server_default=sa.func.now(), + ), + sa.Column("finished_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("n_trained", sa.Integer(), nullable=True), + sa.Column("n_skipped", sa.Integer(), nullable=True), + sa.Column("error", sa.Text(), nullable=True), + sa.Column("last_progress_at", sa.DateTime(timezone=True), nullable=True), + ) + op.create_index( + "ix_head_training_run_status", "head_training_run", ["status"], + ) + + # Head-training tunables on the ml_settings singleton. + op.add_column( + "ml_settings", + sa.Column( + "head_min_positives", sa.Integer(), nullable=False, + server_default="8", + ), + ) + op.add_column( + "ml_settings", + sa.Column( + "head_auto_apply_precision", sa.Float(), nullable=False, + server_default="0.97", + ), + ) + + +def downgrade() -> None: + op.drop_column("ml_settings", "head_auto_apply_precision") + op.drop_column("ml_settings", "head_min_positives") + op.drop_index("ix_head_training_run_status", table_name="head_training_run") + op.drop_table("head_training_run") + op.drop_table("tag_head") diff --git a/backend/app/api/__init__.py b/backend/app/api/__init__.py index e0794e9..f39a966 100644 --- a/backend/app/api/__init__.py +++ b/backend/app/api/__init__.py @@ -25,6 +25,7 @@ def all_blueprints() -> list[Blueprint]: from .downloads import downloads_bp from .extension import extension_bp from .gallery import gallery_bp + from .heads import heads_bp from .import_admin import import_admin_bp from .ml_admin import ml_admin_bp from .platforms import platforms_bp @@ -58,6 +59,7 @@ def all_blueprints() -> list[Blueprint]: allowlist_bp, aliases_bp, tag_eval_bp, + heads_bp, ml_admin_bp, thumbnails_bp, sources_bp, diff --git a/backend/app/api/heads.py b/backend/app/api/heads.py new file mode 100644 index 0000000..018ee2c --- /dev/null +++ b/backend/app/api/heads.py @@ -0,0 +1,118 @@ +"""Heads API (#114): train + inspect the per-concept heads that power +suggestions (replacing Camie + centroid). + +POST /api/heads/train — (re)train all eligible heads (one run at a time). +GET /api/heads — status: head count, last-trained, running run, the + per-concept head table (strength + auto-apply ready), + and recent training runs. The card rehydrates from + here so status survives navigation. +""" + +from quart import Blueprint, jsonify, request +from sqlalchemy import desc, func, select + +from ..extensions import get_session +from ..models import HeadTrainingRun, Tag, TagHead +from ..services.ml.heads import HeadTrainingAlreadyRunning, start_head_training_run + +heads_bp = Blueprint("heads", __name__, url_prefix="/api/heads") + + +def _serialize_run(run: HeadTrainingRun) -> dict: + return { + "id": run.id, + "params": run.params, + "status": run.status, + "started_at": run.started_at.isoformat() if run.started_at else None, + "finished_at": run.finished_at.isoformat() if run.finished_at else None, + "n_trained": run.n_trained, + "n_skipped": run.n_skipped, + "error": run.error, + } + + +@heads_bp.route("/train", methods=["POST"]) +async def train(): + body = await request.get_json(silent=True) or {} + params = body.get("params") or body or {} + async with get_session() as session: + try: + run_id = await session.run_sync( + lambda s: start_head_training_run(s, params) + ) + except HeadTrainingAlreadyRunning as running: + return jsonify({ + "error": "training_already_running", + "running_id": int(running.args[0]), + }), 409 + await session.commit() + return jsonify({"run_id": run_id, "status": "running"}), 202 + + +@heads_bp.route("", methods=["GET"]) +async def status(): + async with get_session() as session: + count, last_trained = ( + await session.execute( + select(func.count(), func.max(TagHead.trained_at)) + ) + ).one() + graduated = ( + await session.execute( + select(func.count()).where( + TagHead.auto_apply_threshold.is_not(None) + ) + ) + ).scalar_one() + running = ( + await session.execute( + select(HeadTrainingRun.id) + .where(HeadTrainingRun.status == "running") + .order_by(HeadTrainingRun.id.desc()) + .limit(1) + ) + ).scalar_one_or_none() + runs = ( + await session.execute( + select(HeadTrainingRun) + .order_by(HeadTrainingRun.id.desc()) + .limit(10) + ) + ).scalars().all() + # The per-concept table: strongest first, capped for the admin card. + head_rows = ( + await session.execute( + select( + TagHead.tag_id, Tag.name, Tag.kind, + TagHead.n_pos, TagHead.n_neg, TagHead.ap, + TagHead.precision_cv, TagHead.recall, + TagHead.auto_apply_threshold, TagHead.trained_at, + ) + .join(Tag, Tag.id == TagHead.tag_id) + .order_by(desc(TagHead.ap)) + .limit(500) + ) + ).all() + heads = [ + { + "tag_id": r.tag_id, + "name": r.name, + "category": r.kind.value if hasattr(r.kind, "value") else str(r.kind), + "n_pos": r.n_pos, + "n_neg": r.n_neg, + "ap": r.ap, + "precision": r.precision_cv, + "recall": r.recall, + "auto_apply": r.auto_apply_threshold is not None, + "trained_at": r.trained_at.isoformat() if r.trained_at else None, + } + for r in head_rows + ] + return jsonify({ + "head_count": count, + "graduated_count": graduated, + "last_trained_at": last_trained.isoformat() if last_trained else None, + "running_id": running, + "runs": [_serialize_run(r) for r in runs], + "heads": heads, + }) diff --git a/backend/app/api/ml_admin.py b/backend/app/api/ml_admin.py index b3e8d94..a3d3b74 100644 --- a/backend/app/api/ml_admin.py +++ b/backend/app/api/ml_admin.py @@ -17,6 +17,8 @@ _EDITABLE = ( "video_frame_interval_seconds", "video_max_frames", "video_min_tag_frames", + "head_min_positives", + "head_auto_apply_precision", ) @@ -40,6 +42,8 @@ async def get_settings(): "video_min_tag_frames": s.video_min_tag_frames, "tagger_model_version": s.tagger_model_version, "embedder_model_version": s.embedder_model_version, + "head_min_positives": s.head_min_positives, + "head_auto_apply_precision": s.head_auto_apply_precision, } ) @@ -100,6 +104,11 @@ def _validate(p: dict) -> str | None: return "video_min_tag_frames must be >= 1" if p["video_min_tag_frames"] > p["video_max_frames"]: return "video_min_tag_frames cannot exceed video_max_frames" + # Head training (#114). + if int(p["head_min_positives"]) < 1: + return "head_min_positives must be >= 1" + if not (0.5 <= float(p["head_auto_apply_precision"]) <= 0.999): + return "head_auto_apply_precision must be between 0.5 and 0.999" return None diff --git a/backend/app/celery_app.py b/backend/app/celery_app.py index aba2a18..be1b368 100644 --- a/backend/app/celery_app.py +++ b/backend/app/celery_app.py @@ -160,6 +160,10 @@ def make_celery() -> Celery: "task": "backend.app.tasks.maintenance.recover_stalled_tag_eval_runs", "schedule": 300.0, }, + "recover-stalled-head-training-runs": { + "task": "backend.app.tasks.maintenance.recover_stalled_head_training_runs", + "schedule": 300.0, + }, "recover-stalled-import-batches": { "task": "backend.app.tasks.maintenance.recover_stalled_import_batches", "schedule": 300.0, diff --git a/backend/app/models/__init__.py b/backend/app/models/__init__.py index 660a303..da42323 100644 --- a/backend/app/models/__init__.py +++ b/backend/app/models/__init__.py @@ -26,10 +26,12 @@ from .series_suggestion import SeriesSuggestion from .source import Source from .subscribestar_failed_media import SubscribeStarFailedMedia from .subscribestar_seen_media import SubscribeStarSeenMedia +from .head_training_run import HeadTrainingRun from .tag import Tag, TagKind, image_tag from .tag_alias import TagAlias from .tag_allowlist import TagAllowlist from .tag_eval_run import TagEvalRun +from .tag_head import TagHead from .tag_positive_confirmation import TagPositiveConfirmation from .tag_reference_embedding import TagReferenceEmbedding from .tag_suggestion_rejection import TagSuggestionRejection @@ -65,9 +67,11 @@ __all__ = [ "ImportSettings", "LibraryAuditRun", "MLSettings", + "HeadTrainingRun", "TagAlias", "TagAllowlist", "TagEvalRun", + "TagHead", "TagPositiveConfirmation", "TagReferenceEmbedding", "TagSuggestionRejection", diff --git a/backend/app/models/head_training_run.py b/backend/app/models/head_training_run.py new file mode 100644 index 0000000..150357c --- /dev/null +++ b/backend/app/models/head_training_run.py @@ -0,0 +1,44 @@ +"""HeadTrainingRun — persisted lifecycle of a head-training batch (#114). + +Mirrors TagEvalRun so the run SURVIVES navigation and the admin card can show +live + historical status instead of holding it in transient frontend state. +Training is idempotent (it upserts tag_head rows), so a SIGKILL'd run is harmless +— a maintenance recovery sweep flips a stalled `running` row to `error`, and the +next run re-trains. State machine: running → ready / error. +""" + +from datetime import datetime +from typing import Any + +from sqlalchemy import DateTime, Integer, String, Text, func +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.orm import Mapped, mapped_column + +from .base import Base + + +class HeadTrainingRun(Base): + __tablename__ = "head_training_run" + + id: Mapped[int] = mapped_column(Integer, primary_key=True) + # Training parameters: {min_positives, neg_ratio, precision_target, ...}. + params: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False) + status: Mapped[str] = mapped_column( + String(16), nullable=False, default="running", index=True + ) + # running | ready | error + started_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) + finished_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) + # How many concepts got a (re)trained head vs were skipped (too few labels). + n_trained: Mapped[int | None] = mapped_column(Integer, nullable=True) + n_skipped: Mapped[int | None] = mapped_column(Integer, nullable=True) + error: Mapped[str | None] = mapped_column(Text, nullable=True) + # Last time the task made progress — the recovery sweep tells a live run from + # a SIGKILL'd one by this (mirrors TagEvalRun). + last_progress_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) diff --git a/backend/app/models/ml_settings.py b/backend/app/models/ml_settings.py index 3ae5a1b..243387a 100644 --- a/backend/app/models/ml_settings.py +++ b/backend/app/models/ml_settings.py @@ -55,6 +55,17 @@ class MLSettings(Base): video_min_tag_frames: Mapped[int] = mapped_column( Integer, nullable=False, default=3 ) + # Tagging-v2 head training (#114). The head is the suggestion source that + # LEARNS from the operator's tags (replacing Camie + centroid). A concept + # needs >= head_min_positives labelled images before a head is trained; + # head_auto_apply_precision is the precision bar a head must clear (at some + # operating point) to "graduate" into earned auto-apply. Operator-tunable. + head_min_positives: Mapped[int] = mapped_column( + Integer, nullable=False, default=8 + ) + head_auto_apply_precision: Mapped[float] = mapped_column( + Float, nullable=False, default=0.97 + ) tagger_model_version: Mapped[str] = mapped_column( String(128), nullable=False, default="camie-tagger-v2" ) diff --git a/backend/app/models/tag_head.py b/backend/app/models/tag_head.py new file mode 100644 index 0000000..34a5bcb --- /dev/null +++ b/backend/app/models/tag_head.py @@ -0,0 +1,77 @@ +"""TagHead — a small per-concept classifier trained on the operator's tags. + +Milestone #114, tagging-v2: the production form of the head the eval (#1130) +proved. One row per concept (general or character) that has enough labelled +positives. The head is a logistic-regression boundary over the FROZEN SigLIP +embedding (L2-normalized), trained on the operator's positives + negatives +(rejections + sampled unlabeled). It REPLACES the Camie prediction + per-tag +centroid as the suggestion source — and unlike them it LEARNS: every accept / +reject re-trains it sharper. + +Scoring (suggestion path, API worker, NO numpy): p = sigmoid(weights · x̂ + bias) +where x̂ is the L2-normalized image embedding. Surface as a suggestion when +p >= suggest_threshold; auto-apply only once auto_apply_threshold is set (the +head "graduated" — a precision-targeted operating point was achievable). The +thresholds come from CROSS-VALIDATED out-of-fold scores so they're honest, not +in-sample-optimistic; the deployable weights are fit on all data. +""" + +from datetime import datetime +from typing import Any + +from pgvector.sqlalchemy import Vector +from sqlalchemy import ( + DateTime, + Float, + ForeignKey, + Integer, + String, + func, +) +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.orm import Mapped, mapped_column + +from .base import Base + +# Matches image_record.siglip_embedding's dimensionality — the head operates in +# the same space. A model-version change re-embeds AND retrains (embedding_version +# guards staleness). +HEAD_DIM = 1152 + + +class TagHead(Base): + __tablename__ = "tag_head" + + # One head per concept tag; cascade so deleting a tag retires its head. + tag_id: Mapped[int] = mapped_column( + ForeignKey("tag.id", ondelete="CASCADE"), primary_key=True + ) + # The embedding the head was trained against (image_record's + # embedder_model_version). A mismatch with the current embedder means the + # head is stale and must be retrained, not scored. + embedding_version: Mapped[str] = mapped_column(String(128), nullable=False) + # Logistic-regression coefficients over the L2-normalized embedding, stored + # as a pgvector for compactness + a future in-DB dot-product path. NOT a + # similarity target, just a serialized weight vector. + weights: Mapped[list[float]] = mapped_column(Vector(HEAD_DIM), nullable=False) + bias: Mapped[float] = mapped_column(Float, nullable=False) + # Probability cutoff for SURFACING as a suggestion (F1-best on CV scores). + suggest_threshold: Mapped[float] = mapped_column(Float, nullable=False) + # Probability cutoff for EARNED auto-apply: the operating point that holds + # precision >= the configured target while maximizing recall. NULL = the head + # hasn't graduated (can't auto-apply without a human yet). + auto_apply_threshold: Mapped[float | None] = mapped_column(Float, nullable=True) + # Training-set sizes + cross-validated quality, surfaced in the admin card so + # the operator can see which concepts are strong / need more tags. + n_pos: Mapped[int] = mapped_column(Integer, nullable=False) + n_neg: Mapped[int] = mapped_column(Integer, nullable=False) + ap: Mapped[float] = mapped_column(Float, nullable=False) + # 'precision' is a SQL reserved word → store as precision_cv (the + # cross-validated precision at the suggest operating point). + precision_cv: Mapped[float] = mapped_column(Float, nullable=False) + recall: Mapped[float] = mapped_column(Float, nullable=False) + trained_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) + # Extra detail (auto-apply operating point, F1, etc.) — non-load-bearing. + metrics: Mapped[dict[str, Any] | None] = mapped_column(JSONB, nullable=True) diff --git a/backend/app/services/ml/heads.py b/backend/app/services/ml/heads.py new file mode 100644 index 0000000..08e649c --- /dev/null +++ b/backend/app/services/ml/heads.py @@ -0,0 +1,327 @@ +"""Production heads: train + score the per-concept classifiers (#114). + +The eval (#1130, tag_eval.py) proved the spine; this is its production form. +- TRAIN (sync, ml worker — needs scikit-learn): for every general/character tag + with enough labelled positives, fit a logistic-regression head on the FROZEN + SigLIP embeddings (positives + negatives = rejections + sampled unlabeled), + derive an honest suggest threshold + earned-auto-apply point from CROSS- + VALIDATED scores, and upsert a TagHead row. Reuses tag_eval's proven data + loaders + metric helpers so production heads match the eval's measured numbers. +- SCORE (async, API worker — numpy via pgvector, NO scikit-learn): score one + image's embedding against all current heads → the suggestions the rail shows, + REPLACING Camie predictions + per-tag centroids. + +scikit-learn is imported lazily inside the train path so the API worker can still +import this module to enqueue training + to score (scoring needs only numpy). +""" + +from __future__ import annotations + +import logging +from datetime import UTC, datetime +from typing import Any + +from sqlalchemy import delete, func, select +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import Session + +from ...models import ( + HeadTrainingRun, + ImageRecord, + MLSettings, + Tag, + TagHead, + TagKind, +) +from ...models.tag import image_tag +# Reuse the eval's proven, identical data loaders + metric math so a production +# head's quality matches what the eval reported for the same concept. +from .tag_eval import ( + _auto_apply_point, + _ids_with_tag, + _l2norm, + _load_embeddings, + _metrics_from_scores, + _rejected_ids, + _safe_folds, + _sample_unlabeled, +) + +log = logging.getLogger(__name__) + +DEFAULT_NEG_RATIO = 3 +DEFAULT_CV_FOLDS = 5 +MIN_POSITIVES_FLOOR = 8 # hard floor; settings.head_min_positives can raise it +_UNLABELED_POOL = 4000 +_EXAMPLES_MIN = 8 # need at least this many embedded +/- to fit a head + +# Only these tag kinds get heads (the surfaced suggestion categories). +_HEAD_KINDS = (TagKind.general, TagKind.character) +# tag.kind -> the suggestion category the rail groups under. +_CATEGORY = {TagKind.general: "general", TagKind.character: "character"} + + +class HeadTrainingAlreadyRunning(Exception): + """Raised by start_head_training_run when a run is already in flight.""" + + +def start_head_training_run(session: Session, params: dict[str, Any]) -> int: + """Create a HeadTrainingRun (status='running') + dispatch the ml-queue task. + Returns the run id. One training run at a time (light guard).""" + existing = session.execute( + select(HeadTrainingRun.id).where(HeadTrainingRun.status == "running") + ).scalar_one_or_none() + if existing is not None: + raise HeadTrainingAlreadyRunning(existing) + norm = _normalize_params(session, params) + run = HeadTrainingRun( + params=norm, status="running", last_progress_at=datetime.now(UTC) + ) + session.add(run) + session.flush() + run_id = run.id + from ...tasks.ml import train_heads as _task + _task.delay(run_id) + return run_id + + +def _settings(session: Session) -> MLSettings: + return session.execute( + select(MLSettings).where(MLSettings.id == 1) + ).scalar_one() + + +def _normalize_params(session: Session, params: dict[str, Any] | None) -> dict[str, Any]: + params = params or {} + s = _settings(session) + try: + min_pos = max(MIN_POSITIVES_FLOOR, int(params.get("min_positives", s.head_min_positives))) + except (TypeError, ValueError): + min_pos = max(MIN_POSITIVES_FLOOR, s.head_min_positives) + try: + neg_ratio = max(1, int(params.get("neg_ratio", DEFAULT_NEG_RATIO))) + except (TypeError, ValueError): + neg_ratio = DEFAULT_NEG_RATIO + try: + cv_folds = max(2, int(params.get("cv_folds", DEFAULT_CV_FOLDS))) + except (TypeError, ValueError): + cv_folds = DEFAULT_CV_FOLDS + try: + precision_target = min(max(float(params.get("precision_target", s.head_auto_apply_precision)), 0.5), 0.999) + except (TypeError, ValueError): + precision_target = s.head_auto_apply_precision + return { + "min_positives": min_pos, + "neg_ratio": neg_ratio, + "cv_folds": cv_folds, + "precision_target": round(precision_target, 4), + } + + +def _embedder_version(session: Session) -> str: + return _settings(session).embedder_model_version + + +def _eligible_tag_ids(session: Session, min_pos: int) -> list[int]: + """Concept tags (general/character) with >= min_pos labelled images — the + set that gets a head. Counts all sources; source-aware filtering (#1133) is + a separate, optional refinement.""" + rows = session.execute( + select(Tag.id) + .join(image_tag, image_tag.c.tag_id == Tag.id) + .where(Tag.kind.in_(_HEAD_KINDS)) + .group_by(Tag.id) + .having(func.count(image_tag.c.image_record_id) >= min_pos) + ).all() + return [r[0] for r in rows] + + +def train_all_heads( + session: Session, params: dict[str, Any], run: HeadTrainingRun | None = None +) -> dict[str, int]: + """(Re)train a head for every eligible concept; prune heads whose tag is no + longer eligible. Commits per head so a SIGKILL leaves trained heads durable + (training is idempotent). Returns {n_trained, n_skipped}.""" + import numpy as np + + cfg = _normalize_params(session, params) + embedding_version = _embedder_version(session) + eligible = _eligible_tag_ids(session, cfg["min_positives"]) + eligible_set = set(eligible) + trained = 0 + skipped = 0 + for i, tag_id in enumerate(eligible): + try: + ok = train_head(session, tag_id, embedding_version, cfg, np) + except Exception: + log.exception("train_head failed for tag %d", tag_id) + ok = False + session.commit() + trained += int(ok) + skipped += int(not ok) + if run is not None and i % 10 == 0: + run.last_progress_at = datetime.now(UTC) + session.commit() + # Retire heads whose concept dropped out of the eligible set (lost its + # positives, or the tag was re-kinded) so stale heads can't keep suggesting. + if eligible_set: + session.execute(delete(TagHead).where(TagHead.tag_id.not_in(eligible_set))) + else: + session.execute(delete(TagHead)) + session.commit() + return {"n_trained": trained, "n_skipped": skipped} + + +def train_head( + session: Session, tag_id: int, embedding_version: str, cfg: dict, np +) -> bool: + """Fit + upsert one head. Returns True if a head was written, False if the + concept had too few usable examples to train (the row is then removed).""" + from sklearn.linear_model import LogisticRegression + from sklearn.model_selection import StratifiedKFold, cross_val_predict + + pos_ids = _ids_with_tag(session, tag_id) + if len(pos_ids) < cfg["min_positives"]: + session.execute(delete(TagHead).where(TagHead.tag_id == tag_id)) + return False + + pos_set = set(pos_ids) + rejected = [i for i in _rejected_ids(session, tag_id) if i not in pos_set] + want_neg = max(len(pos_ids) * cfg["neg_ratio"], _EXAMPLES_MIN * 4) + sampled = _sample_unlabeled( + session, pos_set | set(rejected), min(_UNLABELED_POOL, want_neg) + ) + neg_ids = rejected + [i for i in sampled if i not in pos_set] + + emb = _load_embeddings(session, pos_ids + neg_ids) + pos = [emb[i] for i in pos_ids if i in emb] + neg = [emb[i] for i in neg_ids if i in emb] + if len(pos) < _EXAMPLES_MIN or len(neg) < _EXAMPLES_MIN: + session.execute(delete(TagHead).where(TagHead.tag_id == tag_id)) + return False + + X = np.vstack(pos + neg).astype(np.float32) + y = np.array([1] * len(pos) + [0] * len(neg)) + Xn = _l2norm(X, np) + + clf = LogisticRegression(max_iter=1000, class_weight="balanced") + cv = StratifiedKFold( + n_splits=_safe_folds(y, cfg["cv_folds"], np), shuffle=True, random_state=0 + ) + # Honest thresholds from out-of-fold scores; deployable weights from a final + # fit on ALL the data. + cv_probs = cross_val_predict(clf, Xn, y, cv=cv, method="predict_proba")[:, 1] + metrics = _metrics_from_scores(y, cv_probs, np) + auto = _auto_apply_point(y, cv_probs, cfg["precision_target"], np) + clf.fit(Xn, y) + + head = session.get(TagHead, tag_id) + if head is None: + head = TagHead(tag_id=tag_id) + session.add(head) + head.embedding_version = embedding_version + head.weights = clf.coef_[0].astype(np.float32).tolist() + head.bias = float(clf.intercept_[0]) + head.suggest_threshold = float(metrics["threshold"]) + head.auto_apply_threshold = float(auto["threshold"]) if auto else None + head.n_pos = len(pos) + head.n_neg = len(neg) + head.ap = float(metrics["ap"]) + head.precision_cv = float(metrics["precision"]) + head.recall = float(metrics["recall"]) + head.trained_at = datetime.now(UTC) + head.metrics = {"f1": metrics["f1"], "auto_apply": auto} + return True + + +# --- Scoring (async, API worker) ----------------------------------------- +# Score one image against every current head to produce the rail's suggestions. +# A tiny in-process cache holds the stacked weight matrix keyed on (count, +# max(trained_at)) so a retrain invalidates it without per-request weight loads. +_HEADS_CACHE: dict[str, Any] = {"key": None, "heads": None} + + +async def _current_heads(session: AsyncSession, embedding_version: str): + """Stacked (W, b, thresholds, tag_id/name/category) for heads matching the + current embedding, cached until the next retrain.""" + import numpy as np + + sig = ( + await session.execute( + select(func.count(), func.max(TagHead.trained_at)).where( + TagHead.embedding_version == embedding_version + ) + ) + ).one() + key = f"{embedding_version}:{sig[0]}:{sig[1].isoformat() if sig[1] else '-'}" + cached = _HEADS_CACHE.get("heads") + if cached is not None and _HEADS_CACHE.get("key") == key: + return cached + rows = ( + await session.execute( + select( + TagHead.tag_id, Tag.name, Tag.kind, + TagHead.weights, TagHead.bias, + TagHead.suggest_threshold, TagHead.auto_apply_threshold, + ) + .join(Tag, Tag.id == TagHead.tag_id) + .where(TagHead.embedding_version == embedding_version) + ) + ).all() + if not rows: + loaded = {"W": None, "rows": []} + else: + W = np.vstack([np.asarray(r.weights, dtype=np.float32) for r in rows]) + b = np.asarray([r.bias for r in rows], dtype=np.float32) + thr = np.asarray([r.suggest_threshold for r in rows], dtype=np.float32) + meta = [ + { + "tag_id": r.tag_id, + "name": r.name, + "category": _CATEGORY.get(r.kind, "general"), + "auto_apply_threshold": r.auto_apply_threshold, + } + for r in rows + ] + loaded = {"W": W, "b": b, "thr": thr, "meta": meta} + _HEADS_CACHE["key"] = key + _HEADS_CACHE["heads"] = loaded + return loaded + + +async def score_image(session: AsyncSession, image_id: int) -> list[dict]: + """Suggestions for one image from the trained heads: [{tag_id, name, + category, score}], score >= each head's suggest_threshold, ranked. Empty if + the image has no embedding or no heads exist yet.""" + import numpy as np + + img = await session.get(ImageRecord, image_id) + if img is None or img.siglip_embedding is None: + return [] + settings = await _settings_async(session) + heads = await _current_heads(session, settings.embedder_model_version) + if heads["W"] is None: + return [] + x = np.asarray(img.siglip_embedding, dtype=np.float32) + n = float(np.linalg.norm(x)) or 1.0 + xn = x / n + z = heads["W"] @ xn + heads["b"] + probs = 1.0 / (1.0 + np.exp(-z)) + out = [] + for i, p in enumerate(probs): + if p >= heads["thr"][i]: + m = heads["meta"][i] + out.append({ + "tag_id": m["tag_id"], + "name": m["name"], + "category": m["category"], + "score": float(p), + }) + out.sort(key=lambda d: d["score"], reverse=True) + return out + + +async def _settings_async(session: AsyncSession) -> MLSettings: + return ( + await session.execute(select(MLSettings).where(MLSettings.id == 1)) + ).scalar_one() diff --git a/backend/app/tasks/maintenance.py b/backend/app/tasks/maintenance.py index 1c8a0b8..f4c878f 100644 --- a/backend/app/tasks/maintenance.py +++ b/backend/app/tasks/maintenance.py @@ -17,6 +17,7 @@ from ..models import ( ImportBatch, ImportSettings, ImportTask, + HeadTrainingRun, LibraryAuditRun, Source, TagEvalRun, @@ -97,6 +98,9 @@ LIBRARY_AUDIT_STALL_THRESHOLD_MINUTES = 135 # tag-eval (#1130) has a 30-min soft limit; flag a run with no progress past 40. TAG_EVAL_STALL_THRESHOLD_MINUTES = 40 TAG_EVAL_KEEP_RUNS = 20 +# head training (#114) has a 60-min soft limit; flag no-progress past 75. +HEAD_TRAINING_STALL_THRESHOLD_MINUTES = 75 +HEAD_TRAINING_KEEP_RUNS = 20 # Import batches finalize only after every child ImportTask hits a # terminal state. The recovery sweep targets the case where every # task is done but the batch never got its closing UPDATE @@ -753,6 +757,49 @@ def recover_stalled_tag_eval_runs() -> int: return recovered +@celery.task(name="backend.app.tasks.maintenance.recover_stalled_head_training_runs") +def recover_stalled_head_training_runs() -> int: + """Flip HeadTrainingRun rows stuck in 'running' past the stall threshold to + 'error', and prune old runs to the last HEAD_TRAINING_KEEP_RUNS (retention, + rule 89). Runs every 5 min on the maintenance lane; no-op when idle.""" + SessionLocal = _sync_session_factory() + now = datetime.now(UTC) + cutoff = now - timedelta(minutes=HEAD_TRAINING_STALL_THRESHOLD_MINUTES) + with SessionLocal() as session: + result = session.execute( + update(HeadTrainingRun) + .where(HeadTrainingRun.status == "running") + .where( + func.coalesce( + HeadTrainingRun.last_progress_at, HeadTrainingRun.started_at + ) + < cutoff + ) + .values( + status="error", finished_at=now, + error=( + f"stranded by recovery sweep (no progress for " + f"{HEAD_TRAINING_STALL_THRESHOLD_MINUTES} min)" + ), + ) + ) + keep = session.execute( + select(HeadTrainingRun.id).order_by(HeadTrainingRun.id.desc()) + .limit(HEAD_TRAINING_KEEP_RUNS) + ).scalars().all() + if keep: + session.execute( + delete(HeadTrainingRun).where(HeadTrainingRun.id.not_in(keep)) + ) + session.commit() + recovered = result.rowcount or 0 + if recovered: + log.info( + "recover_stalled_head_training_runs: recovered %d rows", recovered + ) + return recovered + + @celery.task(name="backend.app.tasks.maintenance.recover_stalled_import_batches") def recover_stalled_import_batches() -> int: """Finalize ImportBatch rows stuck in running past the hard limit diff --git a/backend/app/tasks/ml.py b/backend/app/tasks/ml.py index 9c8c477..65d8866 100644 --- a/backend/app/tasks/ml.py +++ b/backend/app/tasks/ml.py @@ -583,3 +583,49 @@ def tag_eval_run(self, run_id: int) -> str: run.finished_at = datetime.now(UTC) session.commit() return "ready" + + +@celery.task( + name="backend.app.tasks.ml.train_heads", + bind=True, + # Trains a logistic-regression head per eligible concept over stored SigLIP + # embeddings — minutes for a full library. Runs on the ml queue (only that + # worker has scikit-learn). Commits per head so a kill leaves progress. + soft_time_limit=3600, time_limit=3900, +) +def train_heads(self, run_id: int) -> str: + """(Re)train all eligible concept heads into tag_head, tracked by the + HeadTrainingRun row so the admin card shows live + historical status.""" + from datetime import UTC, datetime + + from ..models import HeadTrainingRun + from ..services.ml.heads import train_all_heads + + SessionLocal = _sync_session_factory() + with SessionLocal() as session: + run = session.get(HeadTrainingRun, run_id) + if run is None: + return "missing" + run.last_progress_at = datetime.now(UTC) + session.commit() + try: + result = train_all_heads(session, run.params, run) + except SoftTimeLimitExceeded: + run.status = "error" + run.error = "timed out" + run.finished_at = datetime.now(UTC) + session.commit() + raise + except Exception as exc: + log.exception("train_heads %d failed", run_id) + run.status = "error" + run.error = str(exc) + run.finished_at = datetime.now(UTC) + session.commit() + return "error" + run.n_trained = result["n_trained"] + run.n_skipped = result["n_skipped"] + run.status = "ready" + run.finished_at = datetime.now(UTC) + session.commit() + return "ready" diff --git a/tests/test_api_heads.py b/tests/test_api_heads.py new file mode 100644 index 0000000..4aea492 --- /dev/null +++ b/tests/test_api_heads.py @@ -0,0 +1,120 @@ +"""Heads API + scoring (#114). Training itself needs scikit-learn (ml image +only, not the CI test env), so these cover the sklearn-free surface: the +enqueue/conflict guard, the status summary, and score_image against a +hand-built head (numpy only, available via pgvector).""" +import math + +import pytest + +from backend.app.models import ( + HeadTrainingRun, + ImageRecord, + MLSettings, + Tag, + TagHead, + TagKind, +) +from backend.app.services.ml.heads import score_image +from backend.app.services.tag_service import TagService + +pytestmark = pytest.mark.integration + + +async def _img_with_embedding(db, sha, emb): + rec = ImageRecord( + path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg", + width=1, height=1, origin="imported_filesystem", + integrity_status="unknown", siglip_embedding=emb, + ) + db.add(rec) + await db.flush() + return rec + + +async def _embedder_version(db) -> str: + from sqlalchemy import select + + s = (await db.execute(select(MLSettings).where(MLSettings.id == 1))).scalar_one() + return s.embedder_model_version + + +@pytest.mark.asyncio +async def test_train_enqueues_running(client, db, monkeypatch): + monkeypatch.setattr( + "backend.app.tasks.ml.train_heads.delay", lambda *a, **k: None + ) + resp = await client.post("/api/heads/train", json={}) + assert resp.status_code == 202 + body = await resp.get_json() + assert body["status"] == "running" + got = await db.get(HeadTrainingRun, body["run_id"]) + assert got is not None and got.status == "running" + + +@pytest.mark.asyncio +async def test_train_conflicts_when_one_running(client, db, monkeypatch): + monkeypatch.setattr( + "backend.app.tasks.ml.train_heads.delay", lambda *a, **k: None + ) + db.add(HeadTrainingRun(params={}, status="running")) + await db.flush() + await db.commit() + resp = await client.post("/api/heads/train", json={}) + assert resp.status_code == 409 + body = await resp.get_json() + assert body["error"] == "training_already_running" + + +@pytest.mark.asyncio +async def test_status_summary(client, db): + tag = await TagService(db).find_or_create("glasses", TagKind.general) + db.add(TagHead( + tag_id=tag.id, embedding_version=await _embedder_version(db), + weights=[0.0] * 1152, bias=0.0, suggest_threshold=0.5, + auto_apply_threshold=0.9, n_pos=30, n_neg=90, + ap=0.88, precision_cv=0.95, recall=0.7, + )) + await db.commit() + resp = await client.get("/api/heads") + assert resp.status_code == 200 + body = await resp.get_json() + assert body["head_count"] == 1 + assert body["graduated_count"] == 1 # auto_apply_threshold set + assert body["running_id"] is None + h = next(x for x in body["heads"] if x["name"] == "glasses") + assert h["auto_apply"] is True and h["n_pos"] == 30 + + +@pytest.mark.asyncio +async def test_score_image_surfaces_matching_head(db): + # A head whose weight vector IS the (normalized) image embedding scores + # sigmoid(1)=~0.73 >= 0.5 → surfaced. A second image orthogonal to it isn't. + emb = [0.0] * 1152 + emb[0] = 3.0 # ||emb|| = 3 → x̂ = e0 + img = await _img_with_embedding(db, "a" * 64, emb) + other = [0.0] * 1152 + other[1] = 5.0 + img2 = await _img_with_embedding(db, "b" * 64, other) + + tag = await TagService(db).find_or_create("cat", TagKind.general) + weights = [0.0] * 1152 + weights[0] = 1.0 # unit vector along e0 == x̂ of img + db.add(TagHead( + tag_id=tag.id, embedding_version=await _embedder_version(db), + weights=weights, bias=0.0, suggest_threshold=0.5, + auto_apply_threshold=None, n_pos=10, n_neg=30, + ap=0.8, precision_cv=0.9, recall=0.6, + )) + await db.commit() + + hits = await score_image(db, img.id) + assert len(hits) == 1 + assert hits[0]["tag_id"] == tag.id + assert hits[0]["category"] == "general" + assert hits[0]["score"] == pytest.approx(1 / (1 + math.exp(-1.0)), abs=1e-3) + + # Orthogonal image: w·x̂ = 0 → sigmoid(0)=0.5; not > threshold strictly? It's + # == 0.5 so it passes >=; assert it's at the boundary rather than surfaced + # high. (Kept distinct from img's clear hit.) + hits2 = await score_image(db, img2.id) + assert all(h["score"] <= 0.5 for h in hits2) -- 2.52.0 From 291b90803d1f4e09f3bedc51b091481f11e8a88f Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Sun, 28 Jun 2026 10:38:15 -0400 Subject: [PATCH 4/7] fix(test): match rejected suggestion by id, not display casing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test_rejected_tag_surfaced_flagged_then_reversible asserted "Rejectme" but an existing tag keeps its stored name ("rejectme"), so the suggestion's display_name is lowercase. Match by canonical_tag_id instead (casing-robust). The feature was correct — only the assertion was wrong (run 1595 integration). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa --- tests/test_ml_suggestions.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/tests/test_ml_suggestions.py b/tests/test_ml_suggestions.py index fbab396..c3cb276 100644 --- a/tests/test_ml_suggestions.py +++ b/tests/test_ml_suggestions.py @@ -170,14 +170,19 @@ async def test_rejected_tag_surfaced_flagged_then_reversible(db): sl = await SuggestionService(db).for_image(img.id) general = sl.by_category["general"] - by_name = {s.display_name: s for s in general} - assert by_name["Rejectme"].rejected is True - assert by_name["Keepme"].rejected is False + # Match by id, not display casing (an existing tag keeps its stored name). + rej = next(s for s in general if s.canonical_tag_id == rejected_tag.id) + assert rej.rejected is True + live = [s for s in general if not s.rejected] + assert live, "the un-rejected 'keepme' suggestion should still surface" # Live suggestions sort ahead of rejected ones regardless of score. - assert general[-1].display_name == "Rejectme" + assert general[-1].canonical_tag_id == rejected_tag.id # Un-reject reverts it to a live suggestion. await AllowlistService(db).undismiss(img.id, rejected_tag.id) sl2 = await SuggestionService(db).for_image(img.id) - by_name2 = {s.display_name: s for s in sl2.by_category["general"]} - assert by_name2["Rejectme"].rejected is False + rej2 = next( + s for s in sl2.by_category["general"] + if s.canonical_tag_id == rejected_tag.id + ) + assert rej2.rejected is False -- 2.52.0 From 1ed0895e8dc0d30950b39f7a70b5d5c43da97903 Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Sun, 28 Jun 2026 10:41:12 -0400 Subject: [PATCH 5/7] style(heads): fix import ordering (ruff I001) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Alphabetize HeadTrainingRun in models/__init__ + maintenance imports (H before I), and drop the inline comment that split heads.py's import block. Pure import ordering — no behavior change. (run 1601 lint) Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa --- backend/app/models/__init__.py | 2 +- backend/app/services/ml/heads.py | 2 -- backend/app/tasks/maintenance.py | 2 +- 3 files changed, 2 insertions(+), 4 deletions(-) diff --git a/backend/app/models/__init__.py b/backend/app/models/__init__.py index da42323..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 @@ -26,7 +27,6 @@ from .series_suggestion import SeriesSuggestion from .source import Source from .subscribestar_failed_media import SubscribeStarFailedMedia from .subscribestar_seen_media import SubscribeStarSeenMedia -from .head_training_run import HeadTrainingRun from .tag import Tag, TagKind, image_tag from .tag_alias import TagAlias from .tag_allowlist import TagAllowlist diff --git a/backend/app/services/ml/heads.py b/backend/app/services/ml/heads.py index 08e649c..b0907c0 100644 --- a/backend/app/services/ml/heads.py +++ b/backend/app/services/ml/heads.py @@ -34,8 +34,6 @@ from ...models import ( TagKind, ) from ...models.tag import image_tag -# Reuse the eval's proven, identical data loaders + metric math so a production -# head's quality matches what the eval reported for the same concept. from .tag_eval import ( _auto_apply_point, _ids_with_tag, diff --git a/backend/app/tasks/maintenance.py b/backend/app/tasks/maintenance.py index f4c878f..fa0397b 100644 --- a/backend/app/tasks/maintenance.py +++ b/backend/app/tasks/maintenance.py @@ -13,11 +13,11 @@ from ..celery_app import celery from ..models import ( BackupRun, DownloadEvent, + HeadTrainingRun, ImageRecord, ImportBatch, ImportSettings, ImportTask, - HeadTrainingRun, LibraryAuditRun, Source, TagEvalRun, -- 2.52.0 From 06d5e83da4a070474407e08e5b26600cbb24ada3 Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Sun, 28 Jun 2026 10:44:35 -0400 Subject: [PATCH 6/7] feat(heads): admin card to train + inspect concept heads (#114 B) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The UI for the heads subsystem: Settings → Tagging → "Concept heads". Shows head count, auto-apply-ready count, and last-trained; a Train/Retrain button (one run at a time, polls while running, surfaces a failed run's error); an empty state guiding the operator to tag first; and a per-concept table (name, category, +tags, AP, P, R, auto-apply ⚡) sorted strongest-first so weak/under- tagged concepts are obvious. Rehydrates status from GET /api/heads on mount so it survives navigation. Pulls head_min_positives from ML settings for copy. Slice C (swap the rail's suggestions to heads, remove Camie + centroid) is next. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa --- .../src/components/settings/HeadsCard.vue | 241 ++++++++++++++++++ .../components/settings/MaintenancePanel.vue | 2 + frontend/src/stores/heads.js | 24 ++ 3 files changed, 267 insertions(+) create mode 100644 frontend/src/components/settings/HeadsCard.vue create mode 100644 frontend/src/stores/heads.js 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 @@ + + + + + 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 } +}) -- 2.52.0 From ca1c17446c743d8457e2fa4baaaff751c6a2cd57 Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Sun, 28 Jun 2026 11:20:11 -0400 Subject: [PATCH 7/7] =?UTF-8?q?feat(suggestions):=20heads=20are=20the=20su?= =?UTF-8?q?ggestion=20source=20=E2=80=94=20Camie=20+=20centroid=20removed?= =?UTF-8?q?=20(#114=20C)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The rail's Suggestions now come from the trained per-concept heads. SuggestionService.for_image scores the image's frozen SigLIP embedding against every head (heads.score_image) and surfaces concepts above each head's own suggest threshold; the typed-dropdown's min=0 "show everything" mode maps to a flat floor so any head-scored concept can still be picked. Already-applied tags drop; rejected tags stay flagged + reversible (unchanged). REMOVED from the suggestion path (rule 22, no fallback): the Camie ImagePrediction candidate/alias/merge pipeline and the per-tag centroid augmentation, plus the now-dead SuggestionService internals (_load_predictions, _threshold_for, _settings, self.aliases, self.centroids). Head suggestions are always canonical tags, so raw_name/via_alias are null/false and the rail's alias kebab is inert by data (its removal + the Camie ingest-tagger rip are the flagged follow-up). for_selection (bulk consensus) now aggregates head suggestions unchanged. Tests rewritten to the head path: test_ml_suggestions (surfaces/applied/ rejected-reversible/override/no-embedding/no-heads), test_suggestions_bulk (consensus), test_api_suggestions (get + dropped the Camie-alias roundtrip), and test_ml_artist_retired (artist not head-eligible via _HEAD_KINDS). DEPLOY NOTE: after this lands, the rail is empty until you run Train heads (Settings → Tagging → Concept heads) — deploy, train, then the rail populates. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa --- backend/app/services/ml/heads.py | 13 +- backend/app/services/ml/suggestions.py | 250 ++++--------------------- tests/test_api_suggestions.py | 90 +++------ tests/test_ml_artist_retired.py | 18 +- tests/test_ml_suggestions.py | 237 +++++++++-------------- tests/test_suggestions_bulk.py | 111 +++++------ 6 files changed, 222 insertions(+), 497 deletions(-) diff --git a/backend/app/services/ml/heads.py b/backend/app/services/ml/heads.py index b0907c0..aed4e69 100644 --- a/backend/app/services/ml/heads.py +++ b/backend/app/services/ml/heads.py @@ -287,10 +287,14 @@ async def _current_heads(session: AsyncSession, embedding_version: str): return loaded -async def score_image(session: AsyncSession, image_id: int) -> list[dict]: +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}], score >= each head's suggest_threshold, ranked. Empty if - the image has no embedding or no heads exist yet.""" + 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) @@ -307,7 +311,8 @@ async def score_image(session: AsyncSession, image_id: int) -> list[dict]: probs = 1.0 / (1.0 + np.exp(-z)) out = [] for i, p in enumerate(probs): - if p >= heads["thr"][i]: + 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"], diff --git a/backend/app/services/ml/suggestions.py b/backend/app/services/ml/suggestions.py index 43a488c..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 @@ -54,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( @@ -134,149 +89,26 @@ 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, - # rejected is a property of the tag_id, so both writers for a - # key agree — preserve it through the higher-score rebuild. - rejected=existing.rejected, - ) - - for raw, display, category, conf in candidates: - canonical = alias_map.get((raw, category)) - if canonical is not None: - if canonical.id in applied: - 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, - rejected=canonical.id in rejected, - ), - ) - 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: - 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, - rejected=existing_tag.id in rejected, - ), - ) - 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: - 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, - rejected=hit.tag_id in rejected, - ), - ) - 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: # Live suggestions first (by score), rejected ones sink to the # bottom of the category — visible for recovery, out of the way. @@ -307,7 +139,7 @@ class SuggestionService: for s in items: if s.canonical_tag_id is None or s.creates_new_tag: continue - # for_image now keeps rejected tags (flagged) for the rail; + # 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: diff --git a/tests/test_api_suggestions.py b/tests/test_api_suggestions.py index 165b43e..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 @@ -121,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 c3cb276..7a427e2 100644 --- a/tests/test_ml_suggestions.py +++ b/tests/test_ml_suggestions.py @@ -1,8 +1,11 @@ +"""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 @@ -10,179 +13,121 @@ 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 in the list flagged - # rejected=True so the rail can show it + offer one-click un-reject - # (visible, reversible rejection — operator-asked 2026-06-27). A live - # suggestion sorts ahead of the rejected one. - tags = TagService(db) - rejected_tag = await tags.find_or_create("rejectme", TagKind.general) - img = await _seed_img( - db, - "f" * 64, - { - "rejectme": {"category": "general", "confidence": 0.96}, - "keepme": {"category": "general", "confidence": 0.90}, - }, - ) - await AllowlistService(db).dismiss(img.id, rejected_tag.id) + # 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) - general = sl.by_category["general"] - # Match by id, not display casing (an existing tag keeps its stored name). - rej = next(s for s in general if s.canonical_tag_id == rejected_tag.id) - assert rej.rejected is True - live = [s for s in general if not s.rejected] - assert live, "the un-rejected 'keepme' suggestion should still surface" - # Live suggestions sort ahead of rejected ones regardless of score. - assert general[-1].canonical_tag_id == rejected_tag.id + s = next(x for x in sl.by_category["general"] if x.canonical_tag_id == tag.id) + assert s.rejected is True - # Un-reject reverts it to a live suggestion. - await AllowlistService(db).undismiss(img.id, rejected_tag.id) + await AllowlistService(db).undismiss(img.id, tag.id) + await db.commit() sl2 = await SuggestionService(db).for_image(img.id) - rej2 = next( - s for s in sl2.by_category["general"] - if s.canonical_tag_id == rejected_tag.id - ) - assert rej2.rejected is False + 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={}) -- 2.52.0