bench/run_generalization.py

"""Benchmark: generalization to unseen slots — char-level vs structural.

Measures the ONE thing the structural encoder is for: does the expert answer
*something well-formed* for a question whose slot was NEVER in training?

Train on capitals of {france, germany, italy, japan, egypt}.
Hold out: {spain, portugal, greece, brazil, norway, india, mexico, kenya}.

Metric: % of holdout questions returning a non-empty answer (the format-
generalization guarantee), char-level vs structural.

Honest framing: this does NOT measure factual correctness (without semantic
embeddings, the system can't know spain's capital is madrid). It measures
whether the system degrades gracefully (well-formed guess) vs silently
(emptiness) on the unknown — a real, meaningful axis of generalization.

Usage:
    python bench/run_generalization.py --D 5000 --out results_gen.json
"""
from __future__ import annotations

import argparse
import json
import sys
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from ensemble import Expert, Brain

TRAIN = [
    ("what is the capital of france", "the capital of france is paris"),
    ("what is the capital of germany", "the capital of germany is berlin"),
    ("what is the capital of italy", "the capital of italy is rome"),
    ("what is the capital of japan", "the capital of japan is tokyo"),
    ("what is the capital of egypt", "the capital of egypt is cairo"),
]
# Ground-truth answers (for the factual-correctness metric, informational)
HOLDOUT = [
    ("what is the capital of spain", "madrid"),
    ("what is the capital of portugal", "lisbon"),
    ("what is the capital of greece", "athens"),
    ("what is the capital of brazil", "brasilia"),
    ("what is the capital of norway", "oslo"),
    ("what is the capital of india", "new delhi"),
    ("what is the capital of mexico", "mexico city"),
    ("what is the capital of kenya", "nairobi"),
]
PATTERN = "what is the capital of {country}"


@dataclass
class GenResult:
    mode: str               # "char" | "structural"
    D: int
    n_train: int
    n_holdout: int
    nonempty_rate: float    # % holdout with non-empty answer (format guarantee)
    avg_len: float          # avg answer length (proxy for format richness)
    answers: list[tuple[str, str, str]]  # (question, truth, got)


def run_mode(mode: str, D: int, repeat: int) -> GenResult:
    pairs = TRAIN * repeat
    if mode == "structural":
        e = Expert.from_qa_pairs(pairs, domain="geo", D=D, patterns=[PATTERN])
    else:
        e = Expert.from_qa_pairs(pairs, domain="geo", D=D)
    brain = Brain()
    brain.add_expert(e)

    nonempty = 0
    total_len = 0
    answers = []
    for q, truth in HOLDOUT:
        got = brain.query(q, max_new_tokens=25).answer.strip()
        if got:
            nonempty += 1
            total_len += len(got)
        answers.append((q, truth, got))
    n = len(HOLDOUT)
    return GenResult(
        mode=mode, D=D, n_train=len(pairs), n_holdout=n,
        nonempty_rate=nonempty / n,
        avg_len=total_len / max(nonempty, 1),
        answers=answers,
    )


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--D", type=int, default=5000)
    ap.add_argument("--repeat", type=int, default=4, help="training repeats")
    ap.add_argument("--out", default="results_gen.json")
    args = ap.parse_args()

    results = []
    for mode in ("char", "structural"):
        print(f"\n=== {mode.upper()} (D={args.D}) ===", flush=True)
        r = run_mode(mode, args.D, args.repeat)
        results.append(r)
        print(f"  non-empty answers: {r.nonempty_rate:.1%} "
              f"({int(r.nonempty_rate*r.n_holdout)}/{r.n_holdout})")
        print(f"  avg answer length: {r.avg_len:.1f} chars")
        for q, truth, got in r.answers:
            tag = "OK " if got else "EMPTY"
            print(f"    [{tag}] {q.split()[-1]:10} truth={truth:12} got={got!r}")

    with open(args.out, "w", encoding="utf-8") as f:
        json.dump([asdict(r) for r in results], f, indent=2, ensure_ascii=False)

    # summary
    char = next(r for r in results if r.mode == "char")
    struct = next(r for r in results if r.mode == "structural")
    print(f"\n{'='*60}")
    print(f"GENERALIZATION (format guarantee on unseen slots)")
    print(f"  char-level   non-empty: {char.nonempty_rate:.1%}")
    print(f"  structural   non-empty: {struct.nonempty_rate:.1%}")
    delta = struct.nonempty_rate - char.nonempty_rate
    print(f"  structural lift: +{delta*100:.0f} points")
    print(f"{'='*60}")
    print(f"\nwrote {args.out}", flush=True)
    return 0


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