Spaces:
Running on Zero
Running on Zero
| """Per-chunk knowledge-graph extraction with verbatim provenance. | |
| Each chunk is sent to the extractor model once. Every returned triple or result must | |
| carry a quote that anchors verbatim in that chunk (``quotes.locate``); anything else is | |
| rejected with a recorded reason. Accepted records keep the chunk ID and character | |
| offsets of their evidence. Raw extractor replies live in the LLM cache; this module | |
| writes one JSONL record per chunk with accepted and rejected items, so precision of the | |
| extraction stage can be audited and measured. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| from collections import Counter | |
| from concurrent.futures import ThreadPoolExecutor | |
| from pathlib import Path | |
| from .canonical import is_generic | |
| from .llm import LLMClient, LLMError, UsageMeter, call_structured | |
| from .prompts import EXTRACT_SYSTEM, THIS_PAPER, ExtractionOut, extract_user | |
| from .quotes import locate | |
| from .schemas import SIGNATURES, Chunk, Document | |
| from .verification import strip_headings | |
| def _validate_triple(t, chunk_text: str) -> tuple[dict | None, str | None]: | |
| subj_type = "Paper" if t.subject.strip() == THIS_PAPER else t.subject_type | |
| obj_type = t.object_type | |
| if t.object.strip() == THIS_PAPER: | |
| return None, "paper_as_object" | |
| if (subj_type, obj_type) not in SIGNATURES[t.relation]: | |
| return None, "signature" | |
| if subj_type == "Paper" and t.subject.strip() != THIS_PAPER: | |
| return None, "foreign_paper_subject" | |
| if subj_type != "Paper" and is_generic(t.subject): | |
| return None, "generic_name" | |
| if is_generic(t.object): | |
| return None, "generic_name" | |
| loc = locate(t.quote, chunk_text) | |
| if loc is None: | |
| return None, "unanchored_quote" | |
| return { | |
| "subject": t.subject.strip(), | |
| "subject_type": subj_type, | |
| "relation": t.relation, | |
| "object": t.object.strip(), | |
| "object_type": obj_type, | |
| "quote": t.quote, | |
| "start_char": loc[0], | |
| "end_char": loc[1], | |
| "elided": loc[2], | |
| }, None | |
| def _validate_result(r, chunk_text: str) -> tuple[dict | None, str | None]: | |
| # THIS_PAPER as the method is resolved to the paper's proposed method at graph build | |
| # (or dropped there); it is never allowed as a dataset or metric. | |
| if (r.method.strip() != THIS_PAPER and is_generic(r.method)) or is_generic(r.dataset) or is_generic(r.metric): | |
| return None, "generic_name" | |
| if not any(ch.isdigit() for ch in r.value): | |
| return None, "non_numeric_value" | |
| loc = locate(r.quote, chunk_text) | |
| if loc is None: | |
| return None, "unanchored_quote" | |
| if r.value.strip().rstrip("%") not in chunk_text[loc[0] : loc[1]]: | |
| return None, "value_not_in_quote" | |
| return { | |
| "method": r.method.strip(), | |
| "dataset": r.dataset.strip(), | |
| "metric": r.metric.strip(), | |
| "value": r.value.strip(), | |
| "quote": r.quote, | |
| "start_char": loc[0], | |
| "end_char": loc[1], | |
| "elided": loc[2], | |
| }, None | |
| def extract_chunk(client: LLMClient, doc: Document, chunk: Chunk, meter: UsageMeter | None = None) -> dict: | |
| """One chunk -> record with accepted/rejected triples and results.""" | |
| # Quotes are anchored against the chunk text exactly as stored (headings included). | |
| rec: dict = {"chunk_id": chunk.chunk_id, "doc_id": doc.doc_id, "extractor_model": client.spec.model} | |
| try: | |
| out = call_structured( | |
| client, meter, "extract", EXTRACT_SYSTEM, extract_user(doc.title, chunk.section_heading, strip_headings(chunk.text)), ExtractionOut, max_tokens=4096 | |
| ) | |
| except LLMError as e: | |
| rec.update(error=str(e), triples=[], results=[], rejected=[]) | |
| return rec | |
| triples, results, rejected = [], [], [] | |
| for t in out.triples: | |
| ok, why = _validate_triple(t, chunk.text) | |
| (triples.append(ok) if ok else rejected.append({"kind": "triple", "reason": why, **t.model_dump()})) | |
| for r in out.results: | |
| ok, why = _validate_result(r, chunk.text) | |
| (results.append(ok) if ok else rejected.append({"kind": "result", "reason": why, **r.model_dump()})) | |
| rec.update(error=None, triples=triples, results=results, rejected=rejected) | |
| return rec | |
| def extract_corpus( | |
| client: LLMClient, | |
| docs: dict[str, Document], | |
| chunks: list[Chunk], | |
| out_path: str | Path, | |
| workers: int = 4, | |
| ) -> dict: | |
| """Resumable: chunks already present in ``out_path`` are skipped.""" | |
| out_path = Path(out_path) | |
| out_path.parent.mkdir(parents=True, exist_ok=True) | |
| done: set[str] = set() | |
| if out_path.exists(): | |
| for line in out_path.read_text(encoding="utf-8").splitlines(): | |
| if line.strip(): | |
| r = json.loads(line) | |
| if not r.get("error"): | |
| done.add(r["chunk_id"]) | |
| todo = [c for c in chunks if c.chunk_id not in done] | |
| meter = UsageMeter() | |
| with out_path.open("a", encoding="utf-8") as f, ThreadPoolExecutor(max_workers=workers) as pool: | |
| for rec in pool.map(lambda c: extract_chunk(client, docs[c.doc_id], c, meter), todo): | |
| f.write(json.dumps(rec, ensure_ascii=False) + "\n") | |
| f.flush() | |
| return {"chunks_total": len(chunks), "chunks_skipped": len(done), "chunks_processed": len(todo), "usage": meter.summary()} | |
| def extraction_stats(records: list[dict]) -> dict: | |
| reasons = Counter(r["reason"] for rec in records for r in rec.get("rejected", [])) | |
| accepted_t = sum(len(rec.get("triples", [])) for rec in records) | |
| accepted_r = sum(len(rec.get("results", [])) for rec in records) | |
| rejected = sum(reasons.values()) | |
| proposed = accepted_t + accepted_r + rejected | |
| return { | |
| "chunks": len(records), | |
| "chunk_errors": sum(1 for r in records if r.get("error")), | |
| "triples_accepted": accepted_t, | |
| "results_accepted": accepted_r, | |
| "items_rejected": rejected, | |
| "rejection_reasons": dict(reasons), | |
| "unanchored_rate": (reasons.get("unanchored_quote", 0) / proposed) if proposed else None, | |
| "relations": dict(Counter(t["relation"] for rec in records for t in rec.get("triples", []))), | |
| } | |