eval/build_dataset.py

"""Generate the retrieval evaluation dataset.

Samples chunks and asks an LLM for questions each one answers. The chunk a
question came from is that question's expected answer.

    uv run python -m eval.build_dataset
    uv run python -m eval.build_dataset --sample 50
"""

import argparse
import json
import os
import random
import re
import sys
from concurrent.futures import ThreadPoolExecutor

from dotenv import load_dotenv
from llama_index.core.base.llms.types import ChatMessage, MessageRole

from rag.config import CHUNKS_PATH, EVAL_DATASET_PATH, PROVIDERS
from rag.providers import make_llm
from rag.types import Chunk, EvalRow

SAMPLE_SIZE = 200
SEED = 20260804
CONCURRENCY = 16
MIN_CHARS = 400
CANDIDATES = 4

WORD = re.compile(r"[a-z_][a-z0-9_]+")

# Real Stack Overflow titles used as few-shot prompt examples 
#
#   stackoverflow.com/questions/29483365   multiline string literal
#   stackoverflow.com/questions/24990520   integer to string
#   stackoverflow.com/questions/39204908   release / debug builds with cfg
#   stackoverflow.com/questions/57756927   main.rs and lib.rs
#   stackoverflow.com/questions/26469715   asserting a panic in a test
#   stackoverflow.com/questions/31192956   reading and writing files
EXAMPLES = (
    "What is the syntax for a multiline string literal?",
    "How do I convert from an integer to a string?",
    "How to check release / debug builds using cfg in Rust?",
    "Rust modules confusion when there is main.rs and lib.rs",
    "How do I write a Rust unit test that ensures that a panic has occurred?",
    "What's the de-facto way of reading and writing files in Rust 1.x?",
)

PROMPT = """Below is an excerpt from the official Rust documentation.

---------------------
{context}
---------------------

Write {count} different questions that this excerpt answers, as a Rust
programmer would type them into a search box before finding this page.

Match the register of these real questions:
{examples}

Rules:
- Under twelve words each.
- One thing per question. No compound questions.
- Reach for the words a programmer would use before reading this page, not the
  excerpt's own phrasing.
- Do not mention "the excerpt", "the text", or "the documentation".

Reply with one question per line and nothing else."""

JUDGE = """Below is an excerpt from the official Rust documentation, followed by
questions someone might search for.

---------------------
{context}
---------------------

{questions}

For each question, reply KEEP or DROP.

KEEP if this excerpt is where a reader searching that question should land: it
answers them directly, and more specifically than the rest of the documentation
would.

DROP if the excerpt only touches the subject in passing, or if the question is
broad enough that several other pages would answer it just as well.

Reply with one verdict per line, in order, as "1. KEEP" or "1. DROP", and
nothing else."""


def lexical_overlap(question: str, text: str) -> float:
    asked = set(WORD.findall(question.lower()))
    if not asked:
        return 1.0
    return len(asked.intersection(WORD.findall(text.lower()))) / len(asked)


def sample_chunks(sample_size: int) -> list[Chunk]:
    with CHUNKS_PATH.open(encoding="utf-8") as handle:
        chunks = [json.loads(line) for line in handle]
    usable = [chunk for chunk in chunks if len(chunk["text"]) >= MIN_CHARS]
    print(f"{len(chunks):,} chunks, {len(usable):,} long enough to sample from")
    random.seed(SEED)
    return random.sample(usable, min(sample_size, len(usable)))


def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--sample", type=int, default=SAMPLE_SIZE)
    parser.add_argument("--provider", default="openai", choices=sorted(PROVIDERS))
    args = parser.parse_args(argv[1:])

    load_dotenv(".env")
    spec = PROVIDERS[args.provider]
    api_key = os.environ.get(spec.env_var, "").strip()
    if not api_key:
        print(f"Set {spec.env_var} to generate the dataset.", file=sys.stderr)
        return 2
    if not CHUNKS_PATH.exists():
        print(f"No {CHUNKS_PATH}. Run: uv run python -m ingest.parse_books", file=sys.stderr)
        return 2

    llm = make_llm(args.provider, api_key, spec.default_model)
    chunks = sample_chunks(args.sample)
    by_id = {chunk["id"]: chunk for chunk in chunks}

    prompt = PROMPT.replace("{examples}", "\n".join(f"  {q}" for q in EXAMPLES))

    def say(content: str, chunk_id: str) -> str | None:
        try:
            return str(llm.chat([ChatMessage(role=MessageRole.USER, content=content)]).message.content)
        except Exception as error:
            print(f"  {chunk_id}: {error}", file=sys.stderr)
            return None

    def numbered(reply: str) -> list[str]:
        """Strip whatever list markers the model chose to use."""
        return [
            re.sub(r"^\s*(?:[-*\d.)]+\s*)+", "", line).strip()
            for line in reply.splitlines()
            if line.strip()
        ]

    def ask(chunk: Chunk) -> str | None:
        reply = say(prompt.format(context=chunk["text"], count=CANDIDATES), chunk["id"])
        if reply is None:
            return None
        candidates = [q for q in numbered(reply) if q.endswith("?") or len(q.split()) >= 4]
        if not candidates:
            return None

        listing = "\n".join(f"{n}. {q}" for n, q in enumerate(candidates, 1))
        verdicts = say(JUDGE.format(context=chunk["text"], questions=listing), chunk["id"])
        if verdicts is None:
            return None
        kept = [
            question
            for question, verdict in zip(candidates, numbered(verdicts))
            if verdict.upper().startswith("KEEP")
        ]
        if not kept:
            return None
        return min(kept, key=lambda q: lexical_overlap(q, chunk["text"]))

    rows: list[EvalRow] = []
    with ThreadPoolExecutor(max_workers=CONCURRENCY) as pool:
        for position, (chunk, question) in enumerate(zip(chunks, pool.map(ask, chunks)), 1):
            if question:
                rows.append(
                    {
                        "question": question,
                        "chunk_id": chunk["id"],
                        "book": chunk["metadata"]["book"],
                    }
                )
            if position % 25 == 0 or position == len(chunks):
                print(f"  {position}/{len(chunks)}", flush=True)

    failed = len(chunks) - len(rows)
    if failed:
        print(f"{failed} chunk(s) produced no question and were dropped", file=sys.stderr)

    overlaps = sorted(
        lexical_overlap(row["question"], by_id[row["chunk_id"]]["text"]) for row in rows
    )
    words = sorted(len(row["question"].split()) for row in rows)
    print(
        f"\nmedian overlap with the source passage {overlaps[len(overlaps) // 2]:.2f}"
        f", median length {words[len(words) // 2]} words"
    )

    with EVAL_DATASET_PATH.open("w", encoding="utf-8") as handle:
        json.dump(rows, handle, indent=2, ensure_ascii=False)
    print(f"\nwrote {len(rows)} question/chunk pairs to {EVAL_DATASET_PATH}")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))