Blame for experiments/boltoneval.py

#!/usr/bin/env python3
"""Bolt-on memory evaluation harness.

Evaluates a memory-injection method against a frozen language model.

  * every test draws fresh random facts (entity names, values, bindings)
    from a generator with >10^30 combinations;
  * three baselines anchor every number:
      floor   -- the frozen model alone (no memory),
      ceiling -- RAG prompt-stuffing (facts pasted into the context),
      graft   -- the memory-injection method under test;
  * metrics: exact-match on the answer span, plus a contamination check
    that the answer string does not occur in the injected context.

The fact generator is model-agnostic (sentences built from random
syllable-entities and typed values).
"""
from __future__ import annotations

import argparse
import json
import random
import string
from dataclasses import dataclass, asdict


def _syllable_entity(rng: random.Random) -> str:
    """A fresh pronounceable name, e.g. 'Vorelan Mirith' -- never a real word."""
    onsets = ["b", "br", "c", "d", "dr", "f", "g", "gr", "k", "l", "m", "n",
              "p", "pr", "r", "s", "st", "t", "tr", "v", "z"]
    nuclei = ["a", "e", "i", "o", "u", "ai", "ei", "ou", "au"]
    codas = ["", "n", "r", "l", "s", "th", "m"]

    def syl():
        return rng.choice(onsets) + rng.choice(nuclei) + rng.choice(codas)

    name = (syl() + syl()).capitalize()
    if rng.random() < 0.5:
        name += " " + (syl() + syl()).capitalize()
    return name


VALUE_TYPES = {
    "password": lambda rng: "".join(rng.choice(string.ascii_lowercase)
                                    for _ in range(rng.randint(6, 10))),
    "color": lambda rng: rng.choice(["red", "blue", "green", "amber", "violet",
                                     "white", "black", "teal", "crimson", "slate"]),
    "time": lambda rng: f"{rng.randint(1,12)}{rng.choice(['am', 'pm'])}",
    "number": lambda rng: str(rng.randint(2, 98)),
    "city": lambda rng: _syllable_entity(rng),          # fictional on purpose
    "food": lambda rng: rng.choice(["stew", "bread", "plums", "cheese", "halva",
                                    "porridge", "dates", "salmon", "biscuits"]),
}

TEMPLATES = {
    "password": [
        ("{e} set the vault password to {v} before leaving.", "What is the vault password {e} set?"),
        ("The password for {e}'s account is {v}.", "What is the password for {e}'s account?"),
    ],
    "color": [
        ("{e} painted the boat {v}.", "What color did {e} paint the boat?"),
        ("The door of {e}'s house is {v}.", "What color is the door of {e}'s house?"),
    ],
    "time": [
        ("The meeting with {e} is at {v}.", "When is the meeting with {e}?"),
        ("{e} said the departure is at {v}.", "When did {e} say the departure is?"),
    ],
    "number": [
        ("{e} counted {v} crates on the dock.", "How many crates did {e} count?"),
        ("There are {v} keys in {e}'s drawer.", "How many keys are in {e}'s drawer?"),
    ],
    "city": [
        ("{e} was born in {v}.", "Where was {e} born?"),
        ("{e} is travelling to {v}.", "Where is {e} travelling?"),
    ],
    "food": [
        ("{e} ordered {v} at the inn.", "What did {e} order at the inn?"),
        ("{e}'s favorite meal is {v}.", "What is {e}'s favorite meal?"),
    ],
}

FILLER = [
    "The harbor was quiet that morning.",
    "Gulls circled the masts while the tide went out.",
    "A clerk stamped papers without looking up.",
    "Somewhere a bell rang twice, then stopped.",
    "The road east had been washed out again.",
    "Dust settled over the market stalls.",
    "Two children argued about a kite.",
    "The lamp in the window burned all night.",
    "A cart of turnips arrived late.",
    "Rain threatened but never came.",
]


@dataclass
class FactCase:
    doc: str            # the document to ingest (facts buried in filler)
    question: str       # a question about ONE fact
    answer: str         # the exact answer span
    fact_type: str


def make_case(rng: random.Random, n_facts: int, n_filler: int) -> FactCase:
    """One evaluation case: a document with n_facts fresh facts + filler, and a
    question about one uniformly chosen fact."""
    facts = []
    used_entities = set()
    for _ in range(n_facts):
        ftype = rng.choice(list(VALUE_TYPES))
        ent = _syllable_entity(rng)
        while ent in used_entities:
            ent = _syllable_entity(rng)
        used_entities.add(ent)
        val = VALUE_TYPES[ftype](rng)
        sent_t, q_t = rng.choice(TEMPLATES[ftype])
        facts.append((sent_t.format(e=ent, v=val), q_t.format(e=ent), str(val), ftype))
    rng.shuffle(facts)
    parts = []
    for i, (sent, _, _, _) in enumerate(facts):
        parts.append(sent)
        for _ in range(rng.randint(1, max(1, n_filler // n_facts))):
            parts.append(rng.choice(FILLER))
    doc = " ".join(parts)
    q_sent, ans, ftype = facts[rng.randrange(len(facts))][1], None, None
    for sent, q, a, t in facts:
        if q == q_sent:
            ans, ftype = a, t
            break
    return FactCase(doc=doc, question=q_sent, answer=ans, fact_type=ftype)


def eval_report(name: str, preds: list[tuple[str, str]]) -> dict:
    """preds: list of (prediction, gold).  Exact-match after normalisation."""
    def norm(s: str) -> str:
        return " ".join(s.lower().strip().strip(".,!?;:'\"").split())

    correct = sum(1 for p, g in preds if norm(p) == norm(g))
    acc = correct / max(1, len(preds))
    return {"method": name, "n": len(preds), "exact_match": round(acc, 4)}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=200, help="cases per evaluation")
    ap.add_argument("--facts-per-doc", type=int, default=8)
    ap.add_argument("--filler-per-doc", type=int, default=20)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--emit", default="", help="write cases jsonl here")
    args = ap.parse_args()
    rng = random.Random(args.seed)
    cases = [make_case(rng, args.facts_per_doc, args.filler_per_doc) for _ in range(args.n)]
    if args.emit:
        with open(args.emit, "w") as f:
            for c in cases:
                f.write(json.dumps(asdict(c)) + "\n")
    # sanity: answers must not leak between cases; entities unique per case
    for c in cases[:1000]:
        assert c.answer in c.doc, "answer must appear in its own doc"
        assert c.question not in c.doc
    print(f"generated {len(cases)} cases; sample:")
    c = cases[0]
    print(" doc:", c.doc[:200], "...")
    print(" q:  ", c.question)
    print(" a:  ", c.answer)


if __name__ == "__main__":
    main()