llm-ingestion-okf/tools/okf_consume_measure.py

"""Measure a consumption pre-pass: hit@k per question, against a chance baseline.

Measure, don't build. This instrument calls `tools/okf_consume.py` exactly as a
consumer would and scores what comes back; it contains no ranking of its own,
because a scorer sharing code with the thing it scores measures its own
agreement rather than the ranker's performance.

**The gold set is an INPUT, never a constant here.** It arrives as a JSON file
naming, per question, the gold document. Two reasons, and the second is the
binding one:

- An instrument with a corpus baked in measures one corpus.
- This file is tracked in a PUBLIC repository. A gold set names documents in a
  consumer's corpus, and this repository's standing rule keeps corpus paths and
  document titles out of tracked files. The instrument is publishable; the
  answer key is not, so it lives beside the run and not beside the code.

**hit@k is scored at DOCUMENT granularity by default**, which is a weaker claim
than concept granularity and is labelled as such wherever the number is used: a
gold document counts as hit at k when any of the top-k excerpts has that
document as its `concept_id` path prefix, and the reported rank is the position
of the first such excerpt. A gold document holding one concept is the only case
where document and concept granularity coincide.

**Every row carries its own chance baseline**, computed two ways, because a
document-prefix hit is easier for a large gold document: a gold set spanning one
to fifty concepts makes a raw hit count uninterpretable without the baseline it
should be read against.

The questions file:

```json
{"questions": [{"question": "...", "gold_document": "...", "note": "optional"}],
 "negatives": [{"question": "...", "label": "..."}]}
```
"""

from __future__ import annotations

import argparse
import json
import random
import sys
import time
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

import okf_consume  # noqa: E402
from llm_ingestion_okf.materialize import parse_frontmatter  # noqa: E402

#: The seed every empirical baseline in a published table is drawn with. Fixed
#: so a reader re-running this reproduces the same figure; the analytic form
#: beside it is what makes the empirical one checkable at all.
DEFAULT_SEED = 20260907

#: Trials for the empirical baseline. Enough that it lands within about half a
#: percentage point of the analytic value, which is the agreement actually
#: observed -- NOT three decimal places, and this instrument does not claim it.
DEFAULT_TRIALS = 20_000


@dataclass(frozen=True)
class Row:
    """One question, measured."""

    question: str
    gold_document: str
    gold_concepts: int
    hit: bool
    rank: int | None
    considered: int
    withheld: int
    delivered: int
    spent: int
    payload_bytes: int
    seconds: float
    chance_analytic: float
    chance_empirical: float


def hit_rank(excerpts: Sequence[Mapping[str, object]], gold_document: str) -> int | None:
    """The rank of the first excerpt inside `gold_document`, or `None`.

    Prefix on the `concept_id`, which in this library is the concept's
    bundle-relative path with the suffix removed -- so the test is a string
    operation on a field the payload already carries. The `+ "/"` is not
    cosmetic: without it, a gold document `bilag-7` would match a concept in
    `bilag-70`.
    """
    for excerpt in excerpts:
        concept_id = str(excerpt.get("concept_id", ""))
        if concept_id == gold_document or concept_id.startswith(f"{gold_document}/"):
            rank = excerpt.get("rank")
            return int(rank) if isinstance(rank, int) else None
    return None


def chance_analytic(gold_concepts: int, corpus: int, k: int) -> float:
    """P(a uniform random k-of-corpus draw touches any of `gold_concepts`).

    The complement of drawing k items that all miss, computed as a product of
    hypergeometric terms rather than with factorials, so it stays exact for the
    sizes a corpus actually reaches.
    """
    if gold_concepts <= 0 or k <= 0 or corpus <= 0:
        return 0.0
    miss = 1.0
    for drawn in range(min(k, corpus)):
        remaining = corpus - drawn
        available = corpus - gold_concepts - drawn
        if available <= 0:
            return 1.0
        miss *= available / remaining
    return 1.0 - miss


def chance_empirical(
    gold_concepts: int,
    corpus: int,
    k: int,
    *,
    trials: int = DEFAULT_TRIALS,
    seed: int = DEFAULT_SEED,
) -> float:
    """The same probability, drawn rather than derived.

    Kept beside the analytic form because two routes to one number is what makes
    either believable: an off-by-one in the closed form and a bug in the sampler
    do not agree by accident.
    """
    rng = random.Random(seed)
    pool = range(corpus)
    gold = set(range(gold_concepts))
    hits = sum(1 for _ in range(trials) if gold & set(rng.sample(pool, k)))
    return hits / trials


def document_sizes(concept_ids: Sequence[str]) -> dict[str, int]:
    """How many concepts each top-level document holds."""
    sizes: dict[str, int] = {}
    for concept_id in concept_ids:
        document = concept_id.split("/", 1)[0]
        sizes[document] = sizes.get(document, 0) + 1
    return sizes


