agent/ret_index.py

#!/usr/bin/env python3
"""Parametrized Iconclass retriever index — swap the embedding model & doc representation.

The production retriever (`grpo-tools/iconclass_tools.py`) is hardwired to
`all-MiniLM-L6-v2` and indexes each code by its LEAF text only. That makes
sibling codes collapse together ("head turned to the left" vs "...right") and
caps exact-leaf retrieval at ~18% — the bottleneck the retriever fine-tune
attacks.

This module builds the SAME catalogue of Iconclass codes (reusing
`_get_all_iconclass_codes` so coverage matches the baseline exactly, 40,457
codes) but lets us choose:
  * `model_name`   — any sentence-transformers model (off-the-shelf or our FT one)
  * `doc_mode`     — how each code is rendered to text for indexing:
        'leaf'        : the leaf description only            (baseline behaviour)
        'parent_leaf' : "<parent leaf>: <leaf>"              (one level of context)
        'path'        : full root->leaf path, " > "-joined   (max context)
  * `query_prompt` / `doc_prompt` — optional instruction prefixes (BGE/E5/GTE)

Doc embeddings are cached to ./ret_cache/<slug>.npz so repeated evals are fast.
Used by ret_eval.py (intrinsic metrics + ceiling probe) and ret_retrieve_pool.py
(regenerate the fusion pool with a new retriever, reusing cached descriptions).
"""

from __future__ import annotations

import hashlib
import os
import sys
from functools import lru_cache

import numpy as np

# Reuse the EXACT catalogue the production retriever indexes, so the only thing
# that changes vs the baseline is the embedding model + doc representation.
sys.path.insert(0, os.path.abspath(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 _get_all_iconclass_codes  # noqa: E402

from iconclass import init as _ic_init  # noqa: E402

_ic = _ic_init()

CACHE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ret_cache")
os.makedirs(CACHE_DIR, exist_ok=True)


# --------------------------------------------------------------------------- catalogue
@lru_cache(maxsize=1)
def code_catalog():
    """[(code, leaf_text)] — the 40,457 codes the baseline retriever indexes."""
    return list(_get_all_iconclass_codes())


@lru_cache(maxsize=200_000)
def _leaf_text(code):
    try:
        n = _ic[code]
        return (n() if callable(n) else str(n)) or code
    except Exception:
        return code


@lru_cache(maxsize=200_000)
def path_text(code, joiner=" > "):
    """Full root->leaf textual path for a code (drops any (KEY)/(+x) suffix)."""
    base = code.split("(")[0].strip()
    try:
        p = _ic[base].obj.get("p", []) or [base]
    except Exception:
        p = [base]
    parts = [_leaf_text(a) for a in p]
    return joiner.join(x for x in parts if x)


def doc_text(code, leaf, mode):
    """Render a code to the string we index it by."""
    if mode == "leaf":
        return leaf
    if mode == "parent_leaf":
        base = code.split("(")[0].strip()
        try:
            p = _ic[base].obj.get("p", [])
        except Exception:
            p = []
        if len(p) >= 2:
            return f"{_leaf_text(p[-2])}: {leaf}"
        return leaf
    if mode == "path":
        return path_text(code)
    raise ValueError(f"unknown doc_mode {mode!r}")


# --------------------------------------------------------------------------- index
def _slug(model_name, doc_mode, doc_prompt):
    h = hashlib.md5(
        f"{model_name}|{doc_mode}|{doc_prompt}|{len(code_catalog())}".encode()
    ).hexdigest()[:10]
    safe = model_name.replace("/", "__")
    return f"{safe}__{doc_mode}__{h}"


class RetrieverIndex:
    def __init__(
        self,
        model_name="all-MiniLM-L6-v2",
        doc_mode="leaf",
        query_prompt=None,
        doc_prompt=None,
        device=None,
        trust_remote_code=False,
        batch_size=512,
    ):
        from sentence_transformers import SentenceTransformer

        self.model_name = model_name
        self.doc_mode = doc_mode
        self.query_prompt = query_prompt or ""
        self.doc_prompt = doc_prompt or ""
        self.batch_size = batch_size

        cat = code_catalog()
        self.codes = [c for c, _ in cat]
        self.leaves = [leaf for _, leaf in cat]
        self.doc_texts = [doc_text(c, leaf, doc_mode) for c, leaf in cat]

        print(f"[index] loading {model_name} (doc_mode={doc_mode})…", flush=True)
        self.model = SentenceTransformer(
            model_name, device=device, trust_remote_code=trust_remote_code
        )

        self.embeddings = self._load_or_build_doc_embeddings()
        print(
            f"[index] ready: {len(self.codes)} codes, dim={self.embeddings.shape[1]}",
            flush=True,
        )

    def _load_or_build_doc_embeddings(self):
        slug = _slug(self.model_name, self.doc_mode, self.doc_prompt)
        path = os.path.join(CACHE_DIR, f"{slug}.npz")
        if os.path.exists(path):
            print(f"[index] cached doc embeddings: {path}", flush=True)
            return np.load(path)["emb"]
        print(
            f"[index] encoding {len(self.doc_texts)} doc texts (one-time)…", flush=True
        )
        texts = (
            [self.doc_prompt + t for t in self.doc_texts]
            if self.doc_prompt
            else self.doc_texts
        )
        emb = self.model.encode(
            texts,
            batch_size=self.batch_size,
            show_progress_bar=True,
            convert_to_numpy=True,
            normalize_embeddings=True,
        ).astype(np.float32)
        np.savez_compressed(path, emb=emb)
        return emb

    def encode_queries(self, queries):
        texts = (
            [self.query_prompt + q for q in queries] if self.query_prompt else queries
        )
        return self.model.encode(
            texts,
            batch_size=self.batch_size,
            convert_to_numpy=True,
            normalize_embeddings=True,
        ).astype(np.float32)

    def search_many(self, queries, k=10):
        """Batch search. Returns list (per query) of [(code, leaf, sim)]."""
        if not queries:
            return []
        q = self.encode_queries(queries)  # (Q, d), normalized
        sims = q @ self.embeddings.T  # (Q, N) cosine (both normalized)
        out = []
        topk = min(k, sims.shape[1])
        idx = np.argpartition(-sims, topk - 1, axis=1)[:, :topk]
        for i in range(sims.shape[0]):
            row = idx[i]
            row = row[np.argsort(-sims[i, row])]
            out.append([(self.codes[j], self.leaves[j], float(sims[i, j])) for j in row])
        return out

    def search(self, query, k=10):
        return self.search_many([query], k)[0]