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:
@@ -24,3 +24,11 @@ WATCH_DB_PATH=./data/processed_files.db
|
||||
# Optional: GROBID for citation parsing
|
||||
GROBID_ENABLED=false
|
||||
GROBID_URL=http://localhost:8070
|
||||
|
||||
# Optional: Neo4j for `ocr-pipeline kg export -f neo4j`
|
||||
NEO4J_URI=bolt://localhost:7687
|
||||
NEO4J_USER=neo4j
|
||||
NEO4J_PASSWORD=
|
||||
|
||||
# Optional: LLM provider keys for kg_ocr.rag (ask_wllm)
|
||||
# OPENROUTER_API_KEY=
|
||||
@@ -8,6 +8,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
- Knowledge graph: `kg_ocr.graph` builds a NetworkX graph (documents, chunks, entities, citations, entity co-occurrence) from pipeline chunk markdown; `kg_ocr.graph.analyzer` provides summaries, top entities/citations, and anomaly detection (low OCR confidence, empty chunks/documents, entity hubs).
|
||||
- Graph traversal: `chunks_for_entity`, `chunks_for_citation`, `related_entities`, `expand_context` (one-hop CO_OCCURS expansion for graph-RAG), exposed via `ocr-pipeline kg query`.
|
||||
- `kg_ocr.embeddings.config_loader.KgConfig` for traversal tunables and Neo4j credentials from env.
|
||||
- `kg_ocr.export`: node-link JSON (lossless round-trip), GraphML for Gephi/Cytoscape, and batched MERGE export to Neo4j via `Neo4jExporter`.
|
||||
- CLI: `ocr-pipeline kg build|stats|query|export` subcommands with lazy kg_ocr imports.
|
||||
- Optional extras: `kg` (networkx + txtai + litellm + python-dotenv), `neo4j` (driver).
|
||||
- `init` command: writes a minimal `config.yaml` with platform-correct defaults and checks for Tesseract.
|
||||
- `demo` command: runs the pipeline on a bundled sample screenshot, prints the output Markdown, and verifies the install in <30s.
|
||||
- `setup [basic|full]` command: executes install steps with confirmation (`--yes`/`--dry-run`).
|
||||
@@ -25,7 +31,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
- OCR engine default is now `tesseract` (was `paddleocr`).
|
||||
- `run` results table now shows Skipped count.
|
||||
- Run directories are created lazily on first write (no empty `run_NNN`).
|
||||
- `kg_ocr` package marked experimental; deps under `--extra kg`, helpful ImportError on bare import.
|
||||
- `kg_ocr` is now a functional knowledge-graph package (graph/export need only networkx); txtai/litellm imports are lazy so the package imports in slim environments.
|
||||
|
||||
### Fixed
|
||||
- ISBN regex no longer false-positives on years, URLs, job IDs, or dilutions. Now requires an explicit `ISBN` prefix and matches strict 13- or 10-digit structures.
|
||||
@@ -33,6 +39,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
- `import ocr_pipeline` no longer depends on the caller's cwd.
|
||||
|
||||
### Removed
|
||||
- Dead `src/ocr_pipeline/watch.py` shim (shadowed by the `watch/` package).
|
||||
- Unused core dependencies: `tqdm`, `scikit-image`, `sqlite-utils`, `xxhash`, `python-slugify`, `python-magic`, `legacy-cgi`.
|
||||
- Core dependency on `spacy` (moved to `scientific` extra).
|
||||
- Core dependency on `paddleocr`/`paddlepaddle` (moved to `paddle` extra).
|
||||
|
||||
193
README.md
193
README.md
@@ -1,90 +1,118 @@
|
||||
# OCR Pipeline for Life Science Screenshots
|
||||
|
||||
[](./.github/workflows/tests.yml)
|
||||
|
||||
Turns scientific screenshots into RAG-ready Markdown, with metadata attached.
|
||||
|
||||
- OCR via PaddleOCR, falls back to Tesseract
|
||||
- Preprocessing: deskew, denoise, CLAHE, line removal
|
||||
- Figure/table detection
|
||||
- Entity extraction with scispaCy
|
||||
- Citation matching
|
||||
- Citation matching (DOI, PMID, arXiv, ISBN)
|
||||
- Chunking for embedding
|
||||
- Knowledge graph over the output (`kg` commands)
|
||||
|
||||
## Quick start
|
||||
|
||||
The default install is small: Tesseract OCR only, no multi-GB ML downloads.
|
||||
|
||||
```bash
|
||||
uv sync # install the core pipeline
|
||||
uv run ocr-pipeline setup # installs Tesseract if missing (asks first)
|
||||
uv run ocr-pipeline init # write a config for this machine
|
||||
uv run ocr-pipeline demo # verify everything works on a sample screenshot
|
||||
uv sync
|
||||
uv run ocr-pipeline setup # installs Tesseract if missing
|
||||
uv run ocr-pipeline init # write a config
|
||||
uv run ocr-pipeline demo # check it works
|
||||
uv run ocr-pipeline run # process your screenshots
|
||||
```
|
||||
|
||||
Config is found in this order: `$OCR_PIPELINE_CONFIG`, `./config.yaml`, then the
|
||||
per-user config directory (`~/Library/Application Support/ocr-pipeline` on macOS,
|
||||
`~/.config/ocr-pipeline` on Linux). Running `run` with neither a config nor
|
||||
`--input-dir` stops with guidance instead of scanning a default directory.
|
||||
Default install is small: Tesseract only, no multi-GB ML downloads.
|
||||
|
||||
Use `--force` to reprocess files already recorded in the processing database. The
|
||||
`--exclude` option is repeatable and prevents recursive scans from entering bundles
|
||||
such as macOS `.photoslibrary` directories.
|
||||
## Commands
|
||||
|
||||
### Optional capabilities
|
||||
- `run` / `watch` — process a directory, or watch for new screenshots
|
||||
- `status` / `stats` — what's been processed
|
||||
- `doctor` — check engines, models, config
|
||||
- `kg build|stats|query|export` — knowledge graph over the output
|
||||
- `setup [basic|full]` — install dependencies step by step
|
||||
|
||||
Heavy ML features are opt-in extras and degrade gracefully when absent (a warning
|
||||
is logged and that step is skipped):
|
||||
Use `--force` to reprocess. `--exclude "*.photoslibrary/*"` keeps recursive
|
||||
scans out of macOS photo bundles.
|
||||
|
||||
<details>
|
||||
<summary><b>Optional extras</b></summary>
|
||||
|
||||
Heavy ML features are opt-in and degrade gracefully when missing (logged, step
|
||||
skipped, run continues):
|
||||
|
||||
| Extra | Provides |
|
||||
|-------|----------|
|
||||
| `paddle` | PaddleOCR engine (`ocr.engine: paddleocr` or `auto`) |
|
||||
| `scientific` | spaCy/scispaCy NER (`entities.enabled: true`) |
|
||||
| `paddle` | PaddleOCR engine |
|
||||
| `scientific` | spaCy/scispaCy NER |
|
||||
| `tables` | Table Transformer detection |
|
||||
| `figures` | LayoutParser figure/caption detection (needs a Detectron2 build) |
|
||||
| `figures` | LayoutParser figure/caption detection |
|
||||
| `full` | All of the above |
|
||||
| `kg` | Knowledge graph (networkx, txtai, litellm) |
|
||||
| `neo4j` | Neo4j export driver |
|
||||
|
||||
```bash
|
||||
uv run ocr-pipeline setup full # runs: uv sync --extra full + downloads en_core_sci_lg
|
||||
uv sync --extra full
|
||||
# or guided:
|
||||
uv run ocr-pipeline setup full
|
||||
```
|
||||
|
||||
`en_core_sci_lg` is distributed separately from scispaCy. The compatible scispaCy
|
||||
0.5.4 release requires spaCy 3.7.x; if the model install conflicts with your spaCy
|
||||
version, use a dedicated Python 3.11 environment as described in
|
||||
`ocr-pipeline setup full --dry-run` output, and enable `entities` only there.
|
||||
`en_core_sci_lg` is distributed separately from scispaCy and needs spaCy 3.7.x.
|
||||
If the model install conflicts, use a dedicated Python 3.11 venv and enable
|
||||
`entities` only there.
|
||||
|
||||
## kg_ocr (experimental)
|
||||
</details>
|
||||
|
||||
The top-level `kg_ocr` package is a prototype txtai/litellm RAG interface over the
|
||||
pipeline's output. It is not the supported interface (that is the `ocr-pipeline`
|
||||
CLI) and its dependencies are not installed by default:
|
||||
<details>
|
||||
<summary><b>Knowledge graph</b></summary>
|
||||
|
||||
Build a graph from pipeline output: documents, chunks, entities, citations,
|
||||
entity co-occurrence.
|
||||
|
||||
```bash
|
||||
uv sync --extra kg
|
||||
|
||||
uv run ocr-pipeline kg build -d data/ocr_output
|
||||
uv run ocr-pipeline kg stats data/ocr_output/kg_graph.json
|
||||
|
||||
# retrieval
|
||||
uv run ocr-pipeline kg query data/ocr_output/kg_graph.json --entity BRCA1
|
||||
uv run ocr-pipeline kg query data/ocr_output/kg_graph.json --entity BRCA1 --expand
|
||||
uv run ocr-pipeline kg query data/ocr_output/kg_graph.json --citation 10.1038/nature12345
|
||||
|
||||
# export
|
||||
uv run ocr-pipeline kg export data/ocr_output/kg_graph.json -f graphml
|
||||
uv run ocr-pipeline kg export data/ocr_output/kg_graph.json -f neo4j
|
||||
```
|
||||
|
||||
## Docker
|
||||
Anomaly detection flags low-confidence documents, empty chunks, and entity
|
||||
hubs. Neo4j export needs `uv sync --extra neo4j` plus `NEO4J_URI`,
|
||||
`NEO4J_USER`, `NEO4J_PASSWORD` env vars (or `--uri/--user/--password`).
|
||||
|
||||
The image ships the Tesseract-only core with a ready-made config that reads from
|
||||
`/data/screenshots` and writes to `/data/output`:
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>Docker</b></summary>
|
||||
|
||||
The image ships the Tesseract-only core. It reads `/data/screenshots` and
|
||||
writes `/data/output`:
|
||||
|
||||
```bash
|
||||
docker build -t ocr-pipeline .
|
||||
docker run -v ~/Pictures:/data/screenshots -v "$PWD/ocr-output:/data/output" ocr-pipeline run
|
||||
docker run -v ~/Pictures:/data/screenshots -v "$PWD/ocr-output:/data/output" ocr-pipeline watch
|
||||
|
||||
# ML extras baked in:
|
||||
# with ML extras:
|
||||
docker build --build-arg EXTRAS="--extra full" -t ocr-pipeline:full .
|
||||
```
|
||||
|
||||
## Config
|
||||
|
||||
`ocr-pipeline init` writes a minimal working config; the example below shows every
|
||||
knob for reference.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><code>config.yaml</code></summary>
|
||||
<summary><b>Config</b></summary>
|
||||
|
||||
`ocr-pipeline init` writes a minimal config. Lookup order:
|
||||
`$OCR_PIPELINE_CONFIG` → `./config.yaml` → per-user config dir.
|
||||
|
||||
Every knob, for reference:
|
||||
|
||||
```yaml
|
||||
input:
|
||||
@@ -134,20 +162,6 @@ output:
|
||||
organize_by: "date_run" # date_run | source_dir | flat
|
||||
write_consolidated: true
|
||||
consolidated_filename: "all_ocr.md"
|
||||
frontmatter:
|
||||
- source_path
|
||||
- source_hash
|
||||
- timestamp
|
||||
- ocr_engine
|
||||
- ocr_confidence_mean
|
||||
- language
|
||||
- detected_entities
|
||||
- entity_extraction_backend
|
||||
- has_figures
|
||||
- has_tables
|
||||
- citations_found
|
||||
- chunk_index
|
||||
- total_chunks
|
||||
|
||||
watch:
|
||||
enabled: true
|
||||
@@ -157,7 +171,8 @@ watch:
|
||||
|
||||
</details>
|
||||
|
||||
## Output
|
||||
<details>
|
||||
<summary><b>Output example</b></summary>
|
||||
|
||||
Each chunk is a Markdown file with YAML frontmatter:
|
||||
|
||||
@@ -172,71 +187,61 @@ language: "en"
|
||||
detected_entities: ["GENE", "PROTEIN", "CHEMICAL"]
|
||||
has_figures: true
|
||||
has_tables: false
|
||||
citations_found: ["DOI:10.1038/nature12345", "PMID:12345678"]
|
||||
citations_found: ["doi:10.1038/nature12345", "pmid:PMID:12345678"]
|
||||
chunk_index: 0
|
||||
total_chunks: 2
|
||||
---
|
||||
|
||||
# Screenshot: SCR-20250115-gel.png
|
||||
|
||||
## Figures
|
||||
|
||||
### Figure 1
|
||||
- BBox: [100, 200, 800, 600]
|
||||
- Confidence: 0.92
|
||||
- Caption: "Western blot showing BRCA1 expression..."
|
||||
|
||||
## Detected Entities
|
||||
|
||||
- BRCA1
|
||||
- CRISPR
|
||||
- β-actin
|
||||
|
||||
## Citations
|
||||
|
||||
- DOI: 10.1038/nature12345
|
||||
- PMID: 12345678
|
||||
|
||||
## Extracted Text
|
||||
|
||||
**Western Blot Analysis of BRCA1 Expression**
|
||||
|
||||
Lane 1: WT control
|
||||
Lane 2: BRCA1 KO (CRISPR)
|
||||
Lane 3: BRCA1 KO + pBRCA1-WT rescue
|
||||
Lane 4: BRCA1 KO + pBRCA1-C61G mutant
|
||||
|
||||
Anti-BRCA1 (1:1000), Anti-β-actin (1:5000)
|
||||
...
|
||||
```
|
||||
|
||||
Files land in `data/ocr_output/<date>/run_NNN/`. Individual chunk files are retained for RAG indexing. Set `output.write_consolidated: true` to also write `all_ocr.md` containing one frontmatter block and all source text grouped by image.
|
||||
Files land in `<output>/<date>/run_NNN/`. Chunk files are kept for RAG
|
||||
indexing. `output.write_consolidated: true` also writes `all_ocr.md` with one
|
||||
frontmatter block and all source text grouped by image.
|
||||
|
||||
## Life science specifics
|
||||
</details>
|
||||
|
||||
Gene/protein names go through scispaCy's `en_core_sci_lg`. Chemical formulas and units (µM, ng/mL, kb/Mb/Gb, °C, ×g) get normalized, scientific notation gets cleaned up (`1.5×10⁻³` → `1.5×10^-3`), and gel/blot figures get their captions pulled out separately. Citations are matched by regex for DOI, PMID, arXiv, PMC, and ISBN.
|
||||
<details>
|
||||
<summary><b>Models and dependencies</b></summary>
|
||||
|
||||
## Models
|
||||
|
||||
Model weights are only fetched when the matching extra is installed and enabled.
|
||||
First run downloads what it needs, cached in `~/.cache/ocr_pipeline/`:
|
||||
Weights download on first use, cached in `~/.cache/ocr_pipeline/`:
|
||||
|
||||
- PaddleOCR models (~200MB)
|
||||
- scispaCy `en_core_sci_lg` (~800MB)
|
||||
- Table Transformer (~500MB)
|
||||
- LayoutParser PubLayNet (~300MB)
|
||||
|
||||
## Runtime requirements and health
|
||||
Notes per capability:
|
||||
|
||||
The core pipeline can run with Tesseract alone. PaddleOCR, scientific NER, figure detection, and table detection are optional capabilities with heavyweight, platform-specific dependencies. The pipeline records the OCR engine actually used and the entity-extraction backend in generated frontmatter; inspect them after each run rather than assuming configured models loaded.
|
||||
- **PaddleOCR:** pinned to the legacy 2.x API. Falls back to Tesseract when it
|
||||
can't initialize; the fallback is logged and recorded in frontmatter.
|
||||
- **Scientific NER:** without scispaCy + `en_core_sci_lg`, a conservative
|
||||
regex extractor runs instead and the backend is marked accordingly.
|
||||
- **Figures:** needs a Detectron2 build matching your Torch/Python. No
|
||||
universal wheel exists, so it's never installed automatically.
|
||||
- **Tables:** downloaded by Transformers on first use; needs `timm`.
|
||||
|
||||
- **PaddleOCR:** the project pins the legacy 2.x API used by the pipeline. Install the locked environment with `uv sync`; a startup fallback to Tesseract is logged when Paddle cannot initialize.
|
||||
- **Scientific NER:** install a compatible `scispacy` distribution and the separately distributed `en_core_sci_lg` model before enabling production scientific NER. Without it, the pipeline uses its conservative regex fallback and marks the backend accordingly.
|
||||
- **Figures:** `layoutparser`'s `Detectron2LayoutModel` requires a Detectron2 build matching your Torch/Python platform. It is intentionally not forced as a universal dependency because no single wheel supports every platform.
|
||||
- **Tables:** the Table Transformer model is downloaded by Transformers on first use; ensure the selected model's optional dependencies (including `timm`, when required by that model revision) are installed in the runtime image.
|
||||
`ocr-pipeline doctor` checks all of the above and prints remediation hints.
|
||||
|
||||
Use `ocr-pipeline config` to verify effective settings. A missing optional model is logged and produces empty results for that detector instead of failing a complete batch.
|
||||
</details>
|
||||
|
||||
## Life science specifics
|
||||
|
||||
Gene/protein names go through scispaCy's `en_core_sci_lg`. Chemical formulas
|
||||
and units (µM, ng/mL, kb/Mb/Gb, °C, ×g) get normalized, scientific notation
|
||||
gets cleaned up (`1.5×10⁻³` → `1.5×10^-3`), and gel/blot figures get their
|
||||
captions pulled out separately.
|
||||
|
||||
## TODO
|
||||
|
||||
- [ ] Build a knowledge graph from extracted entities/citations
|
||||
|
||||
- [x] Build a knowledge graph from extracted entities/citations (see `kg` commands)
|
||||
- [ ] Entity linking to MeSH/UniProt identifiers
|
||||
- [ ] GROBID-based structured citation parsing (server optional, regex today)
|
||||
|
||||
@@ -1,17 +1,50 @@
|
||||
"""Experimental txtai/litellm RAG interface built on the OCR pipeline.
|
||||
"""Knowledge graph + RAG interface built on the OCR pipeline.
|
||||
|
||||
This package is a prototype. The supported interface is the `ocr-pipeline` CLI
|
||||
(src/ocr_pipeline). Extra dependencies are required: `uv sync --extra kg`.
|
||||
Lightweight submodules (graph, export) import without optional deps.
|
||||
The txtai/litellm RAG pieces need: uv sync --extra kg.
|
||||
"""
|
||||
|
||||
try:
|
||||
from .ocr import get_screenshots, extract_text
|
||||
from .embeddings import create_and_index
|
||||
from .rag import retrieve, ask_wllm
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"kg_ocr is experimental and needs extra dependencies: "
|
||||
"uv sync --extra kg (or: uv pip install txtai litellm python-dotenv)"
|
||||
) from exc
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = ["get_screenshots", "extract_text", "create_and_index", "retrieve", "ask_wllm"]
|
||||
from typing import Any
|
||||
|
||||
_LAZY: dict[str, tuple[str, str]] = {
|
||||
"get_screenshots": (".ocr", "get_screenshots"),
|
||||
"extract_text": (".ocr", "extract_text"),
|
||||
"create_and_index": (".embeddings", "create_and_index"),
|
||||
"retrieve": (".rag", "retrieve"),
|
||||
"ask_wllm": (".rag", "ask_wllm"),
|
||||
"build_graph": (".graph", "build_graph"),
|
||||
"build_from_directory": (".graph", "build_from_directory"),
|
||||
"summary": (".graph", "summary"),
|
||||
"top_entities": (".graph", "top_entities"),
|
||||
"detect_anomalies": (".graph", "detect_anomalies"),
|
||||
"chunks_for_entity": (".graph", "chunks_for_entity"),
|
||||
"chunks_for_citation": (".graph", "chunks_for_citation"),
|
||||
"related_entities": (".graph", "related_entities"),
|
||||
"expand_context": (".graph", "expand_context"),
|
||||
"export_json": (".export", "export_json"),
|
||||
"export_graphml": (".export", "export_graphml"),
|
||||
"load_json": (".export", "load_json"),
|
||||
"Neo4jExporter": (".export", "Neo4jExporter"),
|
||||
}
|
||||
|
||||
__all__ = list(_LAZY)
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name not in _LAZY:
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
module_name, attr = _LAZY[name]
|
||||
try:
|
||||
from importlib import import_module
|
||||
|
||||
module = import_module(module_name, __name__)
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
f"kg_ocr.{name} needs optional dependencies: uv sync --extra kg "
|
||||
f"(graph/export only need networkx; embeddings/rag also need txtai, litellm)"
|
||||
) from exc
|
||||
value = getattr(module, attr)
|
||||
globals()[name] = value
|
||||
return value
|
||||
|
||||
@@ -1,3 +0,0 @@
|
||||
# to do
|
||||
|
||||
# add entrypoint to setuppy
|
||||
@@ -1 +0,0 @@
|
||||
# to do
|
||||
@@ -1 +0,0 @@
|
||||
# to do
|
||||
@@ -1 +1,31 @@
|
||||
# in future prefer a config, especially for graph traversal
|
||||
"""Config for kg_ocr: graph traversal defaults and Neo4j credentials.
|
||||
|
||||
Environment variables: NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class KgConfig:
|
||||
"""Tunables for graph building/traversal plus Neo4j connection."""
|
||||
|
||||
# Traversal
|
||||
expand_limit: int = 5
|
||||
related_limit: int = 10
|
||||
|
||||
# Neo4j
|
||||
neo4j_uri: str = "bolt://localhost:7687"
|
||||
neo4j_user: str = "neo4j"
|
||||
neo4j_password: str = field(default="", repr=False)
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> KgConfig:
|
||||
return cls(
|
||||
neo4j_uri=os.environ.get("NEO4J_URI", cls.neo4j_uri),
|
||||
neo4j_user=os.environ.get("NEO4J_USER", cls.neo4j_user),
|
||||
neo4j_password=os.environ.get("NEO4J_PASSWORD", ""),
|
||||
)
|
||||
|
||||
@@ -1,18 +1,27 @@
|
||||
from txtai.embeddings import Embeddings
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
DEFAULT_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
|
||||
|
||||
|
||||
def create_and_index(
|
||||
data: list[str], model: str = "sentence-transformers/all-MiniLM-L6-v2"
|
||||
) -> Embeddings:
|
||||
"""Create and index embeddings from text."""
|
||||
embeddings = Embeddings({
|
||||
def create_and_index(data: list[str], model: str = DEFAULT_MODEL) -> Any:
|
||||
"""Create and index embeddings from text.
|
||||
|
||||
Requires txtai (uv sync --extra kg). Returns a txtai Embeddings instance.
|
||||
"""
|
||||
try:
|
||||
from txtai.embeddings import Embeddings
|
||||
except ImportError as exc:
|
||||
raise ImportError("create_and_index needs txtai: uv sync --extra kg") from exc
|
||||
|
||||
embeddings = Embeddings(
|
||||
{
|
||||
"path": model,
|
||||
"content": True,
|
||||
"hybrid": True,
|
||||
"scoring": "bm25",
|
||||
})
|
||||
}
|
||||
)
|
||||
embeddings.index(data)
|
||||
return embeddings
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
from .neo4j_exporter import Neo4jExporter, export_graphml, export_json, load_json
|
||||
|
||||
__all__ = ["Neo4jExporter", "export_graphml", "export_json", "load_json"]
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
"""Serialize a knowledge graph to JSON / GraphML, or push it to Neo4j."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
NEO4J_LABELS = {
|
||||
"document": "Document",
|
||||
"chunk": "Chunk",
|
||||
"entity": "Entity",
|
||||
"citation": "Citation",
|
||||
}
|
||||
|
||||
|
||||
def export_json(graph: Any, path: Path) -> Path:
|
||||
"""Node-link JSON (lossless round-trip with `load_json`)."""
|
||||
import networkx as nx
|
||||
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
data = nx.node_link_data(graph, edges="edges")
|
||||
path.write_text(json.dumps(data, indent=2, default=str), encoding="utf-8")
|
||||
return path
|
||||
|
||||
|
||||
def load_json(path: Path) -> Any:
|
||||
"""Load a graph previously written by `export_json`."""
|
||||
import networkx as nx
|
||||
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
return nx.node_link_graph(data, edges="edges")
|
||||
|
||||
|
||||
def export_graphml(graph: Any, path: Path) -> Path:
|
||||
"""GraphML for Gephi/Cytoscape. Chunk text is dropped (GraphML attr limits)."""
|
||||
import networkx as nx
|
||||
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
slim = nx.DiGraph()
|
||||
for node, data in graph.nodes(data=True):
|
||||
slim.add_node(node, **{k: v for k, v in data.items() if k != "text"})
|
||||
slim.add_edges_from(
|
||||
(u, v, {k: val for k, val in data.items() if isinstance(val, (int, float, str, bool))})
|
||||
for u, v, data in graph.edges(data=True)
|
||||
)
|
||||
nx.write_graphml(slim, str(path))
|
||||
return path
|
||||
|
||||
|
||||
class Neo4jExporter:
|
||||
"""MERGE the graph into a Neo4j database.
|
||||
|
||||
Credentials come from constructor args or NEO4J_URI / NEO4J_USER /
|
||||
NEO4J_PASSWORD environment variables.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
uri: str | None = None,
|
||||
user: str | None = None,
|
||||
password: str | None = None,
|
||||
) -> None:
|
||||
import os
|
||||
|
||||
self.uri = uri or os.environ.get("NEO4J_URI", "bolt://localhost:7687")
|
||||
self.user = user or os.environ.get("NEO4J_USER", "neo4j")
|
||||
self.password = password or os.environ.get("NEO4J_PASSWORD", "")
|
||||
try:
|
||||
from neo4j import GraphDatabase
|
||||
except ImportError as exc:
|
||||
raise ImportError("Neo4j export needs the driver: uv pip install neo4j") from exc
|
||||
self._driver = GraphDatabase.driver(self.uri, auth=(self.user, self.password))
|
||||
|
||||
def close(self) -> None:
|
||||
self._driver.close()
|
||||
|
||||
def __enter__(self) -> Neo4jExporter:
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc_info: object) -> None:
|
||||
self.close()
|
||||
|
||||
def push(self, graph: Any, batch_size: int = 500) -> dict[str, int]:
|
||||
"""MERGE all nodes and edges. Returns counts pushed."""
|
||||
nodes = list(graph.nodes(data=True))
|
||||
edges = list(graph.edges(data=True))
|
||||
with self._driver.session() as session:
|
||||
for start in range(0, len(nodes), batch_size):
|
||||
session.execute_write(self._push_nodes, nodes[start : start + batch_size])
|
||||
for start in range(0, len(edges), batch_size):
|
||||
session.execute_write(self._push_edges, edges[start : start + batch_size])
|
||||
return {"nodes": len(nodes), "edges": len(edges)}
|
||||
|
||||
@staticmethod
|
||||
def _push_nodes(tx: Any, batch: list[tuple[str, dict]]) -> None:
|
||||
for node_id, data in batch:
|
||||
label = NEO4J_LABELS.get(data.get("kind", ""), "Node")
|
||||
props = {k: v for k, v in data.items() if k != "kind"}
|
||||
tx.run(
|
||||
f"MERGE (n:{label} {{id: $id}}) SET n += $props",
|
||||
id=node_id,
|
||||
props=props,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _push_edges(tx: Any, batch: list[tuple[str, str, dict]]) -> None:
|
||||
for source, target, data in batch:
|
||||
rel = data.get("kind", "RELATED")
|
||||
props = {k: v for k, v in data.items() if k != "kind"}
|
||||
tx.run(
|
||||
f"MATCH (a {{id: $source}}), (b {{id: $target}}) "
|
||||
f"MERGE (a)-[r:{rel}]->(b) SET r += $props",
|
||||
source=source,
|
||||
target=target,
|
||||
props=props,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
@@ -1,4 +1,4 @@
|
||||
from .extractor import get_screenshots
|
||||
from .batch_processor import extract_text
|
||||
from .extractor import get_screenshots
|
||||
|
||||
__all__ = ["get_screenshots", "extract_text"]
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
from PIL import Image
|
||||
import pytesseract
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def extract_text(images: list[str]) -> list[str]:
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
from pathlib import Path
|
||||
|
||||
import platform
|
||||
|
||||
def_paths = {
|
||||
"Darwin": Path.home() / "Desktop",
|
||||
"Windows": Path.home() / "Pictures" / "Screenshots",
|
||||
|
||||
@@ -1,11 +1,10 @@
|
||||
import platform
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from .constants import def_paths, sc_pathpatterns
|
||||
|
||||
|
||||
def get_screenshots(path: Optional[str | Path] = None) -> list[str]:
|
||||
def get_screenshots(path: str | Path | None = None) -> list[str]:
|
||||
"""Find screenshot files for the current OS."""
|
||||
if path is None:
|
||||
path = def_paths.get(platform.system(), Path.home())
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
from .query import retrieve, ask_wllm
|
||||
from .query import ask_wllm, retrieve
|
||||
|
||||
__all__ = ["retrieve", "ask_wllm"]
|
||||
|
||||
@@ -1,32 +1,47 @@
|
||||
from txtai.embeddings import Embeddings
|
||||
from txtai import LLM
|
||||
import litellm
|
||||
from dotenv import load_dotenv
|
||||
import os
|
||||
from __future__ import annotations
|
||||
|
||||
load_dotenv()
|
||||
from typing import Any
|
||||
|
||||
def retrieve(embeddings: Embeddings, query: str, limit: int = 3) -> list[dict]:
|
||||
"""Search embeddings and return results with scores"""
|
||||
DEFAULT_CHAT_MODEL = "openrouter/minimax/minimax-m2.5:free"
|
||||
|
||||
SYSTEM_PROMPT = (
|
||||
"Answer ONLY using the provided context. Cite which parts you're drawing "
|
||||
"from. If the context doesn't cover something, say 'not in my documents'."
|
||||
)
|
||||
|
||||
|
||||
def retrieve(embeddings: Any, query: str, limit: int = 3) -> list[dict]:
|
||||
"""Search embeddings and return results with scores."""
|
||||
return embeddings.search(query, limit)
|
||||
|
||||
def ask_wllm(embeddings: Embeddings, question: str, model: str = "openrouter/minimax/minimax-m2.5:free", limit: int = 3) -> str:
|
||||
|
||||
def ask_wllm(
|
||||
embeddings: Any,
|
||||
question: str,
|
||||
model: str = DEFAULT_CHAT_MODEL,
|
||||
limit: int = 3,
|
||||
) -> str:
|
||||
"""RAG: retrieve context from embeddings, then answer with an LLM."""
|
||||
try:
|
||||
import litellm
|
||||
except ImportError as exc:
|
||||
raise ImportError("ask_wllm needs litellm: uv sync --extra kg") from exc
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
results = retrieve(embeddings, question, limit)
|
||||
context = "\n\n".join([r["text"] for r in results])
|
||||
|
||||
response = litellm.completion(
|
||||
model=model,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "Answer ONLY using the provided context. Cite which parts you're drawing from. If the context doesn't cover something, say 'not in my documents'."
|
||||
},
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"Context from my documents:\n{context}\n\nQuestion: {question}"
|
||||
}
|
||||
]
|
||||
"content": f"Context from my documents:\n{context}\n\nQuestion: {question}",
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
return response.choices[0].message.content
|
||||
@@ -63,10 +63,14 @@ scientific = [
|
||||
"spacy>=3.7.0; python_version >= '3.12'",
|
||||
]
|
||||
kg = [
|
||||
"networkx>=3.2",
|
||||
"txtai>=7.4.0",
|
||||
"litellm>=1.40.0",
|
||||
"python-dotenv>=1.0.0",
|
||||
]
|
||||
neo4j = [
|
||||
"neo4j>=5.0",
|
||||
]
|
||||
grobid = []
|
||||
full = [
|
||||
"ocr-pipeline[paddle,tables,figures,scientific]",
|
||||
|
||||
Binary file not shown.
@@ -28,10 +28,27 @@ app = typer.Typer(
|
||||
help="OCR pipeline for screenshots -> RAG-ready Markdown",
|
||||
add_completion=False,
|
||||
)
|
||||
kg_app = typer.Typer(
|
||||
name="kg",
|
||||
help="Knowledge graph over OCR output (needs: uv sync --extra kg)",
|
||||
add_completion=False,
|
||||
)
|
||||
app.add_typer(kg_app, name="kg")
|
||||
console = Console()
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _load_kg():
|
||||
"""Import kg_ocr lazily so the base CLI works without the kg extra."""
|
||||
try:
|
||||
from kg_ocr import export as kg_export
|
||||
from kg_ocr import graph as kg_graph
|
||||
except ImportError as exc:
|
||||
console.print(f"[red]{exc}[/red]")
|
||||
raise typer.Exit(code=2) from exc
|
||||
return kg_graph, kg_export
|
||||
|
||||
|
||||
def _ensure_configured(input_dir: Path | None) -> None:
|
||||
"""Stop early with guidance instead of scanning a default directory."""
|
||||
if input_dir is not None or find_config_file() is not None:
|
||||
@@ -296,6 +313,169 @@ def config():
|
||||
console.print(settings.model_dump_json(indent=2))
|
||||
|
||||
|
||||
@kg_app.command("build")
|
||||
def kg_build(
|
||||
output_dir: Path = typer.Option(
|
||||
..., "--output-dir", "-d", help="Pipeline output directory with chunk markdown files"
|
||||
),
|
||||
save: Path | None = typer.Option(
|
||||
None, "--save", "-s", help="Graph JSON path (default: <output-dir>/kg_graph.json)"
|
||||
),
|
||||
):
|
||||
"""Build a knowledge graph from pipeline output and save it as JSON."""
|
||||
kg_graph, kg_export = _load_kg()
|
||||
if not output_dir.is_dir():
|
||||
console.print(f"[red]Not a directory:[/red] {output_dir}")
|
||||
raise typer.Exit(code=2)
|
||||
graph = kg_graph.build_from_directory(output_dir)
|
||||
save = save or (output_dir / "kg_graph.json")
|
||||
kg_export.export_json(graph, save)
|
||||
|
||||
counts = kg_graph.summary(graph)
|
||||
table = Table(title=f"Knowledge graph -> {save}")
|
||||
table.add_column("Kind", style="cyan")
|
||||
table.add_column("Count", style="green")
|
||||
for kind, count in sorted(counts.items()):
|
||||
table.add_row(kind, str(count))
|
||||
console.print(table)
|
||||
|
||||
|
||||
@kg_app.command("stats")
|
||||
def kg_stats(
|
||||
graph_path: Path = typer.Argument(..., help="Graph JSON written by `kg build`"),
|
||||
top: int = typer.Option(10, "--top", "-n", help="Rows per top-list"),
|
||||
):
|
||||
"""Summarize a graph: counts, top entities/citations, anomalies."""
|
||||
kg_graph, kg_export = _load_kg()
|
||||
if not graph_path.is_file():
|
||||
console.print(f"[red]Graph file not found:[/red] {graph_path}")
|
||||
raise typer.Exit(code=2)
|
||||
graph = kg_export.load_json(graph_path)
|
||||
|
||||
counts = kg_graph.summary(graph)
|
||||
table = Table(title="Graph summary")
|
||||
table.add_column("Kind", style="cyan")
|
||||
table.add_column("Count", style="green")
|
||||
for kind, count in sorted(counts.items()):
|
||||
table.add_row(kind, str(count))
|
||||
console.print(table)
|
||||
|
||||
entities = kg_graph.top_entities(graph, limit=top)
|
||||
if entities:
|
||||
table = Table(title=f"Top {top} entities")
|
||||
table.add_column("Entity", style="cyan")
|
||||
table.add_column("Mentions", style="green")
|
||||
for text, mentions in entities:
|
||||
table.add_row(text, str(mentions))
|
||||
console.print(table)
|
||||
|
||||
citations = kg_graph.top_citations(graph, limit=top)
|
||||
if citations:
|
||||
table = Table(title=f"Top {top} citations")
|
||||
table.add_column("Type", style="cyan")
|
||||
table.add_column("Identifier")
|
||||
table.add_column("Citing chunks", style="green")
|
||||
for ctype, identifier, citing in citations:
|
||||
table.add_row(ctype, identifier, str(citing))
|
||||
console.print(table)
|
||||
|
||||
anomalies = kg_graph.detect_anomalies(graph)
|
||||
if anomalies:
|
||||
console.print(f"\n[yellow]{len(anomalies)} anomalies:[/yellow]")
|
||||
for anomaly in anomalies[:top]:
|
||||
console.print(f" [dim]{anomaly.kind}[/dim] {anomaly.node}: {anomaly.detail}")
|
||||
|
||||
|
||||
@kg_app.command("query")
|
||||
def kg_query(
|
||||
graph_path: Path = typer.Argument(..., help="Graph JSON written by `kg build`"),
|
||||
entity: str | None = typer.Option(None, "--entity", "-e", help="Entity to look up"),
|
||||
citation: str | None = typer.Option(
|
||||
None, "--citation", "-c", help="Citation identifier (e.g. 10.1038/nature12345)"
|
||||
),
|
||||
expand: bool = typer.Option(
|
||||
False, "--expand", "-x", help="Include chunks from co-occurring entities"
|
||||
),
|
||||
limit: int = typer.Option(5, "--limit", "-n", help="Max chunks to show"),
|
||||
):
|
||||
"""Retrieve chunks by entity or citation; --expand adds neighbor chunks."""
|
||||
kg_graph, kg_export = _load_kg()
|
||||
if not graph_path.is_file():
|
||||
console.print(f"[red]Graph file not found:[/red] {graph_path}")
|
||||
raise typer.Exit(code=2)
|
||||
if not entity and not citation:
|
||||
console.print("[red]Give --entity or --citation.[/red]")
|
||||
raise typer.Exit(code=2)
|
||||
graph = kg_export.load_json(graph_path)
|
||||
|
||||
if entity:
|
||||
chunks = (
|
||||
kg_graph.expand_context(graph, entity, limit=limit)
|
||||
if expand
|
||||
else kg_graph.chunks_for_entity(graph, entity)
|
||||
)
|
||||
related = kg_graph.related_entities(graph, entity)
|
||||
if related:
|
||||
console.print(
|
||||
"[dim]Related entities: "
|
||||
+ ", ".join(f"{text} ({weight})" for text, weight in related[:5])
|
||||
+ "[/dim]"
|
||||
)
|
||||
else:
|
||||
chunks = kg_graph.chunks_for_citation(graph, citation or "")
|
||||
|
||||
if not chunks:
|
||||
console.print("[yellow]No matching chunks.[/yellow]")
|
||||
return
|
||||
|
||||
for payload in chunks[:limit]:
|
||||
via = f" via {payload['via_entity']}" if payload.get("via_entity") else ""
|
||||
console.print(
|
||||
f"\n[bold]{Path(str(payload['source_path'])).name}[/bold] "
|
||||
f"chunk {payload['chunk_index']}{via} "
|
||||
f"[dim]({payload['ocr_engine']} {payload['ocr_confidence_mean']:.2f})[/dim]"
|
||||
)
|
||||
excerpt = " ".join(str(payload["text"]).split())[:300]
|
||||
console.print(f" {excerpt}")
|
||||
|
||||
|
||||
@kg_app.command("export")
|
||||
def kg_export_cmd(
|
||||
graph_path: Path = typer.Argument(..., help="Graph JSON written by `kg build`"),
|
||||
fmt: str = typer.Option("graphml", "--format", "-f", help="graphml | neo4j"),
|
||||
out: Path | None = typer.Option(None, "--out", "-o", help="Output path for graphml"),
|
||||
uri: str | None = typer.Option(None, "--uri", help="Neo4j bolt URI (or NEO4J_URI)"),
|
||||
user: str | None = typer.Option(None, "--user", help="Neo4j user (or NEO4J_USER)"),
|
||||
password: str | None = typer.Option(
|
||||
None, "--password", help="Neo4j password (or NEO4J_PASSWORD)"
|
||||
),
|
||||
):
|
||||
"""Export a graph JSON to GraphML or push it into Neo4j."""
|
||||
_, kg_export = _load_kg()
|
||||
if not graph_path.is_file():
|
||||
console.print(f"[red]Graph file not found:[/red] {graph_path}")
|
||||
raise typer.Exit(code=2)
|
||||
graph = kg_export.load_json(graph_path)
|
||||
|
||||
if fmt == "graphml":
|
||||
out = out or graph_path.with_suffix(".graphml")
|
||||
kg_export.export_graphml(graph, out)
|
||||
console.print(f"[green]GraphML written:[/green] {out}")
|
||||
elif fmt == "neo4j":
|
||||
try:
|
||||
with kg_export.Neo4jExporter(uri=uri, user=user, password=password) as exporter:
|
||||
pushed = exporter.push(graph)
|
||||
except ImportError as exc:
|
||||
console.print(f"[red]{exc}[/red]")
|
||||
raise typer.Exit(code=2) from exc
|
||||
console.print(
|
||||
f"[green]Pushed to Neo4j:[/green] {pushed['nodes']} nodes, {pushed['edges']} edges"
|
||||
)
|
||||
else:
|
||||
console.print(f"[red]Unknown format:[/red] {fmt} (choose graphml or neo4j)")
|
||||
raise typer.Exit(code=2)
|
||||
|
||||
|
||||
def main():
|
||||
app()
|
||||
|
||||
|
||||
@@ -1,5 +0,0 @@
|
||||
"""Compatibility module for the canonical watcher package."""
|
||||
|
||||
from .watch.watcher import ScreenshotHandler, Watcher
|
||||
|
||||
__all__ = ["ScreenshotHandler", "Watcher"]
|
||||
@@ -0,0 +1,267 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
nx = pytest.importorskip("networkx", reason="kg extra not installed")
|
||||
|
||||
from kg_ocr.export import export_graphml, export_json, load_json # noqa: E402
|
||||
from kg_ocr.graph import ( # noqa: E402
|
||||
build_from_directory,
|
||||
build_graph,
|
||||
detect_anomalies,
|
||||
find_chunk_files,
|
||||
parse_chunk_file,
|
||||
summary,
|
||||
top_citations,
|
||||
top_cooccurrences,
|
||||
top_entities,
|
||||
)
|
||||
|
||||
CHUNK_MD = """---
|
||||
source_path: /shots/SCR-001.png
|
||||
source_hash: aaaabbbbccccdddd
|
||||
timestamp: '2026-07-19T10:00:00+00:00'
|
||||
ocr_engine: paddleocr
|
||||
ocr_confidence_mean: 0.92
|
||||
language: en
|
||||
detected_entities:
|
||||
- BRCA1
|
||||
- PARP
|
||||
citations_found:
|
||||
- doi:10.1038/nature12345
|
||||
chunk_index: 0
|
||||
total_chunks: 1
|
||||
---
|
||||
|
||||
# Screenshot: SCR-001.png
|
||||
|
||||
## Extracted Text
|
||||
|
||||
BRCA1 and PARP inhibition in ovarian cancer.
|
||||
|
||||
## Citations Found
|
||||
|
||||
| Type | Identifier |
|
||||
"""
|
||||
|
||||
CHUNK_MD_2 = """---
|
||||
source_path: /shots/SCR-002.png
|
||||
source_hash: eeeeffff00001111
|
||||
timestamp: '2026-07-19T10:01:00+00:00'
|
||||
ocr_engine: tesseract
|
||||
ocr_confidence_mean: 0.2
|
||||
language: en
|
||||
detected_entities:
|
||||
- BRCA1
|
||||
citations_found: []
|
||||
chunk_index: 0
|
||||
total_chunks: 1
|
||||
---
|
||||
|
||||
# Screenshot: SCR-002.png
|
||||
|
||||
## Extracted Text
|
||||
|
||||
BRCA1 western blot.
|
||||
"""
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def chunk_dir(tmp_path: Path) -> Path:
|
||||
run_dir = tmp_path / "out" / "2026-07-19" / "run_001"
|
||||
run_dir.mkdir(parents=True)
|
||||
(run_dir / "SCR-001_aaaabbbbcccc.md").write_text(CHUNK_MD, encoding="utf-8")
|
||||
(run_dir / "SCR-002_eeeeffff0000.md").write_text(CHUNK_MD_2, encoding="utf-8")
|
||||
return tmp_path / "out"
|
||||
|
||||
|
||||
def test_parse_chunk_file(chunk_dir: Path) -> None:
|
||||
path = next(chunk_dir.rglob("SCR-001_*.md"))
|
||||
record = parse_chunk_file(path)
|
||||
assert record is not None
|
||||
assert record.source_hash == "aaaabbbbccccdddd"
|
||||
assert record.ocr_confidence_mean == pytest.approx(0.92)
|
||||
assert record.entities == ["BRCA1", "PARP"]
|
||||
assert record.citations == ["doi:10.1038/nature12345"]
|
||||
assert "BRCA1 and PARP" in record.text
|
||||
|
||||
|
||||
def test_parse_chunk_file_rejects_plain_markdown(tmp_path: Path) -> None:
|
||||
plain = tmp_path / "notes.md"
|
||||
plain.write_text("# just notes\n", encoding="utf-8")
|
||||
assert parse_chunk_file(plain) is None
|
||||
|
||||
|
||||
def test_find_chunk_files_skips_consolidated(chunk_dir: Path) -> None:
|
||||
(chunk_dir / "all_ocr.md").write_text("# consolidated\n", encoding="utf-8")
|
||||
names = [p.name for p in find_chunk_files(chunk_dir)]
|
||||
assert "all_ocr.md" not in names
|
||||
assert len(names) == 2
|
||||
|
||||
|
||||
def test_build_graph_structure(chunk_dir: Path) -> None:
|
||||
graph = build_from_directory(chunk_dir)
|
||||
counts = summary(graph)
|
||||
|
||||
assert counts["nodes:document"] == 2
|
||||
assert counts["nodes:chunk"] == 2
|
||||
assert counts["nodes:entity"] == 2 # brca1, parp (normalized)
|
||||
assert counts["nodes:citation"] == 1
|
||||
assert counts["edges:CONTAINS"] == 2
|
||||
assert counts["edges:MENTIONS"] == 3
|
||||
assert counts["edges:CITES"] == 1
|
||||
assert counts["edges:CO_OCCURS"] == 1 # brca1 <-> parp
|
||||
|
||||
assert graph.has_edge("doc:aaaabbbbccccdddd", "chunk:aaaabbbbccccdddd:0")
|
||||
assert graph.has_edge("chunk:aaaabbbbccccdddd:0", "entity:brca1")
|
||||
|
||||
|
||||
def test_top_entities_and_cooccurrence(chunk_dir: Path) -> None:
|
||||
graph = build_from_directory(chunk_dir)
|
||||
top = top_entities(graph)
|
||||
assert top[0] == ("BRCA1", 2)
|
||||
pairs = top_cooccurrences(graph)
|
||||
assert pairs == [("BRCA1", "PARP", 1)]
|
||||
|
||||
|
||||
def test_top_citations(chunk_dir: Path) -> None:
|
||||
graph = build_from_directory(chunk_dir)
|
||||
rows = top_citations(graph)
|
||||
assert rows == [("doi", "10.1038/nature12345", 1)]
|
||||
|
||||
|
||||
def test_detect_anomalies(chunk_dir: Path) -> None:
|
||||
graph = build_from_directory(chunk_dir)
|
||||
kinds = {a.kind for a in detect_anomalies(graph)}
|
||||
# SCR-002 has confidence 0.2 (< 0.4) -> low_confidence
|
||||
assert "low_confidence" in kinds
|
||||
|
||||
|
||||
def test_json_round_trip(chunk_dir: Path, tmp_path: Path) -> None:
|
||||
graph = build_from_directory(chunk_dir)
|
||||
out = export_json(graph, tmp_path / "g.json")
|
||||
restored = load_json(out)
|
||||
assert summary(restored) == summary(graph)
|
||||
assert restored.nodes["entity:brca1"]["mentions"] == 2
|
||||
|
||||
|
||||
def test_graphml_export(chunk_dir: Path, tmp_path: Path) -> None:
|
||||
graph = build_from_directory(chunk_dir)
|
||||
out = export_graphml(graph, tmp_path / "g.graphml")
|
||||
content = out.read_text(encoding="utf-8")
|
||||
assert "graphml" in content
|
||||
|
||||
|
||||
def test_empty_directory(tmp_path: Path) -> None:
|
||||
graph = build_from_directory(tmp_path)
|
||||
assert summary(graph) == {}
|
||||
|
||||
|
||||
def test_build_graph_from_records_directly() -> None:
|
||||
from kg_ocr.graph import ChunkRecord
|
||||
|
||||
record = ChunkRecord(
|
||||
source_hash="x" * 16,
|
||||
source_path="/a.png",
|
||||
chunk_index=0,
|
||||
total_chunks=1,
|
||||
text="hello",
|
||||
entities=["TP53", "tp53"], # dedupe via normalization
|
||||
citations=["pmid:PMID:12345678"],
|
||||
)
|
||||
graph = build_graph([record])
|
||||
counts = summary(graph)
|
||||
assert counts["nodes:entity"] == 1
|
||||
assert graph.nodes["entity:tp53"]["mentions"] == 2
|
||||
assert graph.nodes["citation:PMID:12345678"]["citation_type"] == "pmid"
|
||||
|
||||
|
||||
def test_kg_cli_build_and_stats(chunk_dir: Path, tmp_path: Path) -> None:
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from ocr_pipeline.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
save = tmp_path / "g.json"
|
||||
|
||||
result = runner.invoke(app, ["kg", "build", "-d", str(chunk_dir), "-s", str(save)])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert save.is_file()
|
||||
|
||||
result = runner.invoke(app, ["kg", "stats", str(save)])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "BRCA1" in result.output
|
||||
assert "low_confidence" in result.output
|
||||
|
||||
out_graphml = tmp_path / "g.graphml"
|
||||
result = runner.invoke(
|
||||
app, ["kg", "export", str(save), "-f", "graphml", "-o", str(out_graphml)]
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert out_graphml.is_file()
|
||||
|
||||
|
||||
def test_chunks_for_entity(chunk_dir: Path) -> None:
|
||||
from kg_ocr.graph import chunks_for_entity
|
||||
|
||||
graph = build_from_directory(chunk_dir)
|
||||
chunks = chunks_for_entity(graph, "brca1") # case-insensitive
|
||||
assert len(chunks) == 2
|
||||
assert {c["source_path"] for c in chunks} == {"/shots/SCR-001.png", "/shots/SCR-002.png"}
|
||||
assert chunks_for_entity(graph, "NONEXISTENT") == []
|
||||
|
||||
|
||||
def test_chunks_for_citation(chunk_dir: Path) -> None:
|
||||
from kg_ocr.graph import chunks_for_citation
|
||||
|
||||
graph = build_from_directory(chunk_dir)
|
||||
chunks = chunks_for_citation(graph, "10.1038/nature12345")
|
||||
assert len(chunks) == 1
|
||||
assert "BRCA1 and PARP" in chunks[0]["text"]
|
||||
assert chunks_for_citation(graph, "nope") == []
|
||||
|
||||
|
||||
def test_related_entities(chunk_dir: Path) -> None:
|
||||
from kg_ocr.graph import related_entities
|
||||
|
||||
graph = build_from_directory(chunk_dir)
|
||||
assert related_entities(graph, "BRCA1") == [("PARP", 1)]
|
||||
assert related_entities(graph, "PARP") == [("BRCA1", 1)]
|
||||
assert related_entities(graph, "NONEXISTENT") == []
|
||||
|
||||
|
||||
def test_expand_context(chunk_dir: Path) -> None:
|
||||
from kg_ocr.graph import expand_context
|
||||
|
||||
graph = build_from_directory(chunk_dir)
|
||||
chunks = expand_context(graph, "PARP")
|
||||
# PARP's own chunk + BRCA1 neighbor chunks (SCR-001 deduped)
|
||||
assert len(chunks) == 2
|
||||
via = {c["via_entity"] for c in chunks}
|
||||
assert "PARP" in via
|
||||
assert "BRCA1" in via
|
||||
|
||||
|
||||
def test_kg_cli_query(chunk_dir: Path, tmp_path: Path) -> None:
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from ocr_pipeline.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
save = tmp_path / "g.json"
|
||||
result = runner.invoke(app, ["kg", "build", "-d", str(chunk_dir), "-s", str(save)])
|
||||
assert result.exit_code == 0, result.output
|
||||
|
||||
result = runner.invoke(app, ["kg", "query", str(save), "--entity", "PARP"])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "SCR-001.png" in result.output
|
||||
assert "Related entities" in result.output
|
||||
|
||||
result = runner.invoke(app, ["kg", "query", str(save), "--citation", "10.1038/nature12345"])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "SCR-001.png" in result.output
|
||||
|
||||
result = runner.invoke(app, ["kg", "query", str(save)])
|
||||
assert result.exit_code == 2 # neither --entity nor --citation
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
from kg_ocr.embeddings import create_and_index
|
||||
from kg_ocr.rag import retrieve
|
||||
|
||||
|
||||
class _FakeEmbeddings:
|
||||
instances: list[_FakeEmbeddings] = []
|
||||
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
self.indexed: list[str] | None = None
|
||||
_FakeEmbeddings.instances.append(self)
|
||||
|
||||
def index(self, data):
|
||||
self.indexed = data
|
||||
|
||||
|
||||
def _install_fake_txtai(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
fake_txtai = types.ModuleType("txtai")
|
||||
fake_embeddings_mod = types.ModuleType("txtai.embeddings")
|
||||
fake_embeddings_mod.Embeddings = _FakeEmbeddings # type: ignore[attr-defined]
|
||||
monkeypatch.setitem(sys.modules, "txtai", fake_txtai)
|
||||
monkeypatch.setitem(sys.modules, "txtai.embeddings", fake_embeddings_mod)
|
||||
|
||||
|
||||
def test_create_and_index_config_and_data(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_install_fake_txtai(monkeypatch)
|
||||
emb = create_and_index(["alpha", "beta"], model="some-model")
|
||||
assert emb.config["path"] == "some-model"
|
||||
assert emb.config["content"] is True
|
||||
assert emb.config["hybrid"] is True
|
||||
assert emb.indexed == ["alpha", "beta"]
|
||||
|
||||
|
||||
def test_create_and_index_raises_without_txtai(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setitem(sys.modules, "txtai", None)
|
||||
monkeypatch.setitem(sys.modules, "txtai.embeddings", None)
|
||||
with pytest.raises(ImportError, match="txtai"):
|
||||
create_and_index(["x"])
|
||||
|
||||
|
||||
def test_retrieve_passes_through() -> None:
|
||||
class FakeEmb:
|
||||
def search(self, query, limit):
|
||||
return [{"text": f"{query}-{i}", "score": 1.0} for i in range(limit)]
|
||||
|
||||
out = retrieve(FakeEmb(), "q", limit=2)
|
||||
assert out == [{"text": "q-0", "score": 1.0}, {"text": "q-1", "score": 1.0}]
|
||||
|
||||
27
uv.lock
generated
27
uv.lock
generated
@@ -2349,6 +2349,18 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a9/82/0340caa499416c78e5d8f5f05947ae4bc3cba53c9f038ab6e9ed964e22f1/nbformat-5.10.4-py3-none-any.whl", hash = "sha256:3b48d6c8fbca4b299bf3982ea7db1af21580e4fec269ad087b9e81588891200b", size = 78454 },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "neo4j"
|
||||
version = "6.2.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pytz" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/ce/f4/aaa4ac19adae4b01bc742b63afd2672a77e7351566f02721e713e4b863ee/neo4j-6.2.0.tar.gz", hash = "sha256:e1e246b65b572bd8ea97f9e0e721b7d40a5ce53e53d0007c29aef63e4f9124d9", size = 241459 }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/e6/cf/1c3795866cefaac6e648d4e98c373cafd97810f6e317c307371007ab4abb/neo4j-6.2.0-py3-none-any.whl", hash = "sha256:b87abdd13a5cc2e3bd51026926c2f20ac38fa3febe98c340520dce19e97388d0", size = 327824 },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "networkx"
|
||||
version = "3.6.1"
|
||||
@@ -2692,9 +2704,13 @@ full = [
|
||||
]
|
||||
kg = [
|
||||
{ name = "litellm" },
|
||||
{ name = "networkx" },
|
||||
{ name = "python-dotenv" },
|
||||
{ name = "txtai" },
|
||||
]
|
||||
neo4j = [
|
||||
{ name = "neo4j" },
|
||||
]
|
||||
paddle = [
|
||||
{ name = "paddleocr" },
|
||||
{ name = "paddlepaddle" },
|
||||
@@ -2728,6 +2744,8 @@ requires-dist = [
|
||||
{ name = "litellm", marker = "extra == 'kg'", specifier = ">=1.40.0" },
|
||||
{ name = "mypy", marker = "extra == 'dev'", specifier = ">=1.8.0" },
|
||||
{ name = "nbformat", specifier = ">=5.9.0" },
|
||||
{ name = "neo4j", marker = "extra == 'neo4j'", specifier = ">=5.0" },
|
||||
{ name = "networkx", marker = "extra == 'kg'", specifier = ">=3.2" },
|
||||
{ name = "numpy", marker = "python_full_version < '3.12'", specifier = ">=1.26.0,<2.0" },
|
||||
{ name = "numpy", marker = "python_full_version >= '3.12'", specifier = ">=1.26.0" },
|
||||
{ name = "ocr-pipeline", extras = ["paddle", "tables", "figures", "scientific"], marker = "extra == 'full'" },
|
||||
@@ -3863,6 +3881,15 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/0b/d7/1959b9648791274998a9c3526f6d0ec8fd2233e4d4acce81bbae76b44b2a/python_dotenv-1.2.2-py3-none-any.whl", hash = "sha256:1d8214789a24de455a8b8bd8ae6fe3c6b69a5e3d64aa8a8e5d68e694bbcb285a", size = 22101 },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytz"
|
||||
version = "2026.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/ff/46/dd499ec9038423421951e4fad73051febaa13d2df82b4064f87af8b8c0c3/pytz-2026.2.tar.gz", hash = "sha256:0e60b47b29f21574376f218fe21abc009894a2321ea16c6754f3cad6eb7cdd6a", size = 320861 }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/ec/dd/96da98f892250475bdf2328112d7468abdd4acc7b902b6af23f4ed958ea0/pytz-2026.2-py2.py3-none-any.whl", hash = "sha256:04156e608bee23d3792fd45c94ae47fae1036688e75032eea2e3bf0323d1f126", size = 510141 },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pywin32"
|
||||
version = "312"
|
||||
|
||||
Reference in New Issue
Block a user