def measure_question(
    bundle: Path, question: str, gold_document: str, *, k: int, sizes: Mapping[str, int]
) -> Row:
    started = time.perf_counter()
    payload = okf_consume.build_payload(bundle, question=question, k=k)
    elapsed = time.perf_counter() - started
    text = okf_consume.serialise(payload)
    excerpts = payload["excerpts"]
    counts = payload["denominators"]
    budget = payload["budget"]
    assert isinstance(excerpts, list) and isinstance(counts, dict) and isinstance(budget, dict)
    corpus = int(counts["considered"])
    gold_concepts = sizes.get(gold_document, 0)
    rank = hit_rank(excerpts, gold_document)
    return Row(
        question=question,
        gold_document=gold_document,
        gold_concepts=gold_concepts,
        hit=rank is not None,
        rank=rank,
        considered=corpus,
        withheld=int(counts["withheld"]),
        delivered=int(counts["delivered"]),
        spent=int(budget["spent"]),
        payload_bytes=len(text.encode("utf-8")),
        seconds=round(elapsed, 3),
        chance_analytic=round(chance_analytic(gold_concepts, corpus, k), 3),
        chance_empirical=round(chance_empirical(gold_concepts, corpus, k), 3),
    )


def token_reach(bundle: Path, question: str) -> dict[str, int]:
    """How many concepts each of the question's tokens reaches.

    The known-negative control's own PRECONDITION, measured before any payload
    is built: a question whose tokens all reach real corpus text is not a
    question the bundle has no answer to, and scoring it as one would report a
    property of the question as a property of the ranker.
    """
    root_index = bundle / okf_consume.DEFAULT_PROFILE.index.name
    root_bundle_id = parse_frontmatter(root_index).get("bundle_id", "")
    corpus = [
        set(
            okf_consume.normalise(
                f"{concept.title} {concept.concept_id.replace('/', ' ')} {concept.body}"
            )
        )
        for concept in (
            okf_consume.read_concept(
                bundle / f"{concept_id}{okf_consume.CONCEPT_SUFFIX}",
                bundle_root=bundle,
                root_bundle_id=root_bundle_id,
            )
            for concept_id in okf_consume.enumerate_concepts(bundle)
        )
    ]
    return {
        token: sum(
            1
            for tokens in corpus
            if any(okf_consume.tokens_match(token, other) for other in tokens)
        )
        for token in okf_consume.normalise(question)
    }


def parse_args(argv: list[str] | None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("bundle", type=Path, help="the OKF bundle to measure against")
    parser.add_argument(
        "--questions", type=Path, required=True, help="JSON gold set (see the module docstring)"
    )
    parser.add_argument("--k", type=int, default=okf_consume.DEFAULT_K)
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    try:
        spec = json.loads(args.questions.read_text(encoding="utf-8"))
    except OSError as error:
        print(f"okf_consume_measure: FAILED - {error}", file=sys.stderr)
        return 2
    sizes = document_sizes(okf_consume.enumerate_concepts(args.bundle))
    rows = [
        measure_question(
            args.bundle, entry["question"], entry["gold_document"], k=args.k, sizes=sizes
        )
        for entry in spec.get("questions", [])
    ]
    negatives = []
    for entry in spec.get("negatives", []):
        reach = token_reach(args.bundle, entry["question"])
        started = time.perf_counter()
        # `withheld_full`: the row below reports the SET of rules a question
        # fell under, so it needs every drop and not the nearest N.
        payload = okf_consume.build_payload(
            args.bundle, question=entry["question"], k=args.k, withheld_full=True
        )
        elapsed = time.perf_counter() - started
        counts, budget, block = (
            payload["denominators"],
            payload["budget"],
            payload["withheld"],
        )
        assert isinstance(counts, dict) and isinstance(budget, dict) and isinstance(block, dict)
        withheld = block["nearest"]
        assert isinstance(withheld, list)
        negatives.append(
            {
                "question": entry["question"],
                "label": entry.get("label", ""),
                "token_reach": reach,
                "tokens_reaching_nothing": sum(1 for value in reach.values() if value == 0),
                "tokens": len(reach),
                "delivered": int(counts["delivered"]),
                "withheld": int(counts["withheld"]),
                "considered": int(counts["considered"]),
                "spent": int(budget["spent"]),
                "payload_bytes": len(okf_consume.serialise(payload).encode("utf-8")),
                "seconds": round(elapsed, 3),
                "rules": sorted({str(entry["rule"]) for entry in withheld}),
            }
        )
    print(
        json.dumps(
            {
                "k": args.k,
                "seed": DEFAULT_SEED,
                "trials": DEFAULT_TRIALS,
                "documents": len(sizes),
                "rows": [row.__dict__ for row in rows],
                "hits": sum(1 for row in rows if row.hit),
                "expected_by_chance": round(sum(row.chance_analytic for row in rows), 2),
                "negatives": negatives,
            },
            ensure_ascii=False,
            indent=2,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())