example_pool.py

"""CUE-keyed example pool: retrieval and injection into decoded manuals.

The pool is large (one embedding plus example texts per training session), so it ships
as a separate Hugging Face **dataset** rather than inside the model repo. Pools built by
CUE-clean's ``cue build-example-pool`` are also loadable straight from disk.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

from .faiss_index import EmbeddingIndex
from .log import log

TAG = "example_pool"
POOL_DIRNAME = "example_pool"
META_NAME = "meta.jsonl"
INDEX_NAME = "index.faiss"
MANIFEST_NAME = "manifest.json"
VECTORS_NAME = "cue_vectors.npy"


@dataclass(frozen=True)
class PoolExample:
    text: str
    kind: str  # sim_contrast | human_contrast | style


@dataclass
class ExamplePool:
    session_ids: list[str]
    cue_vectors: np.ndarray
    examples: list[list[PoolExample]]
    index: EmbeddingIndex
    checkpoint_fingerprint: str = ""

    @property
    def dim(self) -> int:
        return int(self.cue_vectors.shape[1])

    @property
    def ntotal(self) -> int:
        return len(self.session_ids)

    def retrieve(
        self,
        cue: list[float] | np.ndarray,
        *,
        k_sessions: int = 8,
        n_general: int = 2,
        n_specific: int = 6,
        exclude_session_id: str | None = None,
    ) -> list[str]:
        query = np.asarray(cue, dtype=np.float32).reshape(1, -1)
        _, idxs = self.index.search(query, min(k_sessions, max(1, self.ntotal)))
        general: list[str] = []
        specific: list[str] = []
        seen: set[str] = set()
        for idx in idxs[0]:
            if idx < 0:
                continue
            i = int(idx)
            if exclude_session_id and self.session_ids[i] == exclude_session_id:
                continue
            for ex in self.examples[i]:
                key = ex.text.casefold()
                if key in seen:
                    continue
                seen.add(key)
                if ex.kind == "sim_contrast" and len(general) < n_general:
                    general.append(ex.text)
                elif ex.kind != "sim_contrast" and len(specific) < n_specific:
                    specific.append(ex.text)
        out = general + specific
        if len(out) < n_general + n_specific:
            for idx in idxs[0]:
                if idx < 0:
                    continue
                i = int(idx)
                for ex in self.examples[i]:
                    key = ex.text.casefold()
                    if key in seen:
                        continue
                    seen.add(key)
                    out.append(ex.text)
                    if len(out) >= n_general + n_specific:
                        break
                if len(out) >= n_general + n_specific:
                    break
        return out[: n_general + n_specific]

    @classmethod
    def from_rows(
        cls,
        session_ids: list[str],
        cue_vectors: np.ndarray,
        examples: list[list[PoolExample]],
        *,
        checkpoint_fingerprint: str = "",
    ) -> ExamplePool:
        vectors = np.asarray(cue_vectors, dtype=np.float32)
        if vectors.ndim != 2:
            raise ValueError(f"cue_vectors must be [N, D], got {vectors.shape}")
        index = EmbeddingIndex(int(vectors.shape[1]))
        index.add(vectors)
        return cls(
            list(session_ids),
            vectors,
            list(examples),
            index,
            checkpoint_fingerprint,
        )

    @classmethod
    def from_directory(cls, directory: str | Path) -> ExamplePool:
        """Load a pool laid out by CUE-clean's ``cue build-example-pool``."""

        directory = Path(directory)
        if (directory / POOL_DIRNAME / META_NAME).is_file():
            directory = directory / POOL_DIRNAME
        vectors = np.load(directory / VECTORS_NAME).astype(np.float32)
        index = EmbeddingIndex.load(directory / INDEX_NAME, int(vectors.shape[1]))
        session_ids: list[str] = []
        examples: list[list[PoolExample]] = []
        with (directory / META_NAME).open(encoding="utf-8") as handle:
            for line in handle:
                if not line.strip():
                    continue
                row = json.loads(line)
                session_ids.append(str(row["session_id"]))
                examples.append(
                    [
                        PoolExample(text=str(e["text"]), kind=str(e["kind"]))
                        for e in (row.get("examples") or [])
                    ]
                )
        fingerprint = ""
        manifest_path = directory / MANIFEST_NAME
        if manifest_path.is_file():
            manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
            fingerprint = str(manifest.get("checkpoint_fingerprint") or "")
        log(TAG, f"loaded {len(session_ids)} pool sessions from {directory}")
        return cls(session_ids, vectors, examples, index, fingerprint)

    @classmethod
    def from_hub_dataset(cls, repo_id: str, *, split: str = "train") -> ExamplePool:
        """Load the pool from a dataset of {session_id, cue_embedding, examples}.

        Accepts a Hub dataset id or a ``save_to_disk`` directory.
        """

        try:
            from datasets import load_dataset, load_from_disk
        except ImportError as exc:  # pragma: no cover - dependency hint
            raise ImportError(
                "example-pool retrieval from the Hub needs `datasets`. "
                "Install with: pip install 'cue-hf[retrieval]'"
            ) from exc
        if Path(repo_id).is_dir():
            dataset = load_from_disk(repo_id)
            if not hasattr(dataset, "column_names") or isinstance(dataset.column_names, dict):
                dataset = dataset[split]
        else:
            dataset = load_dataset(repo_id, split=split)
        session_ids = [str(x) for x in dataset["session_id"]]
        vectors = np.asarray(dataset["cue_embedding"], dtype=np.float32)
        examples = [
            [
                PoolExample(text=str(e["text"]), kind=str(e["kind"]))
                for e in (row or [])
            ]
            for row in dataset["examples"]
        ]
        fingerprint = ""
        info = getattr(dataset, "info", None)
        if info is not None and getattr(info, "description", ""):
            fingerprint = _fingerprint_from_description(info.description)
        log(TAG, f"loaded {len(session_ids)} pool sessions from dataset {repo_id}")
        return cls.from_rows(
            session_ids, vectors, examples, checkpoint_fingerprint=fingerprint
        )


def _fingerprint_from_description(description: str) -> str:
    for line in description.splitlines():
        if line.lower().startswith("checkpoint_fingerprint:"):
            return line.split(":", 1)[1].strip()
    return ""


def inject_examples(
    manual: dict[str, Any] | None,
    *,
    pool: ExamplePool,
    cue_embedding: list[float] | np.ndarray,
    k_sessions: int = 8,
    n_general: int = 2,
    n_specific: int = 6,
    exclude_session_id: str | None = None,
) -> dict[str, Any] | None:
    """Replace a decoded manual's examples with pool neighbors of its CUE vector.

    Mirrors CUE-clean evaluation, which clears the general block's examples and fills the
    user-specific block from retrieval; after merging, that is the manual's example list.
    """

    if manual is None:
        return None
    retrieved = pool.retrieve(
        cue_embedding,
        k_sessions=k_sessions,
        n_general=n_general,
        n_specific=n_specific,
        exclude_session_id=exclude_session_id,
    )
    return {**manual, "examples": retrieved}