Spaces:
Running on Zero
Running on Zero
| """The four compared methods behind one interface. | |
| ============ ================================================================== | |
| text_rag top-ranked source passages (corpus-wide BM25 or dense) | |
| graph_rag facts around entities linked from the question, serialised as triples | |
| hybrid graph facts + the source passages those facts came from (+ BM25 fill) | |
| graphtrace hybrid context + required cited path + verifier and decision policy | |
| ============ ================================================================== | |
| All methods use the same base model, the same answer prompt and output schema, and the | |
| same context token budget measured on the exact text the model sees. Text RAG also has | |
| to state its reasoning path as triples, so path validity is comparable across methods. | |
| Handles: passages are ``[C01]..``, graph facts ``{E01}..``, result records ``{R01}..``. | |
| They are resolved deterministically to chunk IDs / edge IDs / result IDs. | |
| """ | |
| from __future__ import annotations | |
| import re | |
| from collections import defaultdict | |
| from dataclasses import dataclass, field | |
| from typing import Literal | |
| from pydantic import BaseModel | |
| from rank_bm25 import BM25Okapi | |
| from .graph import KG, value_groups, values_agree | |
| from .linking import EntityLinker | |
| from .llm import LLMClient, LLMError, UsageMeter, call_structured | |
| from .paths import check_path | |
| from .prompts import ANSWER_SYSTEM, RETRY_SUFFIX, AnswerOut, answer_user | |
| from .retrieval import tokenize | |
| from .schemas import BenchItem, Chunk, Edge, MethodOutput, ResultRecord | |
| from .support import score_edge_list | |
| from .tokens import count_tokens | |
| from .verification import NLIVerifier, strip_headings | |
| class MethodConfig(BaseModel): | |
| name: str | |
| kind: Literal["text_rag", "graph_rag", "hybrid", "graphtrace"] | |
| context_budget: int = 2048 | |
| max_hops: int = 2 | |
| graph_share: float = 0.5 | |
| min_edge_support: float | None = 0.5 | |
| retry_on_invalid_path: bool = True | |
| override_confidence: float = 0.5 | |
| class Resources: | |
| base: LLMClient | |
| kg: KG | |
| linker: EntityLinker | |
| chunks: dict[str, Chunk] | |
| titles: dict[str, str] # doc_id -> title | |
| nli: NLIVerifier | None = None | |
| _bm25: BM25Okapi | None = field(default=None, repr=False) | |
| _order: list[str] = field(default_factory=list, repr=False) | |
| def bm25_rank(self, question: str) -> list[str]: | |
| if self._bm25 is None: | |
| self._order = sorted(self.chunks) | |
| self._bm25 = BM25Okapi([tokenize(self.titles.get(self.chunks[c].doc_id, "") + " " + self.chunks[c].text) or ["_"] for c in self._order]) | |
| scores = self._bm25.get_scores(tokenize(question)) | |
| return [self._order[i] for i in sorted(range(len(self._order)), key=lambda i: (-scores[i], i))] | |
| # ---- context building -------------------------------------------------------------- | |
| class Context: | |
| text: str = "" | |
| chunk_handles: dict[str, str] = field(default_factory=dict) # C01 -> chunk_id | |
| fact_handles: dict[str, list[str]] = field(default_factory=dict) # E01 -> edge ids | |
| result_handles: dict[str, str] = field(default_factory=dict) # R01 -> result id | |
| seeds: list[str] = field(default_factory=list) | |
| def tokens(self) -> int: | |
| return count_tokens(self.text) | |
| def render_chunk(handle: str, c: Chunk, res: Resources) -> str: | |
| return f"[{handle}] (paper: {res.titles.get(c.doc_id, c.doc_id)}; section: {c.section_heading})\n{strip_headings(c.text)}" | |
| def add_passages(ctx: Context, chunk_ids: list[str], res: Resources, budget: int, parts: list[str]) -> int: | |
| used = 0 | |
| for cid in chunk_ids: | |
| if cid in ctx.chunk_handles.values(): | |
| continue | |
| h = f"C{len(ctx.chunk_handles) + 1:02d}" | |
| block = render_chunk(h, res.chunks[cid], res) | |
| t = count_tokens(block) + 2 # + separator | |
| if used + t > budget: | |
| continue | |
| ctx.chunk_handles[h] = cid | |
| parts.append(block) | |
| used += t | |
| return used | |
| def graph_facts(question: str, res: Resources, hops: int) -> tuple[list[str], list[list[Edge]], list[ResultRecord]]: | |
| links = res.linker.link_question(question) | |
| seeds = [lk.node_id for lk in links] | |
| facts: dict[tuple[str, str, str], list[Edge]] = {} | |
| for e in res.kg.neighbourhood(seeds, hops): | |
| facts.setdefault((e.subject, e.relation, e.object), []).append(e) | |
| methods = [s for s in seeds if res.kg.nodes[s].type == "Method"] | |
| datasets = [s for s in seeds if res.kg.nodes[s].type == "Dataset"] | |
| results: list[ResultRecord] = [] | |
| for m in methods or [None]: | |
| for d in datasets or [None]: | |
| if m is None and d is None: | |
| continue | |
| results.extend(res.kg.results_for(method=m, dataset=d)) | |
| uniq = {r.result_id: r for r in results} | |
| return seeds, list(facts.values()), sorted(uniq.values(), key=lambda r: (r.method, r.dataset, r.metric, r.result_id)) | |
| def add_graph(ctx: Context, fact_groups: list[list[Edge]], results: list[ResultRecord], res: Resources, budget: int, parts: list[str]) -> int: | |
| used = 0 | |
| for r in results: | |
| h = f"R{len(ctx.result_handles) + 1:02d}" | |
| line = f"{{{h}}} Result: {res.kg.name(r.method)} on {res.kg.name(r.dataset)}, {res.kg.name(r.metric)} = {r.value}" | |
| t = count_tokens(line) + 1 | |
| if used + t > budget: | |
| continue | |
| ctx.result_handles[h] = r.result_id | |
| parts.append(line) | |
| used += t | |
| for edges in fact_groups: | |
| e = edges[0] | |
| h = f"E{len(ctx.fact_handles) + 1:02d}" | |
| line = f"{{{h}}} ({res.kg.name(e.subject)}) -[{e.relation}]-> ({res.kg.name(e.object)})" | |
| t = count_tokens(line) + 1 | |
| if used + t > budget: | |
| continue | |
| ctx.fact_handles[h] = [x.edge_id for x in edges] | |
| parts.append(line) | |
| used += t | |
| return used | |
| def build_context(cfg: MethodConfig, item: BenchItem, res: Resources) -> Context: | |
| ctx = Context() | |
| parts: list[str] = [] | |
| if cfg.kind == "text_rag": | |
| add_passages(ctx, res.bm25_rank(item.question), res, cfg.context_budget, parts) | |
| else: | |
| seeds, groups, results = graph_facts(item.question, res, cfg.max_hops) | |
| ctx.seeds = seeds | |
| graph_budget = cfg.context_budget if cfg.kind == "graph_rag" else int(cfg.context_budget * cfg.graph_share) | |
| graph_parts: list[str] = [] | |
| used = add_graph(ctx, groups, results, res, graph_budget - count_tokens("Graph facts:\n"), graph_parts) | |
| if graph_parts: | |
| parts.append("Graph facts:\n" + "\n".join(graph_parts)) | |
| if cfg.kind in ("hybrid", "graphtrace"): | |
| # Source passages of the facts shown, most-cited first, then BM25 to fill the budget. | |
| counts: dict[str, int] = defaultdict(int) | |
| for eids in ctx.fact_handles.values(): | |
| for eid in eids: | |
| counts[res.kg.edges[eid].source_chunk_id] += 1 | |
| for rid in ctx.result_handles.values(): | |
| counts[res.kg.results[rid].source_chunk_id] += 1 | |
| ordered = sorted(counts, key=lambda c: (-counts[c], c)) + res.bm25_rank(item.question) | |
| passage_parts: list[str] = [] | |
| # Header and separators are counted against the budget too. | |
| add_passages(ctx, ordered, res, cfg.context_budget - used - count_tokens("Graph facts:\nSource passages:\n") - 8, passage_parts) | |
| if passage_parts: | |
| parts.append("Source passages:\n" + "\n\n".join(passage_parts)) | |
| ctx.text = "\n\n".join(parts) | |
| return ctx | |
| # ---- answering + verification --------------------------------------------------------- | |
| _H = re.compile(r"([CER])\s*0*(\d+)") | |
| def norm_handle(raw: str, kind: str) -> str | None: | |
| m = _H.search(raw.upper()) | |
| return f"{m.group(1)}{int(m.group(2)):02d}" if m and m.group(1) == kind else None | |
| def resolve(out: AnswerOut, ctx: Context, res: Resources) -> tuple[list[str], list[str], list[str], list[str]]: | |
| """-> (edge ids, chunk ids, result ids, invalid handles).""" | |
| edge_ids: list[str] = [] | |
| result_ids: list[str] = [] | |
| invalid: list[str] = [] | |
| for raw in out.edge_ids: | |
| h = norm_handle(raw, "E") or norm_handle(raw, "R") | |
| if h and h in ctx.fact_handles: | |
| edge_ids.extend(ctx.fact_handles[h]) | |
| elif h and h in ctx.result_handles: | |
| result_ids.append(ctx.result_handles[h]) | |
| else: | |
| invalid.append(raw) | |
| chunk_ids: list[str] = [] | |
| for raw in out.chunk_ids: | |
| h = norm_handle(raw, "C") | |
| if h and h in ctx.chunk_handles: | |
| chunk_ids.append(ctx.chunk_handles[h]) | |
| else: | |
| invalid.append(raw) | |
| # Cited facts and results ground in the passages they were extracted from. | |
| sources = [res.kg.edges[e].source_chunk_id for e in edge_ids] + [res.kg.results[r].source_chunk_id for r in result_ids] | |
| return list(dict.fromkeys(edge_ids)), list(dict.fromkeys(chunk_ids + sources)), list(dict.fromkeys(result_ids)), invalid | |
| def detect_conflict(ctx: Context, res: Resources) -> list[ResultRecord]: | |
| """Result records *shown in this method's budgeted context* for the linked | |
| (method, dataset[, metric]) whose values disagree across passages. | |
| Only records that were placed in the context are considered, so GraphTrace cannot use | |
| graph knowledge the model was not given. The agreement rule (``value_groups``) is the | |
| same one that defines C1 gold items; this overlap is a documented threat to validity | |
| and C1/C2 are reported separately from the headline. | |
| """ | |
| kg = res.kg | |
| seeds = ctx.seeds | |
| methods = {s for s in seeds if kg.nodes[s].type == "Method"} | |
| datasets = {s for s in seeds if kg.nodes[s].type == "Dataset"} | |
| metrics = {s for s in seeds if kg.nodes[s].type == "Metric"} | |
| shown = [kg.results[rid] for rid in ctx.result_handles.values()] | |
| groups: dict[tuple[str, str, str], list[ResultRecord]] = defaultdict(list) | |
| for r in shown: | |
| if r.method in methods and r.dataset in datasets and (not metrics or r.metric in metrics) and r.value_num is not None: | |
| groups[(r.method, r.dataset, r.metric)].append(r) | |
| if not metrics and len({k[2] for k in groups}) > 1: | |
| return [] # several metrics and none named in the question: not a conflict | |
| for key in sorted(groups): | |
| rs = groups[key] | |
| if len(value_groups([r.value for r in rs])) >= 2 and len({r.source_chunk_id for r in rs}) >= 2: | |
| return rs | |
| return [] | |
| def answer_nodes(answer: list[str], res: Resources) -> set[str]: | |
| out = set() | |
| for a in answer: | |
| lk = res.linker.link_name(a) | |
| if lk: | |
| out.add(lk.node_id) | |
| return out | |
| def edge_support(edge_ids: list[str], res: Resources, cfg: MethodConfig) -> dict[str, float]: | |
| """NLI support for cited edges that were not scored at graph-build time.""" | |
| if res.nli is None: | |
| return {} | |
| todo = [res.kg.edges[e] for e in edge_ids if res.kg.edges[e].confidence is None] | |
| return score_edge_list(res.kg, todo, res.nli) | |
| def run_method(cfg: MethodConfig, item: BenchItem, res: Resources) -> MethodOutput: | |
| meter = UsageMeter() | |
| trace: dict = {} | |
| try: | |
| ctx = build_context(cfg, item, res) | |
| trace["seeds"] = ctx.seeds | |
| trace["context"] = ctx.text | |
| trace["handles"] = {"chunks": ctx.chunk_handles, "facts": ctx.fact_handles, "results": ctx.result_handles} | |
| user = answer_user(item.question, ctx.text or "(empty)") | |
| out = call_structured(res.base, meter, "answer", ANSWER_SYSTEM, user, AnswerOut, max_tokens=2048) | |
| trace["raw_answer"] = out.model_dump() | |
| decision, answer, confidence = out.decision, list(out.answer), float(min(1.0, max(0.0, out.confidence))) | |
| path = [[t.subject, t.relation, t.object] for t in out.graph_path] | |
| edge_ids, chunk_ids, result_ids, invalid = resolve(out, ctx, res) | |
| trace["invalid_handles"] = invalid | |
| if cfg.kind == "graphtrace": | |
| report: dict = {"overrides": []} | |
| conflict = detect_conflict(ctx, res) | |
| if conflict: | |
| report["conflict_results"] = [r.result_id for r in conflict] | |
| if decision != "flag_conflict": | |
| report["overrides"].append(f"{decision}->flag_conflict (conflicting result records)") | |
| decision, confidence = "flag_conflict", cfg.override_confidence | |
| answer = sorted({r.value for r in conflict}) | |
| chunk_ids = list(dict.fromkeys(chunk_ids + [r.source_chunk_id for r in conflict])) | |
| elif decision == "answer": | |
| for attempt in range(2): | |
| support = edge_support(edge_ids, res, cfg) | |
| low = {e for e, s in support.items() if cfg.min_edge_support is not None and s < cfg.min_edge_support} | |
| # A value answer is backed only by a result record that was in the context and cited. | |
| is_value = any(values_agree(res.kg.results[r].value, a) for r in result_ids for a in answer) | |
| pr = check_path(res.kg, res.linker, path, set(ctx.seeds), set() if is_value else answer_nodes(answer, res), cfg.min_edge_support) | |
| problems = pr.problems() + [f"cited fact {e} is not supported by its source" for e in sorted(low)] | |
| if not path and not is_value: | |
| problems.append("no reasoning path was given") | |
| report[f"path_check_{attempt}"] = {**pr.to_dict(), "low_support_edges": sorted(low), "problems": problems} | |
| if not problems: | |
| break | |
| if attempt == 0 and cfg.retry_on_invalid_path: | |
| out = call_structured( | |
| res.base, | |
| meter, | |
| "answer:retry", | |
| ANSWER_SYSTEM, | |
| user + RETRY_SUFFIX.format(problems="\n".join(f"- {p}" for p in problems)), | |
| AnswerOut, | |
| max_tokens=2048, | |
| ) | |
| trace["raw_answer_retry"] = out.model_dump() | |
| decision, answer, confidence = out.decision, list(out.answer), float(min(1.0, max(0.0, out.confidence))) | |
| path = [[t.subject, t.relation, t.object] for t in out.graph_path] | |
| edge_ids, chunk_ids, result_ids, invalid = resolve(out, ctx, res) | |
| if decision != "answer": | |
| break | |
| continue | |
| report["overrides"].append("answer->abstain (path failed verification)") | |
| decision, answer, confidence = "abstain", [], cfg.override_confidence | |
| break | |
| trace["verifier"] = report | |
| return MethodOutput( | |
| item_id=item.item_id, | |
| method=cfg.name, | |
| answer=answer if decision != "abstain" else [], | |
| confidence=confidence, | |
| decision=decision, | |
| graph_path=path, | |
| cited_edge_ids=edge_ids, | |
| source_chunk_ids=chunk_ids, | |
| context_tokens=ctx.tokens, | |
| trace=trace, | |
| usage=meter.summary(), | |
| ) | |
| except Exception as e: # noqa: BLE001 - one item must never abort a run; errors are scored as incorrect | |
| kind = "llm" if isinstance(e, LLMError) else "internal" | |
| return MethodOutput( | |
| item_id=item.item_id, | |
| method=cfg.name, | |
| answer=[], | |
| confidence=0.0, | |
| decision="abstain", | |
| error=f"{kind}: {type(e).__name__}: {e}", | |
| trace=trace, | |
| usage=meter.summary(), | |
| ) | |