agent/describe_retrieve.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "datasets",
#     "huggingface_hub",
#     "iconclass",
#     "openai",
#     "pillow",
#     "sentence-transformers",
#     "numpy",
#     "tqdm",
# ]
# ///
"""Describe -> Retrieve: decouple VISION from TAXONOMY for Iconclass classification.

THE HYPOTHESIS
--------------
A thorough prior session found the fine-tuned Qwen3.5-4B is *capability-limited*
on Iconclass: ~25% code recall, ~62-64% corrected H-F1, and that wall is invariant
across reward tuning AND full-label re-SFT. The model can SEE artworks but can't
reliably NAME the right Iconclass codes — the taxonomy is the bottleneck, not the
vision.

This script tests whether RETRIEVAL can supply the taxonomy knowledge the 4B lacks,
with NO fine-tuning at all:

  1. DESCRIBE  — a strong VLM (Qwen3.6-35B-A3B via the HF router) writes a rich
     iconographic description of the image AND extracts discrete concepts
     (figures, scene, attributes, symbols, actions). This plays to the VLM's
     strength: looking at a picture and saying what's in it.
  2. RETRIEVE  — map that text -> candidate Iconclass codes via SEMANTIC SEARCH
     over the Iconclass tree (search_iconclass / build_search_index from
     grpo-tools/iconclass_tools.py). The Iconclass lib has no text->code search;
     this index is the only thing that supplies the "which code is this concept"
     knowledge. We retrieve per-concept and over the whole description, pool by
     similarity, and keep the top candidates. K is tunable.
  3. CONFIRM   — (optional) a binary VLM judge ("does this artwork depict X?")
     filters the retrieved candidates to lift precision (same judge pattern as
     eval_corrected.py).
  4. SCORE     — hierarchical-F1 + code precision/recall vs the FULL clean GT,
     using eval_sft.py's metric logic verbatim.

WHY THIS IS A FAIR / CONTAMINATION-FREE COMPARISON
--------------------------------------------------
We evaluate on `davanstrien/iconclass-vlm-brillfull` split `test` (full clean
labels, contamination-safe held-out by filename hash). The fine-tuned 4B v3 scored
52.0% H-F1 on this same split — but that number is *contaminated-optimistic*
because v3 trained on these very images. describe->retrieve uses NO training, so
if it matches/beats ~52% it is clearly competitive, and we especially watch RECALL
(the 4B's ~25% wall) to see if retrieval breaks through it.

CONSTRAINTS
-----------
Router (free tier) + CPU only, no GPU. The 35B cold-starts on the router (~2 min);
warm up first (--probe). The semantic index build is ~75s one-time on CPU.

USAGE
-----
    # reachability + index sanity, ~free
    uv run describe_retrieve.py --probe

    # small run to pick K / config
    uv run describe_retrieve.py --num-samples 12 --k 8

    # full ~120-image eval, with judge confirm
    uv run describe_retrieve.py --num-samples 120 --k 8 --confirm

    # sweep several K in one go (no judge), pick the best
    uv run describe_retrieve.py --num-samples 60 --sweep-k 4,6,8,12,16

NOTE: if the project's internal uv index can't resolve deps, prepend
``--index-url https://pypi.org/simple/`` to ``uv run``.
"""

import argparse
import asyncio
import base64
import io
import json
import os
import re
import sys
import time
from dataclasses import dataclass, field

from datasets import load_dataset
from huggingface_hub import get_token
from iconclass import init
from openai import AsyncOpenAI

# Bring in the semantic search over the Iconclass tree. This is the component
# that supplies the taxonomy knowledge (text concept -> code) the 4B lacks.
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(
    0,
    os.path.abspath(
        os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "grpo-tools")
    ),
)
from iconclass_tools import build_search_index, search_iconclass  # noqa: E402

ic = init()

# ----------------------------------------------------------------------------- config
DATASET = "davanstrien/iconclass-vlm-brillfull"
SPLIT = "test"

