kg: graph build, traversal queries, neo4j export - kg_ocr.graph builds a networkx graph (docs, chunks, entities, citations, co-occurrence) from chunk markdown - analyzer for summaries, top entities/citations, anomaly checks - traversal: chunks_for_entity/citation, related_entities, expand_context - export: JSON round-trip, GraphML, batched MERGE into neo4j - new CLI: ocr-pipeline kg build|stats|query|export - lazy kg_ocr imports, networkx/neo4j behind extras - dropped dead watch.py shim, added KgConfig stub - trimmed README, updated TODO
Some checks failed
tests / core (macos-latest, py3.11) (push) Has been cancelled
tests / core (macos-latest, py3.12) (push) Has been cancelled
tests / core (macos-latest, py3.13) (push) Has been cancelled
tests / core (ubuntu-latest, py3.11) (push) Has been cancelled
tests / core (ubuntu-latest, py3.12) (push) Has been cancelled
tests / core (ubuntu-latest, py3.13) (push) Has been cancelled
tests / doctor CLI smoke test (push) Has been cancelled
Some checks failed
tests / core (macos-latest, py3.11) (push) Has been cancelled
tests / core (macos-latest, py3.12) (push) Has been cancelled
tests / core (macos-latest, py3.13) (push) Has been cancelled
tests / core (ubuntu-latest, py3.11) (push) Has been cancelled
tests / core (ubuntu-latest, py3.12) (push) Has been cancelled
tests / core (ubuntu-latest, py3.13) (push) Has been cancelled
tests / doctor CLI smoke test (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,39 @@
|
||||
from .analyzer import (
|
||||
Anomaly,
|
||||
detect_anomalies,
|
||||
summary,
|
||||
top_citations,
|
||||
top_cooccurrences,
|
||||
top_entities,
|
||||
)
|
||||
from .builder import (
|
||||
ChunkRecord,
|
||||
build_from_directory,
|
||||
build_graph,
|
||||
find_chunk_files,
|
||||
parse_chunk_file,
|
||||
)
|
||||
from .traversal import (
|
||||
chunks_for_citation,
|
||||
chunks_for_entity,
|
||||
expand_context,
|
||||
related_entities,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Anomaly",
|
||||
"ChunkRecord",
|
||||
"build_from_directory",
|
||||
"build_graph",
|
||||
"chunks_for_citation",
|
||||
"chunks_for_entity",
|
||||
"detect_anomalies",
|
||||
"expand_context",
|
||||
"find_chunk_files",
|
||||
"parse_chunk_file",
|
||||
"related_entities",
|
||||
"summary",
|
||||
"top_citations",
|
||||
"top_cooccurrences",
|
||||
"top_entities",
|
||||
]
|
||||
|
||||
@@ -1 +1,117 @@
|
||||
# also anomaly detection here
|
||||
"""Analyze a knowledge graph built from OCR pipeline output."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
LOW_CONFIDENCE_THRESHOLD = 0.4
|
||||
HUB_MENTIONS_THRESHOLD = 20
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Anomaly:
|
||||
kind: str
|
||||
node: str
|
||||
detail: str
|
||||
|
||||
|
||||
def _nodes_of_kind(graph: Any, kind: str) -> list[tuple[str, dict]]:
|
||||
return [(n, d) for n, d in graph.nodes(data=True) if d.get("kind") == kind]
|
||||
|
||||
|
||||
def summary(graph: Any) -> dict[str, int]:
|
||||
"""Node and edge counts grouped by kind."""
|
||||
counts: dict[str, int] = {}
|
||||
for _, data in graph.nodes(data=True):
|
||||
kind = data.get("kind", "unknown")
|
||||
counts[f"nodes:{kind}"] = counts.get(f"nodes:{kind}", 0) + 1
|
||||
for _, _, data in graph.edges(data=True):
|
||||
kind = data.get("kind", "unknown")
|
||||
counts[f"edges:{kind}"] = counts.get(f"edges:{kind}", 0) + 1
|
||||
return counts
|
||||
|
||||
|
||||
def top_entities(graph: Any, limit: int = 10) -> list[tuple[str, int]]:
|
||||
"""Most-mentioned entities as (text, mentions), descending."""
|
||||
entities = [
|
||||
(data.get("text", ""), int(data.get("mentions", 0)))
|
||||
for _, data in _nodes_of_kind(graph, "entity")
|
||||
]
|
||||
entities.sort(key=lambda item: item[1], reverse=True)
|
||||
return entities[:limit]
|
||||
|
||||
|
||||
def top_cooccurrences(graph: Any, limit: int = 10) -> list[tuple[str, str, int]]:
|
||||
"""Strongest entity co-occurrence pairs as (left, right, weight)."""
|
||||
pairs = []
|
||||
for left, right, data in graph.edges(data=True):
|
||||
if data.get("kind") != "CO_OCCURS":
|
||||
continue
|
||||
left_text = graph.nodes[left].get("text", left)
|
||||
right_text = graph.nodes[right].get("text", right)
|
||||
pairs.append((left_text, right_text, int(data.get("weight", 1))))
|
||||
pairs.sort(key=lambda item: item[2], reverse=True)
|
||||
return pairs[:limit]
|
||||
|
||||
|
||||
def top_citations(graph: Any, limit: int = 10) -> list[tuple[str, str, int]]:
|
||||
"""Most-cited identifiers as (type, identifier, citing chunks)."""
|
||||
rows = []
|
||||
for node, data in _nodes_of_kind(graph, "citation"):
|
||||
citing = sum(
|
||||
1 for _, _, edge in graph.in_edges(node, data=True) if edge.get("kind") == "CITES"
|
||||
)
|
||||
rows.append((data.get("citation_type", ""), data.get("identifier", node), citing))
|
||||
rows.sort(key=lambda item: item[2], reverse=True)
|
||||
return rows[:limit]
|
||||
|
||||
|
||||
def detect_anomalies(
|
||||
graph: Any,
|
||||
low_confidence: float = LOW_CONFIDENCE_THRESHOLD,
|
||||
hub_mentions: int = HUB_MENTIONS_THRESHOLD,
|
||||
) -> list[Anomaly]:
|
||||
"""Flag suspicious graph regions worth re-checking.
|
||||
|
||||
- documents with very low OCR confidence (likely OCR failure)
|
||||
- chunks with empty extracted text
|
||||
- documents with zero entities and zero citations (nothing extracted)
|
||||
- entity hubs (mentioned everywhere; often a junk pattern)
|
||||
"""
|
||||
anomalies: list[Anomaly] = []
|
||||
|
||||
documents = _nodes_of_kind(graph, "document")
|
||||
for node, data in documents:
|
||||
confidence = float(data.get("ocr_confidence_mean") or 0.0)
|
||||
if confidence < low_confidence:
|
||||
anomalies.append(
|
||||
Anomaly(
|
||||
"low_confidence",
|
||||
node,
|
||||
f"ocr_confidence_mean={confidence:.2f} < {low_confidence}",
|
||||
)
|
||||
)
|
||||
|
||||
meaningful_out = 0
|
||||
for _, target, edge in graph.out_edges(node, data=True):
|
||||
if edge.get("kind") != "CONTAINS":
|
||||
continue
|
||||
for _, _, chunk_edge in graph.out_edges(target, data=True):
|
||||
if chunk_edge.get("kind") in ("MENTIONS", "CITES"):
|
||||
meaningful_out += 1
|
||||
if meaningful_out == 0:
|
||||
anomalies.append(Anomaly("empty_document", node, "no entities or citations extracted"))
|
||||
|
||||
for node, data in _nodes_of_kind(graph, "chunk"):
|
||||
if not (data.get("text") or "").strip():
|
||||
anomalies.append(Anomaly("empty_chunk", node, "extracted text is empty"))
|
||||
|
||||
for node, data in _nodes_of_kind(graph, "entity"):
|
||||
mentions = int(data.get("mentions", 0))
|
||||
if mentions >= hub_mentions:
|
||||
anomalies.append(
|
||||
Anomaly("entity_hub", node, f"mentioned {mentions} times; possible junk pattern")
|
||||
)
|
||||
|
||||
return anomalies
|
||||
|
||||
@@ -0,0 +1,193 @@
|
||||
"""Build a knowledge graph from OCR pipeline markdown output.
|
||||
|
||||
Schema
|
||||
------
|
||||
Nodes
|
||||
doc:<source_hash> Document (source image + OCR metadata)
|
||||
chunk:<source_hash>:<index> Chunk (text segment of a document)
|
||||
entity:<normalized text> Entity (gene/protein/chemical mention)
|
||||
citation:<identifier> 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)
|
||||
|
||||
89
kg_ocr/graph/traversal.py
Normal file
89
kg_ocr/graph/traversal.py
Normal file
@@ -0,0 +1,89 @@
|
||||
"""Traversal queries over the knowledge graph: entity/citation -> chunks, neighbors."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from .builder import _normalize_entity
|
||||
|
||||
|
||||
def _entity_key(text: str) -> str:
|
||||
return f"entity:{_normalize_entity(text)}"
|
||||
|
||||
|
||||
def _chunk_payload(graph: Any, chunk_node: str) -> dict[str, Any]:
|
||||
data = graph.nodes[chunk_node]
|
||||
doc_node = next(
|
||||
(src for src, _, e in graph.in_edges(chunk_node, data=True) if e.get("kind") == "CONTAINS"),
|
||||
None,
|
||||
)
|
||||
doc = graph.nodes[doc_node] if doc_node is not None else {}
|
||||
return {
|
||||
"chunk": chunk_node,
|
||||
"chunk_index": data.get("chunk_index", 0),
|
||||
"text": data.get("text", ""),
|
||||
"source_path": doc.get("source_path", ""),
|
||||
"ocr_engine": doc.get("ocr_engine", ""),
|
||||
"ocr_confidence_mean": doc.get("ocr_confidence_mean", 0.0),
|
||||
"timestamp": doc.get("timestamp", ""),
|
||||
}
|
||||
|
||||
|
||||
def chunks_for_entity(graph: Any, entity_text: str) -> list[dict[str, Any]]:
|
||||
"""All chunks mentioning the given entity (case-insensitive)."""
|
||||
key = _entity_key(entity_text)
|
||||
if key not in graph:
|
||||
return []
|
||||
chunks = [
|
||||
src for src, _, edge in graph.in_edges(key, data=True) if edge.get("kind") == "MENTIONS"
|
||||
]
|
||||
return [_chunk_payload(graph, c) for c in sorted(chunks)]
|
||||
|
||||
|
||||
def chunks_for_citation(graph: Any, identifier: str) -> list[dict[str, Any]]:
|
||||
"""All chunks citing the given identifier (e.g. '10.1038/nature12345')."""
|
||||
key = identifier if identifier.startswith("citation:") else f"citation:{identifier}"
|
||||
if key not in graph:
|
||||
return []
|
||||
chunks = [src for src, _, edge in graph.in_edges(key, data=True) if edge.get("kind") == "CITES"]
|
||||
return [_chunk_payload(graph, c) for c in sorted(chunks)]
|
||||
|
||||
|
||||
def related_entities(graph: Any, entity_text: str, limit: int = 10) -> list[tuple[str, int]]:
|
||||
"""Entities co-occurring with the given one, as (text, shared-chunk weight)."""
|
||||
key = _entity_key(entity_text)
|
||||
if key not in graph:
|
||||
return []
|
||||
neighbors: dict[str, int] = {}
|
||||
for _, target, edge in graph.out_edges(key, data=True):
|
||||
if edge.get("kind") == "CO_OCCURS":
|
||||
neighbors[target] = edge.get("weight", 1)
|
||||
for source, _, edge in graph.in_edges(key, data=True):
|
||||
if edge.get("kind") == "CO_OCCURS":
|
||||
neighbors[source] = max(neighbors.get(source, 0), edge.get("weight", 1))
|
||||
ranked = sorted(
|
||||
((graph.nodes[n].get("text", n), w) for n, w in neighbors.items()),
|
||||
key=lambda item: item[1],
|
||||
reverse=True,
|
||||
)
|
||||
return ranked[:limit]
|
||||
|
||||
|
||||
def expand_context(graph: Any, entity_text: str, limit: int = 5) -> list[dict[str, Any]]:
|
||||
"""Chunks mentioning the entity or its strongest co-occurring neighbors.
|
||||
|
||||
This is the graph-RAG primitive: one hop of CO_OCCURS expansion, so a query
|
||||
for 'BRCA1' also surfaces chunks that only mention its interactors.
|
||||
"""
|
||||
seen: set[str] = set()
|
||||
results: list[dict[str, Any]] = []
|
||||
|
||||
seeds = [entity_text] + [text for text, _ in related_entities(graph, entity_text, limit=limit)]
|
||||
for seed in seeds:
|
||||
for payload in chunks_for_entity(graph, seed):
|
||||
if payload["chunk"] in seen:
|
||||
continue
|
||||
seen.add(payload["chunk"])
|
||||
payload = {**payload, "via_entity": seed}
|
||||
results.append(payload)
|
||||
return results
|
||||
Reference in New Issue
Block a user