"""Build a knowledge graph from OCR pipeline markdown output. Schema ------ Nodes doc: Document (source image + OCR metadata) chunk:: Chunk (text segment of a document) entity: Entity (gene/protein/chemical mention) citation: Citation (doi/pmid/arxiv/isbn) Edges doc -[:CONTAINS]-> chunk chunk -[:MENTIONS]-> entity chunk -[:CITES]-> citation entity -[:CO_OCCURS]-> entity (weight = shared chunks) """ from __future__ import annotations import re from collections.abc import Iterable from dataclasses import dataclass, field from pathlib import Path from typing import Any import yaml DEFAULT_CONSOLIDATED_NAME = "all_ocr.md" _TEXT_SECTION = "## Extracted Text" @dataclass class ChunkRecord: """One parsed pipeline markdown chunk file.""" source_hash: str source_path: str chunk_index: int total_chunks: int text: str timestamp: str = "" ocr_engine: str = "" ocr_confidence_mean: float = 0.0 language: str = "" has_figures: bool = False has_tables: bool = False entities: list[str] = field(default_factory=list) citations: list[str] = field(default_factory=list) # "type:identifier" strings @property def doc_key(self) -> str: return f"doc:{self.source_hash}" @property def chunk_key(self) -> str: return f"chunk:{self.source_hash}:{self.chunk_index}" def _split_frontmatter(raw: str) -> tuple[dict[str, Any], str]: """Return (frontmatter dict, body). Empty dict when no frontmatter.""" match = re.match(r"\A---\s*\n(.*?)\n---\s*\n?(.*)\Z", raw, re.DOTALL) if not match: return {}, raw data = yaml.safe_load(match.group(1)) return (data if isinstance(data, dict) else {}), match.group(2) def _extract_text_section(body: str) -> str: """Pull the text between '## Extracted Text' and the next '## ' heading.""" if _TEXT_SECTION not in body: return "" after = body.split(_TEXT_SECTION, 1)[1] after = re.split(r"\n## ", after, maxsplit=1)[0] return after.strip() def parse_chunk_file(path: Path) -> ChunkRecord | None: """Parse one pipeline chunk markdown file. None when it is not a chunk file.""" try: raw = path.read_text(encoding="utf-8") except OSError: return None frontmatter, body = _split_frontmatter(raw) if not frontmatter or "source_hash" not in frontmatter: return None citations = frontmatter.get("citations_found") or [] if not isinstance(citations, list): citations = [str(citations)] return ChunkRecord( source_hash=str(frontmatter["source_hash"]), source_path=str(frontmatter.get("source_path", "")), chunk_index=int(frontmatter.get("chunk_index", 0)), total_chunks=int(frontmatter.get("total_chunks", 1)), text=_extract_text_section(body), timestamp=str(frontmatter.get("timestamp", "")), ocr_engine=str(frontmatter.get("ocr_engine", "")), ocr_confidence_mean=float(frontmatter.get("ocr_confidence_mean") or 0.0), language=str(frontmatter.get("language", "")), has_figures=bool(frontmatter.get("has_figures", False)), has_tables=bool(frontmatter.get("has_tables", False)), entities=[str(e) for e in (frontmatter.get("detected_entities") or [])], citations=citations, ) def find_chunk_files( output_dir: Path, consolidated_name: str = DEFAULT_CONSOLIDATED_NAME ) -> list[Path]: """All chunk markdown files under an output directory, excluding consolidated files.""" return sorted( p for p in output_dir.rglob("*.md") if p.is_file() and p.name != consolidated_name ) def _normalize_entity(text: str) -> str: return " ".join(text.strip().split()).lower() def _split_citation(raw: str) -> tuple[str, str]: """'isbn:ISBN:9780...' -> ('isbn', 'ISBN:9780...').""" if ":" in raw: ctype, identifier = raw.split(":", 1) return ctype.strip().lower(), identifier.strip() return "unknown", raw.strip() def build_graph(records: Iterable[ChunkRecord]) -> Any: """Assemble the knowledge graph from parsed chunk records.""" import networkx as nx graph = nx.DiGraph() for record in records: graph.add_node( record.doc_key, kind="document", source_path=record.source_path, timestamp=record.timestamp, ocr_engine=record.ocr_engine, ocr_confidence_mean=record.ocr_confidence_mean, language=record.language, has_figures=record.has_figures, has_tables=record.has_tables, total_chunks=record.total_chunks, ) graph.add_node( record.chunk_key, kind="chunk", text=record.text, chunk_index=record.chunk_index, total_chunks=record.total_chunks, ) graph.add_edge(record.doc_key, record.chunk_key, kind="CONTAINS") for entity in record.entities: key = f"entity:{_normalize_entity(entity)}" if key not in graph: graph.add_node(key, kind="entity", text=entity.strip(), mentions=0) graph.nodes[key]["mentions"] += 1 graph.add_edge(record.chunk_key, key, kind="MENTIONS") for raw_citation in record.citations: ctype, identifier = _split_citation(raw_citation) key = f"citation:{identifier}" if key not in graph: graph.add_node(key, kind="citation", citation_type=ctype, identifier=identifier) graph.add_edge(record.chunk_key, key, kind="CITES") # Co-occurrence: entities sharing a chunk, one edge per unordered pair. normalized = sorted({_normalize_entity(e) for e in record.entities if e.strip()}) for i, left in enumerate(normalized): for right in normalized[i + 1 :]: a_key, b_key = f"entity:{left}", f"entity:{right}" if graph.has_edge(a_key, b_key): graph[a_key][b_key]["weight"] += 1 else: graph.add_edge(a_key, b_key, kind="CO_OCCURS", weight=1) return graph def build_from_directory( output_dir: Path, consolidated_name: str = DEFAULT_CONSOLIDATED_NAME ) -> Any: """Parse every chunk file under output_dir and build the graph.""" records = ( record for record in (parse_chunk_file(p) for p in find_chunk_files(output_dir, consolidated_name)) if record is not None ) return build_graph(records)