Spaces:
Running on Zero
Running on Zero
| """Experiment runner: predictions (resumable), evaluation, failure examples. | |
| Layout of ``runs/<name>/``:: | |
| manifest.json config snapshot, prompt version, packages, git commit, model specs | |
| predictions.jsonl one MethodOutput per (method, item), with full trace | |
| judgments.jsonl evaluation-judge labels for non-gold cited chunks (optional) | |
| per_item.csv every metric for every (method, item) | |
| metrics.json aggregates with bootstrap CIs, per-category tables, paired tests, calibration | |
| failures.jsonl automatically selected failure examples per method | |
| """ | |
| from __future__ import annotations | |
| import csv | |
| import hashlib | |
| import json | |
| import math | |
| import threading | |
| from collections import defaultdict | |
| from concurrent.futures import ThreadPoolExecutor | |
| from pathlib import Path | |
| from .bench import read_items | |
| from .config import ExperimentConfig, ProjectConfig, load_models, make_client, manifest, write_json | |
| from .evaluation.metrics import by_key, calibration, item_metrics | |
| from .evaluation.stats import bootstrap_ci, paired_test | |
| from .graph import KG | |
| from .linking import EntityLinker | |
| from .llm import ModelSpec | |
| from .methods import MethodConfig, Resources, run_method | |
| from .pipeline import load_corpus, read_jsonl | |
| from .prompts import PROMPT_VERSION | |
| from .schemas import BenchItem, MethodOutput | |
| METRICS = [ | |
| "correct", | |
| "decision_correct", | |
| "answer_f1", | |
| "answer_exact", | |
| "path_valid_fraction", | |
| "path_all_valid", | |
| "path_completeness", | |
| "citation_rate", | |
| "grounding_precision", | |
| "false_premise_rejected", | |
| "conflict_flagged", | |
| "insufficient_abstained", | |
| "answered", | |
| "errored", | |
| "context_tokens", | |
| "calls", | |
| "input_tokens", | |
| "output_tokens", | |
| "cost_usd", | |
| "llm_latency_s", | |
| ] | |
| # Headline = categories whose gold does not rely on the result-record agreement rule that the | |
| # GraphTrace verifier also uses. C1/C2 are reported separately (protocol, threats to validity). | |
| HEADLINE_CATEGORIES = ("single_hop", "multi_hop", "false_premise") | |
| def method_hash(method: MethodConfig, base_spec: ModelSpec, project: ProjectConfig, exp: ExperimentConfig) -> str: | |
| """Identity of everything that determines a method's output for an item. Predictions | |
| with a different hash are stale and are re-run.""" | |
| payload = { | |
| "method": method.model_dump(), | |
| "base_model": base_spec.model_dump(), | |
| "prompt_version": PROMPT_VERSION, | |
| "nli": exp.nli_model if exp.nli_verifier else None, | |
| "min_edge_confidence": project.bench.min_edge_confidence, | |
| } | |
| return hashlib.sha256(json.dumps(payload, sort_keys=True, default=str).encode("utf-8")).hexdigest()[:16] | |
| def load_bench(exp: ExperimentConfig) -> list[BenchItem]: | |
| items = read_items(exp.bench_file) | |
| unreviewed = [i for i in items if i.review_status not in ("accepted", "edited")] | |
| if unreviewed and not exp.allow_unreviewed: | |
| raise ValueError( | |
| f"{len(unreviewed)} items in {exp.bench_file} are not human-accepted. Freeze the benchmark first " | |
| "(`graphtrace freeze`) or set allow_unreviewed: true for a smoke test that must not be reported." | |
| ) | |
| return items[: exp.limit] if exp.limit else items | |
| def resources(project: ProjectConfig, exp: ExperimentConfig, offline: bool = False) -> Resources: | |
| papers, chunks = load_corpus(project) | |
| kg = KG.load(project.paths.graph, project.bench.min_edge_confidence) | |
| nli = None | |
| if exp.nli_verifier: | |
| from .verification import NLIVerifier | |
| nli = NLIVerifier(exp.nli_model) | |
| return Resources( | |
| base=make_client(exp.base_model, project, offline=offline), | |
| kg=kg, | |
| linker=EntityLinker(kg), | |
| chunks={c.chunk_id: c for c in chunks}, | |
| titles={d: p.title for d, p in papers.items()}, | |
| nli=nli, | |
| ) | |
| def run_experiment(project: ProjectConfig, exp: ExperimentConfig, offline: bool = False, res: Resources | None = None) -> Path: | |
| run_dir = Path(project.paths.runs) / exp.name | |
| run_dir.mkdir(parents=True, exist_ok=True) | |
| items = load_bench(exp) | |
| res = res or resources(project, exp, offline) | |
| models = load_models(project.models_file) if Path(project.models_file).exists() else {} | |
| hashes = {m.name: method_hash(m, res.base.spec, project, exp) for m in exp.methods} | |
| write_json( | |
| run_dir / "manifest.json", | |
| manifest( | |
| { | |
| "experiment": exp.model_dump(), | |
| "model_specs": {k: v.model_dump() for k, v in models.items() if k in (exp.base_model, exp.judge_model)}, | |
| "n_items": len(items), | |
| "method_config_hashes": hashes, | |
| "smoke_test_unreviewed_items": exp.allow_unreviewed, | |
| } | |
| ), | |
| ) | |
| pred_path = run_dir / "predictions.jsonl" | |
| # Completed = same method name AND same config hash AND no error. | |
| done = {(r["method"], r["item_id"], r.get("config_hash")) for r in read_jsonl(pred_path) if not r.get("error")} | |
| todo = [(m, it) for m in exp.methods for it in items if (m.name, it.item_id, hashes[m.name]) not in done] | |
| lock = threading.Lock() | |
| def one(t: tuple[MethodConfig, BenchItem]) -> MethodOutput: | |
| return run_method(t[0], t[1], res).model_copy(update={"config_hash": hashes[t[0].name]}) | |
| with pred_path.open("a", encoding="utf-8") as f, ThreadPoolExecutor(max_workers=exp.workers) as pool: | |
| for out in pool.map(one, todo): | |
| with lock: | |
| f.write(out.model_dump_json() + "\n") | |
| f.flush() | |
| return run_dir | |
| def latest_predictions(run_dir: Path, hashes: dict[str, str] | None = None) -> list[MethodOutput]: | |
| """Latest prediction per (method, item); with ``hashes``, only rows of the current config.""" | |
| latest: dict[tuple[str, str], dict] = {} | |
| for r in read_jsonl(run_dir / "predictions.jsonl"): | |
| if hashes is not None and r.get("config_hash") != hashes.get(r["method"]): | |
| continue | |
| latest[(r["method"], r["item_id"])] = r # later lines (successful re-runs) win | |
| return [MethodOutput(**r) for r in latest.values()] | |
| def _agg(rows: list[dict], seed: int = 0) -> dict: | |
| out = {"n": len(rows), "errors": sum(1 for r in rows if r.get("errored"))} | |
| for m in METRICS: | |
| vals = [r[m] for r in rows if r.get(m) is not None and not (isinstance(r[m], float) and math.isnan(r[m]))] | |
| mean, lo, hi = bootstrap_ci(vals, seed=seed) | |
| out[m] = {"mean": mean, "low": lo, "high": hi, "n": len(vals)} | |
| return out | |
| def evaluate_run(project: ProjectConfig, exp: ExperimentConfig, run_judge: bool = True, offline: bool = False) -> dict: | |
| run_dir = Path(project.paths.runs) / exp.name | |
| items = {i.item_id: i for i in load_bench(exp)} | |
| kg = KG.load(project.paths.graph, project.bench.min_edge_confidence) | |
| linker = EntityLinker(kg) | |
| man_path = run_dir / "manifest.json" | |
| hashes = json.loads(man_path.read_text(encoding="utf-8")).get("method_config_hashes") if man_path.exists() else None | |
| outputs = [o for o in latest_predictions(run_dir, hashes) if o.item_id in items] | |
| judged: dict[tuple[str, str], set[str]] = defaultdict(set) | |
| jpath = run_dir / "judgments.jsonl" | |
| reused = "; reused existing judgments.jsonl" if jpath.exists() else "" | |
| if not run_judge: | |
| judge_status = "skipped (--no-judge)" + reused | |
| elif not exp.judge_model: | |
| judge_status = "skipped (no judge_model configured)" + reused | |
| else: | |
| judge_status = f"ran ({exp.judge_model})" | |
| if run_judge and exp.judge_model: | |
| from .evaluation.judge import LLMJudge, judge_outputs | |
| _, chunks = load_corpus(project) | |
| rows = judge_outputs(outputs, items, {c.chunk_id: c for c in chunks}, LLMJudge(make_client(exp.judge_model, project, offline=offline))) | |
| with jpath.open("w", encoding="utf-8") as f: | |
| for r in rows: | |
| f.write(json.dumps(r, ensure_ascii=False) + "\n") | |
| for r in read_jsonl(jpath): | |
| if r.get("label") == "supported": | |
| judged[(r["method"], r["item_id"])].add(r["chunk_id"]) | |
| rows = [item_metrics(o, items[o.item_id], kg, linker, judged.get((o.method, o.item_id)), exp.min_support) for o in outputs] | |
| rows.sort(key=lambda r: (r["method"], r["item_id"])) | |
| with (run_dir / "per_item.csv").open("w", newline="", encoding="utf-8") as f: | |
| w = csv.DictWriter(f, fieldnames=list(rows[0].keys()) if rows else ["item_id"]) | |
| w.writeheader() | |
| w.writerows(rows) | |
| per_method = by_key(rows, "method") | |
| metrics: dict = { | |
| "methods": {}, | |
| "headline": {}, | |
| "conflict_insufficient": {}, | |
| "by_category": {}, | |
| "calibration": {}, | |
| "paired_vs_reference": {}, | |
| "headline_categories": list(HEADLINE_CATEGORIES), | |
| "judge_used": bool(jpath.exists()), | |
| "judge_status": judge_status, | |
| } | |
| for m, rs in per_method.items(): | |
| metrics["methods"][m] = _agg(rs) | |
| metrics["headline"][m] = _agg([r for r in rs if r["category"] in HEADLINE_CATEGORIES]) | |
| metrics["conflict_insufficient"][m] = _agg([r for r in rs if r["category"] not in HEADLINE_CATEGORIES]) | |
| metrics["by_category"][m] = {c: _agg(crs) for c, crs in by_key(rs, "category").items()} | |
| metrics["calibration"][m] = calibration(rs) | |
| ref = per_method.get(exp.reference_method, []) | |
| ref_by_item = {r["item_id"]: r for r in ref} | |
| for m, rs in per_method.items(): | |
| if m == exp.reference_method: | |
| continue | |
| # Paired tests on the headline categories only. | |
| pair = [(r, ref_by_item[r["item_id"]]) for r in rs if r["item_id"] in ref_by_item and r["category"] in HEADLINE_CATEGORIES] | |
| metrics["paired_vs_reference"][m] = { | |
| k: paired_test([a[k] for a, _ in pair], [b[k] for _, b in pair]) | |
| for k in ( | |
| "correct", | |
| "path_valid_fraction", | |
| "path_all_valid", | |
| "citation_rate", | |
| "grounding_precision", | |
| "false_premise_rejected", | |
| "answer_f1", | |
| ) | |
| } | |
| write_json(run_dir / "metrics.json", metrics) | |
| write_failures(run_dir, outputs, items, rows) | |
| return metrics | |
| def write_failures(run_dir: Path, outputs: list[MethodOutput], items: dict[str, BenchItem], rows: list[dict], per_kind: int = 5) -> None: | |
| """Pick concrete failure cases per method, deterministically (sorted by item ID).""" | |
| out_by = {(o.method, o.item_id): o for o in outputs} | |
| kinds = { | |
| "accepted_false_premise": lambda r: r["category"] == "false_premise" and r["decision"] == "answer", | |
| "missed_conflict": lambda r: r["template_id"] == "C1" and r["decision"] != "flag_conflict", | |
| "wrong_answer_with_invalid_path": lambda r: r["gold_decision"] == "answer" and r["correct"] == 0 and r["path_all_valid"] == 0, | |
| "correct_answer_invalid_path": lambda r: r["correct"] == 1 and r["path_all_valid"] == 0, | |
| "over_abstention": lambda r: r["gold_decision"] == "answer" and r["decision"] == "abstain", | |
| "error": lambda r: bool(r["error"]), | |
| } | |
| with (run_dir / "failures.jsonl").open("w", encoding="utf-8") as f: | |
| for method, rs in sorted(by_key(rows, "method").items()): | |
| for kind, pred in kinds.items(): | |
| for r in [r for r in rs if pred(r)][:per_kind]: | |
| o = out_by[(method, r["item_id"])] | |
| it = items[r["item_id"]] | |
| f.write( | |
| json.dumps( | |
| { | |
| "method": method, | |
| "kind": kind, | |
| "item_id": it.item_id, | |
| "question": it.question, | |
| "gold_decision": it.gold_decision, | |
| "gold_answer": it.gold_answer, | |
| "decision": o.decision, | |
| "answer": o.answer, | |
| "confidence": o.confidence, | |
| "graph_path": o.graph_path, | |
| "source_chunk_ids": o.source_chunk_ids, | |
| "verifier": o.trace.get("verifier"), | |
| }, | |
| ensure_ascii=False, | |
| ) | |
| + "\n" | |
| ) | |
| def estimate_cost(exp: ExperimentConfig, project: ProjectConfig, n_items: int, avg_output_tokens: int = 350) -> dict: | |
| """Upper-bound style estimate: one answer call per (method, item) using the full budget, | |
| plus a retry for half of GraphTrace's items, plus ~600 tokens of instructions per call.""" | |
| spec = load_models(project.models_file)[exp.base_model] | |
| calls = 0 | |
| tokens_in = 0 | |
| for m in exp.methods: | |
| per = 1.5 if m.kind == "graphtrace" else 1.0 | |
| calls += int(n_items * per) | |
| tokens_in += int(n_items * per * (m.context_budget + 600)) | |
| tokens_out = calls * avg_output_tokens | |
| cost = (tokens_in * spec.price_in_per_mtok + tokens_out * spec.price_out_per_mtok) / 1e6 | |
| return {"model": spec.model, "calls": calls, "input_tokens": tokens_in, "output_tokens": tokens_out, "usd": round(cost, 2)} | |