ROUTER_BASE_URL = "https://router.huggingface.co/v1"
VLM_MODEL = "Qwen/Qwen3.6-35B-A3B"
VLM_EXTRA_BODY = {"chat_template_kwargs": {"enable_thinking": False}}

# 4B v3 reference on THIS split (contaminated-optimistic — trained on these images).
FT_4B_V3_BRILLFULL_HF1 = 0.520

# ---- DESCRIBE prompt. Asks for a rich iconographic reading AND structured concepts.
# The structured concepts become individual retrieval queries; the prose description
# becomes a whole-image query. We push the VLM toward iconography (story/meaning,
# attributes, symbols) because the 4B's error analysis showed "narrative blindness"
# was its worst failure mode — exactly what a strong VLM should be able to supply.
DESCRIBE_SYSTEM = (
    "You are an expert art historian and iconographer. You look at an artwork and "
    "describe its CONTENT precisely: who/what is depicted, what is happening, the "
    "setting, and any attributes, symbols, gestures, or objects that identify "
    "people, stories, or themes. You name biblical, mythological, historical, and "
    "allegorical subjects when the visual evidence supports them."
)

DESCRIBE_USER = (
    "Describe the CONTENT of this artwork for iconographic cataloguing. "
    "Return ONLY a JSON object with these keys:\n"
    '  "description": a 2-4 sentence rich description of what is depicted '
    "(figures, scene, action, setting, mood),\n"
    '  "figures": list of people/beings depicted (e.g. "Virgin Mary", "a kneeling '
    'man", "an angel", "a lion"),\n'
    '  "scene": short phrase naming the subject/story/theme if identifiable '
    '(e.g. "the Crucifixion", "Annunciation", "a portrait", "a still life with '
    'fruit", "a landscape"),\n'
    '  "attributes": list of identifying objects/symbols/attributes (e.g. "a '
    'crown", "a sword", "a halo", "an open book", "grapes"),\n'
    '  "actions": list of poses/gestures/actions (e.g. "kneeling", "head turned '
    'right", "blessing", "holding a child").\n'
    "Only include what is visually supported. No prose outside the JSON."
)

# ---- CONFIRM (binary judge) prompt — verbatim style from eval_corrected.py, which
# was tuned in the project's judge prompt-engineering session.
JUDGE_PROMPT_TEMPLATE = (
    "Iconclass is a classification system for the content of art images. "
    "Artworks are assigned codes for the concepts they depict — objects, "
    "people, poses, scenes, themes.\n\n"
    "Does this artwork depict: {concept} ({code})\n\n"
    "Answer with a single word: YES or NO."
)


# =====================================================================
# Metrics — inlined VERBATIM from eval_sft.py so numbers are directly
# comparable to the documented SFT / GRPO / 4B-v3 results.
# =====================================================================


def _is_valid_iconclass(code):
    try:
        base_code = code.split("(")[0].strip()
        _ = ic[base_code]
        return True
    except Exception:
        base_code = code.split("(")[0].strip()
        return _check_partial_validity(base_code)


def _check_partial_validity(code):
    if not code:
        return False
    for i in range(len(code), 0, -1):
        try:
            _ = ic[code[:i]]
            return True
        except Exception:
            continue
    return False


def _get_hierarchy_levels(code):
    """E.g. "25F23" -> ["2", "25", "25F", "25F2", "25F23"]."""
    if not code:
        return []
    levels = []
    if len(code) >= 1:
        levels.append(code[0])
    if len(code) >= 2:
        levels.append(code[:2])
    if len(code) >= 3:
        levels.append(code[:3])
    for i in range(4, len(code) + 1):
        levels.append(code[:i])
    return levels


