Spaces:
Sleeping
Sleeping
| """Corpus preparation and graph building (file-based; no database).""" | |
| from __future__ import annotations | |
| import json | |
| from collections import Counter | |
| from pathlib import Path | |
| from .canonical import Canonicalizer | |
| from .chunking import chunk_document | |
| from .config import ProjectConfig, load_yaml, manifest, write_json | |
| from .corpus import metric_hits, select_corpus | |
| from .data.qasper import load_split | |
| from .graph import KG, build_graph, save_graph | |
| from .schemas import Chunk, Document, paper_prefix | |
| def read_jsonl(path: str | Path) -> list[dict]: | |
| p = Path(path) | |
| if not p.exists(): | |
| return [] | |
| return [json.loads(x) for x in p.read_text(encoding="utf-8").splitlines() if x.strip()] | |
| def write_jsonl(path: str | Path, rows) -> None: | |
| p = Path(path) | |
| p.parent.mkdir(parents=True, exist_ok=True) | |
| with p.open("w", encoding="utf-8") as f: | |
| for r in rows: | |
| f.write((r.model_dump_json() if hasattr(r, "model_dump_json") else json.dumps(r, ensure_ascii=False)) + "\n") | |
| def prepare_corpus(cfg: ProjectConfig) -> dict: | |
| docs: dict[str, Document] = {} | |
| for split in cfg.corpus.splits: | |
| d, _ = load_split(cfg.paths.raw, split) | |
| docs.update(d) | |
| sel = select_corpus(docs, cfg.corpus.n_papers, cfg.corpus.min_mentions) | |
| papers = [docs[p] for p in sorted(sel.papers)] | |
| chunks: list[Chunk] = [] | |
| for doc in papers: | |
| chunks.extend(chunk_document(doc, cfg.chunking, id_prefix=paper_prefix(doc.doc_id))) | |
| out = Path(cfg.paths.processed) | |
| write_jsonl(out / "papers.jsonl", papers) | |
| write_jsonl(out / "chunks.jsonl", chunks) | |
| metric_df = Counter(m for d in papers for m, n in metric_hits(d).items() if n >= cfg.corpus.min_mentions) | |
| toks = sorted(c.n_tokens for c in chunks) | |
| stats = { | |
| **sel.stats, | |
| "selection_order": sel.papers, | |
| "chunks": len(chunks), | |
| "chunk_tokens_median": toks[len(toks) // 2] if toks else 0, | |
| "chunk_tokens_p90": toks[int(0.9 * len(toks))] if toks else 0, | |
| "corpus_tokens": sum(toks), | |
| "metric_lexicon_df": dict(metric_df.most_common()), | |
| "paper_datasets": sel.mentions, | |
| } | |
| write_json(out / "corpus_manifest.json", manifest({"corpus_config": cfg.corpus.model_dump(), "chunking": cfg.chunking.model_dump(), "stats": stats})) | |
| return stats | |
| def load_corpus(cfg: ProjectConfig) -> tuple[dict[str, Document], list[Chunk]]: | |
| out = Path(cfg.paths.processed) | |
| papers = {d["doc_id"]: Document(**d) for d in read_jsonl(out / "papers.jsonl")} | |
| chunks = [Chunk(**c) for c in read_jsonl(out / "chunks.jsonl")] | |
| if not papers: | |
| raise FileNotFoundError(f"no processed corpus in {out}; run `graphtrace corpus` first") | |
| return papers, chunks | |
| def build_canonicalizer(cfg: ProjectConfig, papers: dict[str, Document]) -> Canonicalizer: | |
| curated = load_yaml(cfg.extraction.curated_aliases).get("aliases", {}) if Path(cfg.extraction.curated_aliases).exists() else {} | |
| texts = [p for d in papers.values() for s in d.sections for p in s.paragraphs] | |
| return Canonicalizer.build(curated, texts) | |
| def build(cfg: ProjectConfig, score_confidence: bool | None = None) -> dict: | |
| papers, _ = load_corpus(cfg) | |
| records = read_jsonl(Path(cfg.paths.graph) / "extractions.jsonl") | |
| if not records: | |
| raise FileNotFoundError("no extractions found; run `graphtrace extract` first") | |
| canon = build_canonicalizer(cfg, papers) | |
| build_stats: dict = {} | |
| nodes, edges, results = build_graph(records, papers, canon, stats=build_stats) | |
| confidence: dict[str, float] = {} | |
| if score_confidence if score_confidence is not None else cfg.extraction.nli_confidence: | |
| from .support import score_edges | |
| from .verification import NLIVerifier | |
| confidence = score_edges(KG(nodes, edges, results), NLIVerifier(cfg.extraction.nli_model)) | |
| edges = [e.model_copy(update={"confidence": confidence.get(e.edge_id)}) for e in edges] | |
| save_graph(cfg.paths.graph, nodes, edges, results, canon.merges) | |
| kg = KG(nodes, edges, results) | |
| stats = kg.stats() | |
| write_json( | |
| Path(cfg.paths.graph) / "graph_manifest.json", | |
| manifest( | |
| { | |
| "stats": stats, | |
| **build_stats, | |
| "alias_merges": len(canon.merges), | |
| "nli_model": cfg.extraction.nli_model if confidence else None, | |
| "min_edge_confidence_for_questions_and_methods": cfg.bench.min_edge_confidence, | |
| "edges_below_min_confidence": len(kg.edges) - len(kg.filtered(cfg.bench.min_edge_confidence).edges), | |
| } | |
| ), | |
| ) | |
| return stats | |