Metadata-Version: 2.4
Name: document2graph
Version: 0.1.1
Summary: Extract hierarchical document graphs and baseline chunks from PDF
Author-email: Marina Walther <marina.walther@gmail.com>
License: MIT
Project-URL: Homepage, https://github.com/mwalther10/document2graph
Project-URL: Issues, https://github.com/mwalther10/document2graph/issues
Keywords: pdf,document2graph,docling,networkx,chunking,retrieval
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Text Processing :: Indexing
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: docling
Requires-Dist: docling-core
Requires-Dist: docling-parse
Requires-Dist: numpy
Requires-Dist: pydantic>=2
Requires-Dist: regex
Requires-Dist: networkx
Requires-Dist: coloredlogs
Requires-Dist: tqdm
Requires-Dist: transformers
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: httpx; extra == "dev"
Provides-Extra: embeddings
Requires-Dist: sentence-transformers; extra == "embeddings"
Dynamic: license-file

# document2graph

Extract hierarchical document graphs and baseline chunks from PDF files. The package parses PDFs using [Docling](https://github.com/DS4SD/docling) and produces two complementary representations:

- **Document graph** — text, image, and table nodes connected by weighted structural edges, capturing the layout hierarchy of the document. Every graph is guaranteed to be connected: nodes without a resolvable parent are attached to a synthetic document root node.
- **Baseline chunks** — flat, semantically merged chunks suitable for retrieval pipelines, with optional context enrichment.

## Installation

```bash
pip install document2graph
```

Or from source:

```bash
pip install -e .
```

For development dependencies (pytest):

```bash
pip install -e ".[dev]"
```

Requires Python 3.11+.

## Quick start

### Configuration

Both extractors share the same `ExtractorConfig`:

```python
from document2graph.models import ExtractorConfig
from docling.datamodel.pipeline_options import PdfPipelineOptions

config = ExtractorConfig(
    pdf_path="data/pdfs/",        # directory containing input PDFs
    data_path="data/output/",     # root directory for all output files
    save_json=True,               # write results to disk as JSON (default: True)
    document_type="report",       # arbitrary label stored in document metadata
    pdfPipelineOptions=PdfPipelineOptions(),  # optional Docling pipeline config
)
```

#### Metadata extraction

The `metadata_config` field controls how document metadata (title, authors, version, etc.) is
located inside each PDF. Each field is described by a search string and an inclusive page range
to search on. Fields set to `None` are skipped.

```python
from document2graph.models import ExtractorConfig, MetadataExtractionConfig, MetadataFieldConfig

config = ExtractorConfig(
    pdf_path="data/pdfs/",
    data_path="data/output/",
    metadata_config=MetadataExtractionConfig(
        title_page=1,
        version=MetadataFieldConfig(label="edition", pages=(1, 1)),
        authors=MetadataFieldConfig(label="authors", pages=(1, 2)),
        institutions=MetadataFieldConfig(label="affiliations", pages=(1, 2)),
        bibliography=None,        # not present in this document type
        correspondence=MetadataFieldConfig(label="contact", pages=(2, 3)),
    ),
)
```

The default `MetadataExtractionConfig()` targets German-language Praxisempfehlungen PDFs
(version on page 1, authors/institutions/bibliography/correspondence on page 2).

#### Edge weights

Every edge in the document graph carries a `weight` attribute reflecting the strength of the
structural connection. The defaults can be overridden per category via `edge_weights`:

```python
from document2graph.models import ExtractorConfig, EdgeWeightConfig

config = ExtractorConfig(
    pdf_path="data/pdfs/",
    data_path="data/output/",
    edge_weights=EdgeWeightConfig(
        section=1.0,             # heading -> subheading
        text=0.8,                # heading/body -> body text, subtext, footnotes
        list_item=0.6,           # anchor text -> grouped list item (bullets)
        media=0.9,               # referencing text -> image/table (matched via caption)
        unreferenced_media=0.3,  # fallback heading -> image/table without a caption match
        root=0.1,                # document root -> otherwise disconnected node
    ),
)
```

#### Relevancy weights

Optionally, a content-based relevancy weight can be combined with the structural weight
of each edge. Relevancy weights are computed from the parent and child snippet texts and
normalized to [0, 1]: 0 means parent and child are highly similar (the parent adds little
new context for the child), 1 means the parent is very relevant to consider when looking
at the child.

Three metrics are available:

- `bm25` — lexical relevancy: Okapi BM25, clipped at zero and normalized by the maximum
  BM25 score across the document tree. `bm25_scoring` controls the query/document
  asymmetry: `child_query` (default, child text as query against the parent text),
  `symmetric_mean`, or `symmetric_max` (mean/max of both directions)
- `embedding` — semantic relevancy: cosine similarity of sentence embeddings from a
  configurable model (requires the `embeddings` extra: `pip install document2graph[embeddings]`)
- `blend` — linear blend `alpha * embedding + (1 - alpha) * bm25`

```python
from document2graph.models import EdgeWeightConfig, RelevancyWeightConfig

edge_weights = EdgeWeightConfig(
    relevancy=RelevancyWeightConfig(
        enabled=True,
        metric="blend",           # "bm25", "embedding", or "blend"
        embedding_model="sentence-transformers/all-MiniLM-L6-v2",
        alpha=0.5,                # semantic share: 1.0 = fully semantic, 0.0 = fully lexical
        combination="mean",       # combine with the structural weight: "mean" or "multiply"
    ),
)
```

#### Connectivity guarantee

Each graph contains a synthetic root node (`#/document-root`, labelled with the document
title). Any node whose parent cannot be resolved — as well as any component that would
otherwise be disconnected — is attached to this root with the `root` edge weight, so every
document graph is a single (weakly) connected component.

Extracted metadata is available on the `Document` object under the `metadata` sub-object:

```python
result = extractor.run("my_document.pdf")
doc = result["document_metadata"]

print(doc.title)
print(doc.metadata.authors)
print(doc.metadata.version)
```

### Document graph extraction

Processes every PDF in `pdf_path` and builds a NetworkX graph per document.

```python
from document2graph.document2graph_extractor import DocumentGraphExtractor

extractor = DocumentGraphExtractor(config)

# Process all PDFs and save graphs + snippets to disk
extractor.generate_snippets()

# Or collect Snippet objects in memory as well
snippets = extractor.generate_snippets(return_snippets=True)
```

To process a single file and get the raw graph components:

```python
result = extractor.run("my_document.pdf")

text_nodes   = result["text_nodes"]
image_nodes  = result["image_nodes"]
table_nodes  = result["table_nodes"]
edges        = result["edges"]        # list of (parent_id, child_id, weight) tuples
metadata     = result["document_metadata"]
```

Graphs are saved as `.gexf` files under `<data_path>/nx_graphs/` and can be loaded with NetworkX:

```python
import networkx as nx
G = nx.read_gexf("data/output/nx_graphs/my_document.gexf")
```

### Baseline chunk extraction

Produces flat, retrieval-ready chunks using Docling's hybrid chunker.

```python
from document2graph.baseline_extractor import BaselineExtractor

extractor = BaselineExtractor(config)

# Process all PDFs; saves JSON to disk
extractor.generate_baseline_chunks()

# Or return chunks in memory
all_chunks = extractor.generate_baseline_chunks(return_chunks=True)
# all_chunks["my_document.pdf"]["baseline"]  -> List[Chunk]
# all_chunks["my_document.pdf"]["enriched"]  -> List[Chunk] (context-enriched)
```

To process a single file:

```python
chunks = extractor.extract_baseline_chunks(
    filename="my_document.pdf",
    baseline_description="Q1 report chunks",
    enrich=True,   # also produce context-enriched variants
)
```

## Output structure

```
<data_path>/
├── raw_texts/          # intermediate Docling extraction output
├── nx_graphs/          # .gexf graph files (one per document)
├── snippets/           # <filename>_snippets.json (document graph output)
└── baseline_chunks/    # <filename>_baseline_chunks.json
```

## Data models

| Model | Key fields |
|---|---|
| `Snippet` | `snippet_id`, `type` (text/image/table), `document_id`, `sequence_no`, `label`, `level`, `page_no`, `bbox`, `text` |
| `Chunk` | `chunk_id`, `document_id`, `text`, `meta`, `baseline_description`, `embedding` |
| `Document` | `document_id`, `document_type`, `filename`, `title`, `metadata: DocumentMetadata` |
| `DocumentMetadata` | `version`, `authors`, `institutions`, `bibliography`, `correspondence` |
| `MetadataExtractionConfig` | `title_page`, `version`, `authors`, `institutions`, `bibliography`, `correspondence` (each a `MetadataFieldConfig`) |
| `MetadataFieldConfig` | `label` (search string), `pages` (inclusive 1-based page range, e.g. `(1, 3)`) |
| `EdgeWeightConfig` | `section`, `text`, `list_item`, `media`, `unreferenced_media`, `root` (edge weights by category), `relevancy: RelevancyWeightConfig` |
| `RelevancyWeightConfig` | `enabled`, `metric` (bm25/embedding/blend), `embedding_model`, `alpha`, `bm25_k1`, `bm25_b`, `bm25_scoring` (child_query/symmetric_mean/symmetric_max), `combination` (mean/multiply) |

## Running tests

```bash
pytest
```

Skip slow tests that download models:

```bash
pytest -m "not slow"
```
