Spaces:
Running on Zero
Running on Zero
| """Command line: ``uv run graphtrace --help``.""" | |
| from __future__ import annotations | |
| import json | |
| from pathlib import Path | |
| import typer | |
| from rich.console import Console | |
| from rich.table import Table | |
| from .config import ExperimentConfig, load_project, load_yaml, make_client, write_json | |
| app = typer.Typer(add_completion=False, no_args_is_help=True, help="GraphTrace-Bench: provenance-linked KG QA benchmark.") | |
| console = Console() | |
| CONFIG = typer.Option("configs/default.yaml", "--config", "-c", help="Project config YAML.") | |
| def _exp(path: str) -> ExperimentConfig: | |
| return ExperimentConfig(**load_yaml(path)) | |
| def corpus(config: str = CONFIG) -> None: | |
| """Select the corpus from QASPER, chunk it, write data/processed/.""" | |
| from .data.qasper import download | |
| from .pipeline import prepare_corpus | |
| cfg = load_project(config) | |
| download(cfg.paths.raw, cfg.corpus.splits) | |
| stats = prepare_corpus(cfg) | |
| console.print_json(json.dumps({k: v for k, v in stats.items() if k not in ("selection_order", "paper_datasets")}, default=str)) | |
| def demo_questions(config: str = CONFIG) -> None: | |
| """Select the demo's predefined QASPER questions (rule in graphtrace.demo) into data/processed/.""" | |
| from .data.qasper import load_split | |
| from .demo import select_demo_questions | |
| from .pipeline import load_corpus, write_jsonl | |
| cfg = load_project(config) | |
| papers, chunks = load_corpus(cfg) | |
| questions = [] | |
| for split in cfg.corpus.splits: | |
| questions.extend(load_split(cfg.paths.raw, split)[1]) | |
| rows = select_demo_questions(questions, chunks, {d: p.title for d, p in papers.items()}) | |
| write_jsonl(Path(cfg.paths.processed) / "demo_questions.jsonl", rows) | |
| console.print(f"{len(rows)} demo questions written to {Path(cfg.paths.processed) / 'demo_questions.jsonl'}") | |
| def estimate_extraction(config: str = CONFIG, output_tokens_per_chunk: int = 400) -> None: | |
| """Rough cost of KG extraction (one call per chunk) before spending anything.""" | |
| from .config import load_models | |
| from .pipeline import load_corpus | |
| from .prompts import EXTRACT_SYSTEM | |
| from .tokens import count_tokens | |
| cfg = load_project(config) | |
| _, chunks = load_corpus(cfg) | |
| spec = load_models(cfg.models_file)[cfg.extraction.model] | |
| tin = sum(c.n_tokens + count_tokens(EXTRACT_SYSTEM) + 60 for c in chunks) | |
| tout = len(chunks) * output_tokens_per_chunk | |
| usd = (tin * spec.price_in_per_mtok + tout * spec.price_out_per_mtok) / 1e6 | |
| console.print_json(json.dumps({"model": spec.model, "calls": len(chunks), "input_tokens": tin, "output_tokens_assumed": tout, "usd": round(usd, 2)})) | |
| def extract(config: str = CONFIG, limit: int = typer.Option(0, help="Only the first N chunks (0 = all).")) -> None: | |
| """Run LLM extraction over every chunk (resumable; raw calls cached in cache/llm).""" | |
| from .extraction import extract_corpus, extraction_stats | |
| from .pipeline import load_corpus, read_jsonl | |
| cfg = load_project(config) | |
| papers, chunks = load_corpus(cfg) | |
| if limit: | |
| chunks = chunks[:limit] | |
| client = make_client(cfg.extraction.model, cfg) | |
| out = Path(cfg.paths.graph) / "extractions.jsonl" | |
| summary = extract_corpus(client, papers, chunks, out, workers=cfg.extraction.workers) | |
| stats = extraction_stats(read_jsonl(out)) | |
| write_json(Path(cfg.paths.graph) / "extraction_stats.json", {**summary, **stats}) | |
| console.print_json(json.dumps({**summary, **stats}, default=str)) | |
| def build_graph(config: str = CONFIG, nli: bool = typer.Option(True, help="Score edge support with the NLI model.")) -> None: | |
| """Canonicalise entities and write nodes/edges/results JSONL.""" | |
| from .pipeline import build | |
| console.print_json(json.dumps(build(load_project(config), score_confidence=nli), default=str)) | |
| def graph_stats(config: str = CONFIG) -> None: | |
| from .graph import KG | |
| console.print_json(json.dumps(KG.load(load_project(config).paths.graph).stats(), default=str)) | |
| def generate_bench(config: str = CONFIG, paraphrase: bool = typer.Option(False, help="Paraphrase templated questions with the LLM.")) -> None: | |
| """Generate candidate items (all start unreviewed) into data/bench/candidates.jsonl.""" | |
| from collections import defaultdict | |
| from .bench import Generator, write_items | |
| from .bench import paraphrase as para | |
| from .graph import KG | |
| from .pipeline import load_corpus | |
| cfg = load_project(config) | |
| _, chunks = load_corpus(cfg) | |
| by_doc = defaultdict(list) | |
| for c in chunks: | |
| by_doc[c.doc_id].append(c) | |
| kg = KG.load(cfg.paths.graph) | |
| items, stats = Generator(kg, dict(by_doc), cfg.bench).generate() | |
| if paraphrase: | |
| if not cfg.paraphrase_model: | |
| raise typer.BadParameter("set paraphrase_model in the project config") | |
| client = make_client(cfg.paraphrase_model, cfg) | |
| rejected = 0 | |
| new = [] | |
| for it in items: | |
| names = [kg.name(s.subject) for s in it.required_path[:1]] if it.category != "single_hop" else [] | |
| q, why = para(client, it, [n for n in names if n in it.question_template]) | |
| rejected += why is not None | |
| new.append(it.model_copy(update={"question": q})) | |
| items = new | |
| stats["paraphrase_rejected"] = rejected | |
| write_items(Path(cfg.paths.bench) / "candidates.jsonl", items) | |
| write_json(Path(cfg.paths.bench) / "generation_stats.json", {**stats, "bench_config": cfg.bench.model_dump(), "total": len(items)}) | |
| console.print_json(json.dumps({**stats, "total": len(items)})) | |
| def review(config: str = CONFIG, reviewer: str = typer.Option(..., help="Your name or initials, stored with each decision."), category: str = "") -> None: | |
| """Interactively review unreviewed items against their source evidence.""" | |
| from .bench import append_review, apply_reviews, make_review, read_items, read_reviews | |
| from .graph import KG | |
| from .pipeline import load_corpus | |
| cfg = load_project(config) | |
| bench = Path(cfg.paths.bench) | |
| from .linking import EntityLinker | |
| kg = KG.load(cfg.paths.graph, cfg.bench.min_edge_confidence) | |
| _, chunks = load_corpus(cfg) | |
| ctext = {c.chunk_id: c.text for c in chunks} | |
| items = apply_reviews(read_items(bench / "candidates.jsonl"), read_reviews(bench / "reviews.jsonl"), EntityLinker(kg)) | |
| todo = [i for i in items if i.review_status == "unreviewed" and (not category or i.category == category)] | |
| console.print(f"{len(todo)} items to review. Keys: [a]ccept [r]eject [e]dit [s]kip [q]uit") | |
| for it in todo: | |
| console.rule(f"{it.item_id} ({it.category}, {it.template_id}, {it.difficulty} hop)") | |
| console.print(f"[bold]Q:[/bold] {it.question}") | |
| console.print(f"[bold]Gold decision:[/bold] {it.gold_decision} [bold]Gold answer:[/bold] {it.gold_answer}") | |
| console.print(f"[dim]{it.annotation_notes}[/dim]") | |
| for s in it.required_path: | |
| if s.edge_id in kg.edges: | |
| e = kg.edges[s.edge_id] | |
| console.print(f" ({kg.name(e.subject)}) -[{e.relation}]-> ({kg.name(e.object)}) [{e.source_chunk_id}]") | |
| console.print(f' "{ctext.get(e.source_chunk_id, "")[e.start_char : e.end_char]}"') | |
| elif s.edge_id in kg.results: | |
| r = kg.results[s.edge_id] | |
| console.print(f" result {kg.name(r.method)} / {kg.name(r.dataset)} / {kg.name(r.metric)} = {r.value} [{r.source_chunk_id}]") | |
| console.print(f' "{ctext.get(r.source_chunk_id, "")[r.start_char : r.end_char]}"') | |
| choice = typer.prompt("decision", default="s").strip().lower() | |
| if choice == "q": | |
| break | |
| if choice == "s": | |
| continue | |
| notes = typer.prompt("notes", default="") | |
| if choice == "a": | |
| append_review(bench / "reviews.jsonl", make_review(it.item_id, "accepted", reviewer, notes)) | |
| elif choice == "r": | |
| append_review(bench / "reviews.jsonl", make_review(it.item_id, "rejected", reviewer, notes)) | |
| elif choice == "e": | |
| q = typer.prompt("question", default=it.question) | |
| dec = typer.prompt("gold decision (answer/abstain/flag_conflict)", default=it.gold_decision) | |
| ans = typer.prompt("gold answer (separate with ;)", default="; ".join(it.gold_answer)) | |
| append_review( | |
| bench / "reviews.jsonl", | |
| make_review(it.item_id, "edited", reviewer, notes, question=q, gold_decision=dec, gold_answer=[a.strip() for a in ans.split(";") if a.strip()]), | |
| ) | |
| def freeze_cmd(config: str = CONFIG, version: str = "v1") -> None: | |
| """Export only human-accepted/edited items as the release benchmark.""" | |
| from .bench import freeze, read_items, read_reviews, write_items | |
| from .graph import KG | |
| from .linking import EntityLinker | |
| cfg = load_project(config) | |
| bench = Path(cfg.paths.bench) | |
| linker = EntityLinker(KG.load(cfg.paths.graph, cfg.bench.min_edge_confidence)) | |
| release, man = freeze(read_items(bench / "candidates.jsonl"), read_reviews(bench / "reviews.jsonl"), linker) | |
| write_items(bench / "release" / f"benchmark_{version}.jsonl", release) | |
| write_json(bench / "release" / f"benchmark_{version}.manifest.json", man) | |
| console.print_json(json.dumps(man)) | |
| def run(experiment: str, config: str = CONFIG, offline: bool = typer.Option(False, help="Replay cached calls only; fail on a miss.")) -> None: | |
| """Run all methods of an experiment over the benchmark (resumable).""" | |
| from .runner import run_experiment | |
| d = run_experiment(load_project(config), _exp(experiment), offline=offline) | |
| console.print(f"predictions written to {d / 'predictions.jsonl'}") | |
| def evaluate( | |
| experiment: str, | |
| config: str = CONFIG, | |
| judge: bool = typer.Option(True, "--judge/--no-judge", help="Call the evaluation judge for non-gold cited chunks (default on)."), | |
| ) -> None: | |
| from .runner import evaluate_run | |
| m = evaluate_run(load_project(config), _exp(experiment), run_judge=judge) | |
| t = Table("method", "n (headline)", "errors", "accuracy", "path valid", "citation rate", "grounding", "false-premise rej.") | |
| for name, agg in m["headline"].items(): | |
| t.add_row( | |
| name, | |
| str(agg["n"]), | |
| str(agg["errors"]), | |
| *(f"{agg[k]['mean']:.3f}" for k in ("correct", "path_all_valid", "citation_rate", "grounding_precision", "false_premise_rejected")), | |
| ) | |
| console.print(t) | |
| console.print(f"evaluation judge: {m['judge_status']}") | |
| def report(experiment: str, config: str = CONFIG, plots: bool = True) -> None: | |
| from .report import write_report | |
| cfg = load_project(config) | |
| console.print(f"wrote {write_report(Path(cfg.paths.runs) / _exp(experiment).name, with_plots=plots)}") | |
| def estimate_cost_cmd(experiment: str, config: str = CONFIG, n_items: int = 0) -> None: | |
| """Estimate the answer-stage cost of an experiment before running it.""" | |
| from .bench import read_items | |
| from .runner import estimate_cost | |
| cfg = load_project(config) | |
| exp = _exp(experiment) | |
| n = n_items or len(read_items(exp.bench_file)) | |
| console.print_json(json.dumps(estimate_cost(exp, cfg, n))) | |
| def serve(config: str = CONFIG, host: str = "127.0.0.1", port: int = 8000) -> None: | |
| """Serve the JSON API (and the demo UI if built).""" | |
| import uvicorn | |
| from .app.server import create_app | |
| uvicorn.run(create_app(load_project(config)), host=host, port=port) | |
| if __name__ == "__main__": | |
| app() | |