def _hierarchical_match_score(pred, truth):
    pred_clean = pred.split("(")[0].strip() if isinstance(pred, str) else ""
    truth_clean = truth.split("(")[0].strip() if isinstance(truth, str) else ""

    if pred_clean == truth_clean:
        return 1.0

    pred_levels = _get_hierarchy_levels(pred_clean)
    truth_levels = _get_hierarchy_levels(truth_clean)

    common_levels = 0
    for p, t in zip(pred_levels, truth_levels):
        if p == t:
            common_levels += 1
        else:
            break

    if common_levels == 0:
        return 0.0

    credit_by_level = [0.2, 0.4, 0.6, 0.8, 0.9, 0.95, 1.0]
    if common_levels >= len(credit_by_level):
        score = 1.0
    else:
        score = credit_by_level[common_levels - 1]

    if len(pred_levels) != len(truth_levels):
        score *= 0.95

    return score


def _calculate_hierarchical_f1(predicted, ground_truth):
    if not predicted and not ground_truth:
        return 1.0
    if not predicted or not ground_truth:
        return 0.0

    precision_score = 0.0
    for pred in predicted:
        best_match = max(
            (_hierarchical_match_score(pred, gt) for gt in ground_truth),
            default=0.0,
        )
        precision_score += best_match
    precision = precision_score / len(predicted) if predicted else 0.0

    recall_score = 0.0
    for gt in ground_truth:
        best_match = max(
            (_hierarchical_match_score(pred, gt) for pred in predicted),
            default=0.0,
        )
        recall_score += best_match
    recall = recall_score / len(ground_truth) if ground_truth else 0.0

    if precision + recall > 0:
        return 2 * (precision * recall) / (precision + recall)
    return 0.0


def _hierarchical_precision_recall(predicted, ground_truth):
    """The hierarchical (partial-credit) precision & recall behind the F1.

    Reported alongside exact-set P/R because the 4B's documented "~25% recall"
    wall is the exact-set code recall; we report BOTH so the comparison is honest.
    """
    if not predicted or not ground_truth:
        return 0.0, 0.0
    precision_score = sum(
        max((_hierarchical_match_score(p, gt) for gt in ground_truth), default=0.0)
        for p in predicted
    )
    recall_score = sum(
        max((_hierarchical_match_score(p, gt) for p in predicted), default=0.0)
        for gt in ground_truth
    )
    precision = precision_score / len(predicted)
    recall = recall_score / len(ground_truth)
    return precision, recall


def _code_precision_recall(predicted, ground_truth):
    """Exact-set code precision/recall (matches eval_sft.py)."""
    pred_set, gt_set = set(predicted), set(ground_truth)
    precision = len(pred_set & gt_set) / len(pred_set) if pred_set else 0.0
    recall = len(pred_set & gt_set) / len(gt_set) if gt_set else 0.0
    return precision, recall


# =====================================================================
# Iconclass helpers
# =====================================================================


def _code_description(code):
    """Decoded Iconclass text for a code (strips any (KEY)/(+xx) suffix)."""
    try:
        base = code.split("(")[0].strip()
        node = ic[base]
        desc = node() if callable(node) else str(node)
        return desc or None
    except Exception:
        return None


# =====================================================================
# Image handling — verbatim from smoke_test_models.py / eval_corrected.py
# =====================================================================


def image_to_data_url(img, max_size=768):
    if max(img.size) > max_size:
        r = max_size / max(img.size)
        from PIL import Image

        img = img.resize((int(img.size[0] * r), int(img.size[1] * r)), Image.LANCZOS)
    if img.mode != "RGB":
        img = img.convert("RGB")
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=85)
    return f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"


def first_image(example):
    """brillfull rows carry a singular `image` and a list `images`; accept either."""
    img = example.get("image")
    if img is not None:
        return img
    imgs = example.get("images", [])
    if not imgs:
        return None
    return imgs[0] if isinstance(imgs, list) else imgs


def get_gt_codes(example):
    """brillfull rows carry full clean labels in `label` (list[str])."""
    lab = example.get("label")
    if isinstance(lab, list):
        return [c for c in lab if isinstance(c, str) and c.strip()]
    return []


# =====================================================================
# STEP 1 — DESCRIBE (strong VLM)
# =====================================================================


