feat(suggestions): heads are the suggestion source — Camie + centroid removed (#114 C)
CI / lint (push) Successful in 2s
CI / frontend-build (push) Successful in 17s
CI / backend-lint-and-test (push) Successful in 26s
CI / integration (push) Successful in 3m19s

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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
This commit is contained in:
2026-06-28 11:20:11 -04:00
parent 06d5e83da4
commit ca1c17446c
6 changed files with 222 additions and 497 deletions
+9 -4
View File
@@ -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"],
+41 -209
View File
@@ -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: