Spaces:
Running on Zero
Running on Zero
| """Provenance-linked knowledge graph: build from extraction records, persist as JSONL, | |
| load into a NetworkX ``MultiDiGraph``. | |
| One edge per (subject, relation, object, source chunk, quote span): the same fact | |
| stated in two places gives two edges, each with its own evidence. ``KG.has_fact`` asks | |
| whether a fact holds regardless of source. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import re | |
| from collections import Counter, defaultdict | |
| from pathlib import Path | |
| import networkx as nx | |
| from .canonical import Canonicalizer | |
| from .prompts import THIS_PAPER | |
| from .schemas import Document, Edge, Node, ResultRecord | |
| def _hid(prefix: str, *parts: object) -> str: | |
| return prefix + hashlib.sha1("|".join(map(str, parts)).encode("utf-8")).hexdigest()[:12] | |
| def paper_node_id(doc_id: str) -> str: | |
| return f"paper:{doc_id}" | |
| def parse_number(value: str) -> float | None: | |
| m = re.search(r"-?\d+(?:\.\d+)?", value.replace(",", "")) | |
| return float(m.group(0)) if m else None | |
| def _decimals(value: str) -> int: | |
| m = re.search(r"-?\d+(?:\.(\d+))?", value.replace(",", "")) | |
| return len(m.group(1)) if m and m.group(1) else 0 | |
| def values_agree(a: str, b: str) -> bool: | |
| """Same number up to scale (x, 100x, x/100: fraction vs percentage) after rounding both | |
| to the coarser reported precision. "91.2" == "0.912" == "91.20" == "91.23"; "91.2" != "90.4".""" | |
| na, nb = parse_number(a), parse_number(b) | |
| if na is None or nb is None: | |
| return a.strip().lower() == b.strip().lower() | |
| da, db = _decimals(a), _decimals(b) | |
| for x, dx in ((na, da), (na * 100, max(da - 2, 0)), (na / 100, da + 2)): | |
| p = min(dx, db) | |
| if round(x, p) == round(nb, p): | |
| return True | |
| return False | |
| def value_groups(values: list[str]) -> list[list[str]]: | |
| """Greedy clustering of reported values into groups that agree under ``values_agree``.""" | |
| groups: list[list[str]] = [] | |
| for v in values: | |
| for g in groups: | |
| if values_agree(g[0], v): | |
| g.append(v) | |
| break | |
| else: | |
| groups.append([v]) | |
| return groups | |
| def build_graph( | |
| records: list[dict], | |
| docs: dict[str, Document], | |
| canon: Canonicalizer, | |
| confidence: dict[str, float] | None = None, | |
| stats: dict | None = None, | |
| ) -> tuple[list[Node], list[Edge], list[ResultRecord]]: | |
| """``confidence`` optionally maps edge_id -> NLI support score; ``stats`` receives counters.""" | |
| surface: dict[str, Counter] = defaultdict(Counter) | |
| ntype: dict[str, str] = {} | |
| nodes: dict[str, Node] = {} | |
| for doc_id in sorted({r["doc_id"] for r in records}): | |
| pid = paper_node_id(doc_id) | |
| nodes[pid] = Node(node_id=pid, type="Paper", name=docs[doc_id].title, doc_id=doc_id) | |
| def ent(name: str, etype: str, doc_id: str) -> str: | |
| if etype == "Paper": | |
| return paper_node_id(doc_id) | |
| nid = canon.node_id(name, etype) | |
| surface[nid][name] += 1 | |
| ntype[nid] = etype | |
| return nid | |
| # A result whose method is THIS_PAPER means the paper's own method: resolve it only when | |
| # the paper proposes exactly one method; otherwise drop the record and count it. | |
| proposed_by_doc: dict[str, set[str]] = defaultdict(set) | |
| for rec in records: | |
| if rec.get("error"): | |
| continue | |
| for t in rec["triples"]: | |
| if t["relation"] == "proposes" and t["subject_type"] == "Paper": | |
| proposed_by_doc[rec["doc_id"]].add(canon.node_id(t["object"], "Method")) | |
| dropped_this_paper = 0 | |
| edges: dict[str, Edge] = {} | |
| results: dict[str, ResultRecord] = {} | |
| for rec in records: | |
| if rec.get("error"): | |
| continue | |
| for t in rec["triples"]: | |
| s = ent(t["subject"], t["subject_type"], rec["doc_id"]) | |
| o = ent(t["object"], t["object_type"], rec["doc_id"]) | |
| if s == o: | |
| continue | |
| eid = _hid("e_", s, t["relation"], o, rec["chunk_id"], t["start_char"], t["end_char"]) | |
| edges[eid] = Edge( | |
| edge_id=eid, | |
| subject=s, | |
| relation=t["relation"], | |
| object=o, | |
| source_chunk_id=rec["chunk_id"], | |
| quote=t["quote"], | |
| start_char=t["start_char"], | |
| end_char=t["end_char"], | |
| confidence=(confidence or {}).get(eid), | |
| extractor_model=rec["extractor_model"], | |
| ) | |
| for r in rec["results"]: | |
| if r["method"].strip() == THIS_PAPER: | |
| own = proposed_by_doc.get(rec["doc_id"], set()) | |
| if len(own) != 1: | |
| dropped_this_paper += 1 | |
| continue | |
| (m,) = own | |
| else: | |
| m = ent(r["method"], "Method", rec["doc_id"]) | |
| d = ent(r["dataset"], "Dataset", rec["doc_id"]) | |
| k = ent(r["metric"], "Metric", rec["doc_id"]) | |
| rid = _hid("r_", m, d, k, r["value"], rec["chunk_id"], r["start_char"]) | |
| results[rid] = ResultRecord( | |
| result_id=rid, | |
| method=m, | |
| dataset=d, | |
| metric=k, | |
| value=r["value"], | |
| value_num=parse_number(r["value"]), | |
| source_chunk_id=rec["chunk_id"], | |
| quote=r["quote"], | |
| start_char=r["start_char"], | |
| end_char=r["end_char"], | |
| extractor_model=rec["extractor_model"], | |
| ) | |
| if stats is not None: | |
| stats["results_dropped_this_paper"] = dropped_this_paper | |
| for nid, forms in surface.items(): | |
| name = min(forms.items(), key=lambda kv: (-kv[1], kv[0]))[0] | |
| nodes[nid] = Node(node_id=nid, type=ntype[nid], name=name, aliases=sorted(forms), mentions=sum(forms.values())) | |
| # Aliases from the canonicaliser make the linker find abbreviations and long forms. | |
| for m in canon.merges: | |
| for nid, node in nodes.items(): | |
| if node.type != "Paper" and nid.split(":", 1)[1] == m["canonical"] and m["alias"] not in node.aliases: | |
| node.aliases.append(m["alias"]) | |
| return sorted(nodes.values(), key=lambda n: n.node_id), sorted(edges.values(), key=lambda e: e.edge_id), sorted(results.values(), key=lambda r: r.result_id) | |
| def save_graph(out_dir: str | Path, nodes: list[Node], edges: list[Edge], results: list[ResultRecord], merges: list[dict]) -> None: | |
| out = Path(out_dir) | |
| out.mkdir(parents=True, exist_ok=True) | |
| for name, rows in (("nodes", nodes), ("edges", edges), ("results", results)): | |
| with (out / f"{name}.jsonl").open("w", encoding="utf-8") as f: | |
| for r in rows: | |
| f.write(r.model_dump_json() + "\n") | |
| with (out / "alias_merges.jsonl").open("w", encoding="utf-8") as f: | |
| for m in merges: | |
| f.write(json.dumps(m, ensure_ascii=False) + "\n") | |
| def _read_jsonl(p: Path) -> list[dict]: | |
| if not p.exists(): | |
| return [] | |
| return [json.loads(line) for line in p.read_text(encoding="utf-8").splitlines() if line.strip()] | |
| class KG: | |
| def __init__(self, nodes: list[Node], edges: list[Edge], results: list[ResultRecord]): | |
| self.nodes = {n.node_id: n for n in nodes} | |
| self.edges = {e.edge_id: e for e in edges} | |
| self.results = {r.result_id: r for r in results} | |
| self.g = nx.MultiDiGraph() | |
| for n in nodes: | |
| self.g.add_node(n.node_id, type=n.type, name=n.name) | |
| for e in edges: | |
| self.g.add_edge(e.subject, e.object, key=e.edge_id, relation=e.relation) | |
| self._facts: dict[tuple[str, str, str], list[str]] = defaultdict(list) | |
| for e in edges: | |
| self._facts[(e.subject, e.relation, e.object)].append(e.edge_id) | |
| def filtered(self, min_confidence: float | None) -> KG: | |
| """The graph without edges whose NLI support is below ``min_confidence``. | |
| One filtered view is used for benchmark gold *and* for every method's retrieval and | |
| for path checks, so gold and methods always see the same edge set. Unscored edges | |
| (confidence None) are kept; low-support edges stay in edges.jsonl for audit. | |
| """ | |
| if min_confidence is None: | |
| return self | |
| keep = [e for e in self.edges.values() if e.confidence is None or e.confidence >= min_confidence] | |
| return KG(list(self.nodes.values()), keep, list(self.results.values())) | |
| def load(cls, graph_dir: str | Path, min_confidence: float | None = None) -> KG: | |
| d = Path(graph_dir) | |
| return cls( | |
| [Node(**r) for r in _read_jsonl(d / "nodes.jsonl")], | |
| [Edge(**r) for r in _read_jsonl(d / "edges.jsonl")], | |
| [ResultRecord(**r) for r in _read_jsonl(d / "results.jsonl")], | |
| ).filtered(min_confidence) | |
| # ---- queries ---- | |
| def name(self, node_id: str) -> str: | |
| n = self.nodes.get(node_id) | |
| return n.name if n else node_id | |
| def has_fact(self, s: str, r: str, o: str) -> bool: | |
| return (s, r, o) in self._facts | |
| def fact_edges(self, s: str, r: str, o: str) -> list[Edge]: | |
| return [self.edges[i] for i in self._facts.get((s, r, o), [])] | |
| def out_edges(self, node_id: str, relation: str | None = None) -> list[Edge]: | |
| if node_id not in self.g: | |
| return [] | |
| return [self.edges[k] for _, _, k, d in self.g.out_edges(node_id, keys=True, data=True) if relation is None or d["relation"] == relation] | |
| def in_edges(self, node_id: str, relation: str | None = None) -> list[Edge]: | |
| if node_id not in self.g: | |
| return [] | |
| return [self.edges[k] for _, _, k, d in self.g.in_edges(node_id, keys=True, data=True) if relation is None or d["relation"] == relation] | |
| def objects(self, s: str, relation: str) -> set[str]: | |
| return {e.object for e in self.out_edges(s, relation)} | |
| def subjects(self, o: str, relation: str) -> set[str]: | |
| return {e.subject for e in self.in_edges(o, relation)} | |
| def neighbourhood(self, seeds: list[str], hops: int) -> list[Edge]: | |
| """Edges reachable within ``hops`` of the seeds (ignoring direction), BFS order.""" | |
| seen_nodes = {s for s in seeds if s in self.g} | |
| frontier = list(seen_nodes) | |
| out: list[Edge] = [] | |
| seen_edges: set[str] = set() | |
| for _ in range(hops): | |
| nxt = [] | |
| for n in frontier: | |
| for e in self.out_edges(n) + self.in_edges(n): | |
| if e.edge_id in seen_edges: | |
| continue | |
| seen_edges.add(e.edge_id) | |
| out.append(e) | |
| other = e.object if e.subject == n else e.subject | |
| if other not in seen_nodes: | |
| seen_nodes.add(other) | |
| nxt.append(other) | |
| frontier = nxt | |
| return out | |
| def results_for(self, method: str | None = None, dataset: str | None = None, metric: str | None = None) -> list[ResultRecord]: | |
| return [ | |
| r | |
| for r in self.results.values() | |
| if (method is None or r.method == method) and (dataset is None or r.dataset == dataset) and (metric is None or r.metric == metric) | |
| ] | |
| def stats(self) -> dict: | |
| deg = Counter() | |
| for e in self.edges.values(): | |
| deg[e.subject] += 1 | |
| deg[e.object] += 1 | |
| und = nx.Graph(self.g) | |
| comps = sorted((len(c) for c in nx.connected_components(und)), reverse=True) if len(und) else [] | |
| return { | |
| "nodes": len(self.nodes), | |
| "nodes_by_type": dict(Counter(n.type for n in self.nodes.values())), | |
| "edges": len(self.edges), | |
| "distinct_facts": len(self._facts), | |
| "edges_by_relation": dict(Counter(e.relation for e in self.edges.values())), | |
| "results": len(self.results), | |
| "components": len(comps), | |
| "largest_component_nodes": comps[0] if comps else 0, | |
| "top_degree": [(self.name(n), d) for n, d in deg.most_common(10)], | |
| "edges_with_confidence": sum(1 for e in self.edges.values() if e.confidence is not None), | |
| } | |