def _extract_json(text):
    """Pull the first JSON object out of a model reply (handles ```json fences)."""
    if not text:
        return None
    s = text.strip()
    if "```" in s:
        parts = s.split("```")
        for part in parts:
            p = part
            if p.lstrip().lower().startswith("json"):
                p = p.lstrip()[4:]
            p = p.strip()
            if p.startswith("{"):
                s = p
                break
    try:
        return json.loads(s)
    except (json.JSONDecodeError, TypeError):
        pass
    # last resort: greedy braces
    m = re.search(r"\{.*\}", text, re.DOTALL)
    if m:
        try:
            return json.loads(m.group(0))
        except (json.JSONDecodeError, TypeError):
            return None
    return None


@dataclass
class Description:
    description: str = ""
    figures: list = field(default_factory=list)
    scene: str = ""
    attributes: list = field(default_factory=list)
    actions: list = field(default_factory=list)
    raw: str = ""

    def concepts(self):
        """Discrete concept strings to use as individual retrieval queries.

        The scene phrase is repeated (weighted) because subject/story is the
        signal the 4B was worst at and the highest-value thing to retrieve on.
        """
        out = []
        if self.scene:
            out.append(self.scene)
        out.extend(c for c in self.figures if isinstance(c, str) and c.strip())
        out.extend(c for c in self.attributes if isinstance(c, str) and c.strip())
        out.extend(c for c in self.actions if isinstance(c, str) and c.strip())
        # de-dup, preserve order
        seen, deduped = set(), []
        for c in out:
            cs = c.strip()
            key = cs.lower()
            if cs and key not in seen:
                seen.add(key)
                deduped.append(cs)
        return deduped


def parse_description(text):
    data = _extract_json(text) or {}
    return Description(
        description=str(data.get("description", "") or ""),
        figures=list(data.get("figures", []) or []),
        scene=str(data.get("scene", "") or ""),
        attributes=list(data.get("attributes", []) or []),
        actions=list(data.get("actions", []) or []),
        raw=text[:1500],
    )


async def describe_one(client, sem, image_url, max_tokens):
    async with sem:
        try:
            resp = await client.chat.completions.create(
                model=VLM_MODEL,
                messages=[
                    {"role": "system", "content": DESCRIBE_SYSTEM},
                    {
                        "role": "user",
                        "content": [
                            {"type": "image_url", "image_url": {"url": image_url}},
                            {"type": "text", "text": DESCRIBE_USER},
                        ],
                    },
                ],
                max_tokens=max_tokens,
                temperature=0.0,
                extra_body=VLM_EXTRA_BODY,
            )
            raw = (resp.choices[0].message.content or "").strip()
            if not raw:
                rc = resp.choices[0].message.model_dump().get("reasoning_content", "")
                raw = rc or ""
        except Exception as e:
            return {"error": str(e)[:200]}
    return {"description": parse_description(raw)}


# =====================================================================
# STEP 2 — RETRIEVE (semantic search -> candidate codes)
# =====================================================================


def retrieve_candidates(desc: Description, k_per_query, k_whole, pool_keep):
    """Map a Description to candidate Iconclass codes via semantic search.

    Strategy: run one query per extracted concept (k_per_query hits each) PLUS a
    query over the whole prose description (k_whole hits). Pool all hits, keep the
    best similarity seen per code, then return the top `pool_keep` codes by score.

    Returns (ordered_codes, debug_hits) where debug_hits maps code -> best sim.
    """
    pooled = {}  # code -> best similarity

    def add(query, limit):
        if not query or not query.strip():
            return
        try:
            hits = search_iconclass(query.strip(), limit=limit)
        except Exception:
            return
        for h in hits:
            code = h.get("code")
            if not code or code == "error":
                continue
            sim = float(h.get("similarity", 0.0))
            if code not in pooled or sim > pooled[code]:
                pooled[code] = sim

    # whole-description query (broad context)
    add(desc.description, k_whole)
    # per-concept queries (precise, the main driver)
    for concept in desc.concepts():
        add(concept, k_per_query)

    ordered = sorted(pooled.items(), key=lambda kv: kv[1], reverse=True)
    kept = [code for code, _ in ordered[:pool_keep]]
    return kept, dict(ordered[:pool_keep])


# =====================================================================
# STEP 3 — CONFIRM (binary VLM judge, optional) — eval_corrected.py pattern
# =====================================================================


def parse_yes_no(raw):
    if not raw:
        return False
    s = raw.strip().upper()
    tokens = s.replace(",", " ").replace(".", " ").split()
    first = tokens[0] if tokens else ""
    if first.startswith("YES"):
        return True
    if first.startswith("NO"):
        return False
    return "YES" in s


async def judge_one(client, sem, image_url, code, concept, max_tokens):
    prompt = JUDGE_PROMPT_TEMPLATE.format(concept=concept, code=code)
    async with sem:
        try:
            resp = await client.chat.completions.create(
                model=VLM_MODEL,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "image_url", "image_url": {"url": image_url}},
                            {"type": "text", "text": prompt},
                        ],
                    }
                ],
                max_tokens=max_tokens,
                temperature=0.0,
                extra_body=VLM_EXTRA_BODY,
            )
            raw = (resp.choices[0].message.content or "").strip()
            if not raw:
                rc = resp.choices[0].message.model_dump().get("reasoning_content", "")
                raw = rc or ""
        except Exception as e:
            raw = f"__error__: {e}"
    depicted = False if raw.startswith("__error__") else parse_yes_no(raw)
    return code, depicted


# =====================================================================
# Per-sample record + pipeline
# =====================================================================


@dataclass
class SampleRecord:
    index: int
    gt_codes: list
    image_url: str
    description: Description = None
    candidates: list = field(default_factory=list)  # retrieved, pre-confirm
    candidate_scores: dict = field(default_factory=dict)
    confirmed: list = field(default_factory=list)  # judge YES (post-confirm)
    error: str = None


async def describe_all(records, concurrency, max_tokens):
    client = AsyncOpenAI(base_url=ROUTER_BASE_URL, api_key=get_token())
    sem = asyncio.Semaphore(concurrency)
    tasks = [describe_one(client, sem, r.image_url, max_tokens) for r in records]
    # gather preserves order so results line up with `records`.
    results = await asyncio.gather(*tasks)
    for r, res in zip(records, results):
        if "error" in res:
            r.error = res["error"]
        else:
            r.description = res["description"]
    return client


async def confirm_all(client, records, concurrency, max_tokens):
    """Run the binary judge over every (record, candidate) and keep the YESes."""
    sem = asyncio.Semaphore(concurrency)
    tasks = []
    owners = []
    for r in records:
        if r.error or not r.candidates:
            continue
        for code in r.candidates:
            concept = _code_description(code) or code
            tasks.append(judge_one(client, sem, r.image_url, code, concept, max_tokens))
            owners.append(r)
    if not tasks:
        return
    print(f"  confirming {len(tasks)} candidate codes via binary judge…")
    results = await asyncio.gather(*tasks)
    for owner, (code, depicted) in zip(owners, results):
        if depicted:
            owner.confirmed.append(code)


# =====================================================================
# Scoring / aggregation
# =====================================================================


def score_predictions(records, use_confirmed):
    """Score each record's predicted codes vs GT. Returns list of metric dicts."""
    rows = []
    for r in records:
        if r.error:
            continue
        predicted = r.confirmed if use_confirmed else r.candidates
        h_f1 = _calculate_hierarchical_f1(predicted, r.gt_codes)
        h_p, h_r = _hierarchical_precision_recall(predicted, r.gt_codes)
        c_p, c_r = _code_precision_recall(predicted, r.gt_codes)
        rows.append(
            {
                "index": r.index,
                "n_pred": len(predicted),
                "n_gt": len(r.gt_codes),
                "hierarchical_f1": h_f1,
                "hier_precision": h_p,
                "hier_recall": h_r,
                "code_precision": c_p,
                "code_recall": c_r,
                "exact_match": set(predicted) == set(r.gt_codes),
            }
        )
    return rows


def aggregate(rows, label):
    n = len(rows)
    if n == 0:
        return {"label": label, "n": 0}

    def avg(k):
        return sum(r[k] for r in rows) / n

    return {
        "label": label,
        "n": n,
        "h_f1": avg("hierarchical_f1"),
        "hier_precision": avg("hier_precision"),
        "hier_recall": avg("hier_recall"),
        "code_precision": avg("code_precision"),
        "code_recall": avg("code_recall"),
        "exact": avg("exact_match"),
        "avg_pred": avg("n_pred"),
        "avg_gt": avg("n_gt"),
    }


def print_summary(summaries):
    hdr = (
        f"{'config':<26}{'n':>4}{'H-F1':>8}{'hierP':>8}{'hierR':>8}"
        f"{'codeP':>8}{'codeR':>8}{'exact':>7}{'#pred':>7}"
    )
    print("\n" + "=" * len(hdr))
    print("DESCRIBE -> RETRIEVE  vs  fine-tuned 4B v3 (brillfull-test)")
    print("=" * len(hdr))
    print(hdr)
    print("-" * len(hdr))
    for s in summaries:
        if s.get("n", 0) == 0:
            print(f"{s['label']:<26}{'—':>4}  (no scored samples)")
            continue
        print(
            f"{s['label']:<26}{s['n']:>4}{s['h_f1']:>7.1%}{s['hier_precision']:>8.1%}"
            f"{s['hier_recall']:>8.1%}{s['code_precision']:>8.1%}"
            f"{s['code_recall']:>8.1%}{s['exact']:>7.1%}{s['avg_pred']:>7.1f}"
        )
    print("-" * len(hdr))
    print(
        f"{'fine-tuned 4B v3 (ref)':<26}{'~':>4}{FT_4B_V3_BRILLFULL_HF1:>7.1%}"
        f"{'—':>8}{'~25%':>8}{'—':>8}{'~25%':>8}{'—':>7}{'—':>7}"
    )
    print("  (4B v3 H-F1 is contaminated-optimistic: it TRAINED on these images)")
    print("=" * len(hdr))


# =====================================================================
# Loading
# =====================================================================


def load_records(num_samples, seed, first_n=False):
    print(f"Loading {DATASET} split={SPLIT}…")
    ds = load_dataset(DATASET, split=SPLIT)
    if first_n:
        # range(N) selection — matches eval_sft.py's ds.select(range(N)) so the
        # retrieval pool aligns by index with the fine-tuned 4B's predictions.
        idxs = list(range(min(num_samples, len(ds))))
    else:
        # fixed-seed shuffle so the sample is reproducible and not just the first-N
        idxs = list(range(len(ds)))
        import random

        random.Random(seed).shuffle(idxs)
        idxs = idxs[:num_samples]

    records = []
    for i in idxs:
        ex = ds[i]
        img = first_image(ex)
        gt = get_gt_codes(ex)
        if img is None or not gt:
            continue
        records.append(
            SampleRecord(index=i, gt_codes=gt, image_url=image_to_data_url(img))
        )
    print(f"Got {len(records)} usable samples (seed={seed}).")
    return records


# =====================================================================
# Probe
# =====================================================================


async def probe(seed):
    token = get_token()
    if not token:
        sys.exit("ERROR: no HF token (run `hf auth login`).")
    print("Building semantic search index (one-time, ~75s on CPU)…")
    t0 = time.time()
    build_search_index()
    print(f"  index ready in {time.time() - t0:.1f}s")

    print("\nSemantic search sanity (text -> code):")
    for q in ["the crucifixion of Christ", "a pheasant in a landscape"]:
        hits = search_iconclass(q, limit=3)
        print(f"  {q!r}:")
        for h in hits:
            print(
                f"    {h['code']:<14} sim={h.get('similarity', 0):.3f}  "
                f"{h['description'][:50]}"
            )

    records = load_records(1, seed)
    if not records:
        sys.exit("No usable sample to probe.")
    r = records[0]
    client = AsyncOpenAI(base_url=ROUTER_BASE_URL, api_key=token)
    sem = asyncio.Semaphore(1)

    print(f"\nWarming up + describing 1 image (index {r.index})…")
    t0 = time.time()
    res = await describe_one(client, sem, r.image_url, 700)
    if "error" in res:
        sys.exit(f"VLM describe failed: {res['error']}")
    desc = res["description"]
    print(f"  described in {time.time() - t0:.1f}s")
    print(f"  scene:       {desc.scene!r}")
    print(f"  figures:     {desc.figures}")
    print(f"  attributes:  {desc.attributes}")
    print(f"  actions:     {desc.actions}")
    print(f"  description: {desc.description[:200]}")

    cands, scores = retrieve_candidates(desc, k_per_query=4, k_whole=8, pool_keep=8)
    print(f"\n  retrieved {len(cands)} candidate codes:")
    for c in cands:
        print(f"    {c:<16} sim={scores.get(c, 0):.3f}  {_code_description(c)}")
    print(f"\n  GT codes ({len(r.gt_codes)}): {r.gt_codes}")
    h_f1 = _calculate_hierarchical_f1(cands, r.gt_codes)
    _, h_r = _hierarchical_precision_recall(cands, r.gt_codes)
    print(f"  -> H-F1={h_f1:.1%}  hier-recall={h_r:.1%}")


# =====================================================================
# Main
# =====================================================================


def run_config(records, k_per_query, k_whole, pool_keep, label):
    """Retrieve (pre-confirm) for a given K config and score it."""
    for r in records:
        if r.error:
            continue
        cands, scores = retrieve_candidates(
            r.description, k_per_query, k_whole, pool_keep
        )
        r.candidates = cands
        r.candidate_scores = scores
    rows = score_predictions(records, use_confirmed=False)
    return aggregate(rows, label), rows


