bvandeusen
22c3b54746
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
121 lines
4.2 KiB
Python
121 lines
4.2 KiB
Python
"""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)
|