def main():
    sys.stdout.reconfigure(line_buffering=True)
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--num-samples", type=int, default=120)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument(
        "--first-n",
        action="store_true",
        help="select range(N) (matches eval_sft.py) instead of a seed shuffle",
    )
    ap.add_argument("--k", type=int, default=8, help="pool_keep: #codes kept/image")
    ap.add_argument(
        "--k-per-query", type=int, default=4, help="hits per extracted concept"
    )
    ap.add_argument(
        "--k-whole", type=int, default=8, help="hits for the whole-description query"
    )
    ap.add_argument(
        "--sweep-k",
        type=str,
        default=None,
        help="comma-separated pool_keep values to sweep, e.g. 4,6,8,12,16",
    )
    ap.add_argument(
        "--confirm",
        action="store_true",
        help="run the binary VLM judge to confirm retrieved codes (lifts precision)",
    )
    ap.add_argument("--describe-concurrency", type=int, default=8)
    ap.add_argument("--judge-concurrency", type=int, default=10)
    ap.add_argument("--describe-max-tokens", type=int, default=700)
    ap.add_argument("--judge-max-tokens", type=int, default=200)
    ap.add_argument("--probe", action="store_true")
    ap.add_argument("--output", type=str, default="describe_retrieve_results.json")
    args = ap.parse_args()

    if args.probe:
        asyncio.run(probe(args.seed))
        return

    token = get_token()
    if not token:
        sys.exit("ERROR: no HF token (run `hf auth login`).")

    # Build the index up-front (one-time ~75s) so timing of the VLM stage is clean.
    print("Building semantic search index (one-time, ~75s on CPU)…")
    t0 = time.time()
    build_search_index()
    print(f"  index ready in {time.time() - t0:.1f}s")

    records = load_records(args.num_samples, args.seed, first_n=args.first_n)
    if not records:
        sys.exit("No usable samples.")

    # ---- STEP 1: describe every image with the strong VLM.
    print(f"\n[1/3] Describing {len(records)} images with {VLM_MODEL}…")
    t0 = time.time()
    client = asyncio.run(
        describe_all(records, args.describe_concurrency, args.describe_max_tokens)
    )
    n_desc = sum(1 for r in records if r.description is not None)
    n_err = sum(1 for r in records if r.error)
    print(
        f"  described {n_desc}/{len(records)} ({n_err} errors) in {time.time() - t0:.1f}s"
    )

    # ---- STEP 2: retrieve + score. Sweep K if requested, else single config.
    print("\n[2/3] Retrieving candidate codes via semantic search + scoring…")
    summaries = []
    best = None
    if args.sweep_k:
        ks = [int(x) for x in args.sweep_k.split(",") if x.strip()]
        for k in ks:
            summ, _rows = run_config(
                records, args.k_per_query, args.k_whole, k, f"retrieve K={k}"
            )
            summaries.append(summ)
            if best is None or summ.get("h_f1", -1) > best.get("h_f1", -1):
                best, best_k = summ, k
        print(f"  best pool_keep K = {best_k} (H-F1={best['h_f1']:.1%})")
        pool_keep = best_k
    else:
        summ, _rows = run_config(
            records, args.k_per_query, args.k_whole, args.k, f"retrieve K={args.k}"
        )
        summaries.append(summ)
        best, pool_keep = summ, args.k

    # ---- STEP 3: optional judge confirm (uses the best/chosen K candidates).
    confirm_summary = None
    if args.confirm:
        print("\n[3/3] Confirming retrieved codes with binary VLM judge…")
        # ensure candidates correspond to the chosen K
        for r in records:
            if r.error:
                continue
            cands, scores = retrieve_candidates(
                r.description, args.k_per_query, args.k_whole, pool_keep
            )
            r.candidates, r.candidate_scores = cands, scores
        t0 = time.time()
        asyncio.run(
            confirm_all(client, records, args.judge_concurrency, args.judge_max_tokens)
        )
        crows = score_predictions(records, use_confirmed=True)
        confirm_summary = aggregate(crows, f"retrieve+confirm K={pool_keep}")
        summaries.append(confirm_summary)
        print(f"  confirm done in {time.time() - t0:.1f}s")
    else:
        print("\n[3/3] (skipped judge confirm — pass --confirm to enable)")

    # ---- Report.
    print_summary(summaries)

    # Verdict vs the 4B wall.
    headline = confirm_summary or best
    print("\nVERDICT (vs fine-tuned 4B v3 on the clean ruler):")
    df = headline["h_f1"] - FT_4B_V3_BRILLFULL_HF1
    print(
        f"  describe->retrieve H-F1 = {headline['h_f1']:.1%}  "
        f"({'+' if df >= 0 else ''}{df:.1%} vs 4B v3's {FT_4B_V3_BRILLFULL_HF1:.0%})"
    )
    print(
        f"  hier-recall = {headline['hier_recall']:.1%}, "
        f"exact-set code recall = {headline['code_recall']:.1%}  "
        f"(4B wall ≈ 25% exact-set recall)"
    )

    # ---- Persist everything for inspection.
    out = {
        "summary": {
            "dataset": DATASET,
            "split": SPLIT,
            "vlm_model": VLM_MODEL,
            "n_samples": len(records),
            "seed": args.seed,
            "k_per_query": args.k_per_query,
            "k_whole": args.k_whole,
            "confirmed": bool(args.confirm),
            "ft_4b_v3_brillfull_hf1": FT_4B_V3_BRILLFULL_HF1,
            "configs": summaries,
        },
        "records": [
            {
                "index": r.index,
                "gt_codes": r.gt_codes,
                "scene": r.description.scene if r.description else None,
                "figures": r.description.figures if r.description else None,
                "attributes": r.description.attributes if r.description else None,
                "actions": r.description.actions if r.description else None,
                "description": r.description.description if r.description else None,
                "retrieved_codes": r.candidates,
                "retrieved_scores": {
                    k: round(v, 4) for k, v in r.candidate_scores.items()
                },
                "confirmed_codes": r.confirmed if args.confirm else None,
                "error": r.error,
            }
            for r in records
        ],
    }
    with open(args.output, "w") as f:
        json.dump(out, f, indent=2, default=str)
    print(f"\nSaved → {args.output}")


if __name__ == "__main__":
    main()