document2graph 0.1.0__py3-none-any.whl

document2graph/models/Document.py

@@ -0,0 +1,10 @@
+from pydantic import BaseModel
+from .DocumentMetadata import DocumentMetadata
+
+
+class Document(BaseModel):
+    document_id: str = ""
+    document_type: str = ""
+    filename: str = ""
+    title: str = ""
+    metadata: DocumentMetadata = DocumentMetadata()

document2graph/models/DocumentMetadata.py

@@ -0,0 +1,23 @@
+from pydantic import BaseModel
+
+
+class MetadataFieldConfig(BaseModel):
+    label: str               # substring to search for in the document text
+    pages: tuple[int, int]   # inclusive 1-based page range to search on
+
+
+class MetadataExtractionConfig(BaseModel):
+    title_page: int = 2
+    version: MetadataFieldConfig | None = MetadataFieldConfig(label="version", pages=(1, 1))
+    authors: MetadataFieldConfig | None = MetadataFieldConfig(label="korrespondenzadresse", pages=(2, 2))
+    institutions: MetadataFieldConfig | None = MetadataFieldConfig(label="institute", pages=(2, 2))
+    bibliography: MetadataFieldConfig | None = MetadataFieldConfig(label="bibliografie", pages=(2, 2))
+    correspondence: MetadataFieldConfig | None = MetadataFieldConfig(label="korrespondenzadresse", pages=(2, 2))
+
+
+class DocumentMetadata(BaseModel):
+    version: str = ""
+    authors: str = ""
+    institutions: str = ""
+    bibliography: str = ""
+    correspondence: str = ""

document2graph/models/EdgeWeightConfig.py

@@ -0,0 +1,18 @@
+from pydantic import BaseModel
+
+
+class EdgeWeightConfig(BaseModel):
+    """Weights assigned to graph edges by structural category.
+
+    Higher values indicate a stronger structural connection between
+    parent and child node. All weights are stored on the edge as the
+    ``weight`` attribute (both in the returned edge list and in the
+    exported GEXF graph).
+    """
+
+    section: float = 1.0            # heading -> subheading
+    text: float = 0.8               # heading/body -> body text, subtext, footnotes
+    list_item: float = 0.6          # anchor text -> grouped list item (bullets)
+    media: float = 0.9              # referencing text -> image/table (matched via caption)
+    unreferenced_media: float = 0.3 # fallback heading -> image/table without a caption match
+    root: float = 0.1               # synthetic document root -> otherwise disconnected node

document2graph/models/ExtractorConfig.py

@@ -0,0 +1,14 @@
+from pydantic import BaseModel, Field
+from docling.datamodel.pipeline_options import PdfPipelineOptions
+from .DocumentMetadata import MetadataExtractionConfig
+from .EdgeWeightConfig import EdgeWeightConfig
+
+
+class ExtractorConfig(BaseModel):
+    pdf_path: str
+    data_path: str
+    pdfPipelineOptions: PdfPipelineOptions = PdfPipelineOptions()
+    document_type: str = ""
+    save_json: bool = True
+    metadata_config: MetadataExtractionConfig = Field(default_factory=MetadataExtractionConfig)
+    edge_weights: EdgeWeightConfig = Field(default_factory=EdgeWeightConfig)

document2graph/models/ImageSnippetNode.py

@@ -0,0 +1,21 @@
+from pydantic import BaseModel
+from docling_core.types.doc.document import RefItem
+from docling_core.types.doc.base import BoundingBox
+from .TextSnippetNode import TextSnippetNode
+
+
+class ImageSnippetNode(BaseModel):
+    snippet_id: str
+    document_id: str
+    label: str
+    sequence_no: int
+    level: int
+    level_label: str 
+    parent_id: str 
+    docling_parent_ref: RefItem | None
+    docling_self_ref: RefItem | None
+    is_grouped: bool = False
+    caption_nodes: list[TextSnippetNode] | list = []
+    caption_text: str 
+    bbox: BoundingBox
+    page_no: int

document2graph/models/Snippet.py

@@ -0,0 +1,19 @@
+from pydantic import BaseModel
+from docling_core.types.doc.document import RefItem
+from docling_core.types.doc.base import BoundingBox
+
+class Snippet(BaseModel):
+    snippet_id: str
+    type: str
+    document_id: str
+    sequence_no: int
+    label: str
+    level: int
+    level_label: str 
+    parent_id: str | None
+    is_grouped: bool = False
+    page_no: int
+    bbox: BoundingBox
+    text: str
+    docling_parent_ref: RefItem | None
+    docling_self_ref: RefItem | None

document2graph/models/TableSnippetNode.py

@@ -0,0 +1,6 @@
+
+from .ImageSnippetNode import ImageSnippetNode
+
+
+class TableSnippetNode(ImageSnippetNode):
+    markdown_serialization: str

document2graph/models/TextSnippet.py

@@ -0,0 +1,6 @@
+from docling_core.types.doc.document import TextItem
+from pydantic import BaseModel
+
+class TextSnippet(BaseModel):
+    text_item: TextItem
+    line_heights: list[float] = []

document2graph/models/TextSnippetNode.py

@@ -0,0 +1,19 @@
+from pydantic import BaseModel
+from docling_core.types.doc.document import RefItem
+from docling_core.types.doc.base import BoundingBox
+
+class TextSnippetNode(BaseModel):
+    snippet_id: str
+    document_id: str
+    label: str
+    sequence_no: int
+    level: int
+    level_label: str 
+    parent_id: str | None
+    docling_parent_ref: RefItem 
+    docling_self_ref: RefItem 
+    is_grouped: bool = False
+    text: str
+    bbox: BoundingBox
+    charspan: tuple[int, int]
+    page_no: int

document2graph/models/__init__.py

@@ -0,0 +1,25 @@
+from .Chunk import Chunk
+from .EdgeWeightConfig import EdgeWeightConfig
+from .ExtractorConfig import ExtractorConfig
+from .Snippet import Snippet
+from .Document import Document
+from .DocumentMetadata import DocumentMetadata, MetadataExtractionConfig, MetadataFieldConfig
+from .ImageSnippetNode import ImageSnippetNode
+from .TableSnippetNode import TableSnippetNode
+from .TextSnippetNode import TextSnippetNode
+from .TextSnippet import TextSnippet
+
+__all__ = [
+    "Chunk",
+    "EdgeWeightConfig",
+    "ExtractorConfig",
+    "Snippet",
+    "Document",
+    "DocumentMetadata",
+    "MetadataExtractionConfig",
+    "MetadataFieldConfig",
+    "ImageSnippetNode",
+    "TableSnippetNode",
+    "TextSnippetNode",
+    "TextSnippet",
+]

document2graph/utils/base_extractor.py

@@ -0,0 +1,127 @@
+from docling.document_converter import DocumentConverter, PdfFormatOption
+from docling.datamodel.base_models import InputFormat
+from docling_core.types.doc.document import DoclingDocument, TableItem, TextItem, PictureItem
+from docling.datamodel.pipeline_options import PdfPipelineOptions
+from docling_core.types.doc.document import GroupItem, RefItem
+import os
+import pickle
+from .log import Log
+
+class Extractor:
+    def __init__(self, source: str, pipeline_options: PdfPipelineOptions = PdfPipelineOptions()):
+        self.source = source
+        self.converter = DocumentConverter(format_options={
+            InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)
+        })
+        self.docling_doc = self.converter.convert(source)
+        self.doc = self.docling_doc.document
+
+        self.logger = Log("Extractor").logger
+    
+
+    def is_furniture(self, item):
+        """
+        Check if the item is furniture (e.g., title page, table of contents).
+        """
+        return item.content_layer.value == "furniture"
+
+    def is_text(self, item):
+        """
+        Check if the item is a text item.
+        """
+        return isinstance(item, TextItem)
+
+    def is_picture(self, item):
+        """
+        Check if the item is a picture item.
+        """
+        return isinstance(item, PictureItem)
+
+    def is_table(self, item):
+        """
+        Check if the item is a table item.
+        """
+        return isinstance(item, TableItem)
+
+    def is_group(self, item):
+        """
+        Check if the item is a group item.
+        """
+        return isinstance(item, GroupItem)
+
+    def serialize_ref_items(
+        self, ref_items: list[RefItem], doc: DoclingDocument
+    ) -> tuple[list[str], list[str]]:
+        """
+        Serialize the reference items to a string representation.
+        This function is necessary to handle also string representations of groups, tables and images.
+        Only for annotation of extracted doc necessary.
+        """
+        serialized_items = []
+        serialized_labels = []
+        for refItem in ref_items:
+            item = refItem.resolve(doc)
+            if self.is_furniture(item):
+                continue
+            if self.is_text(item):
+                serialized_items.append(item.text)
+                serialized_labels.append(item.label.value)
+            elif self.is_picture(item):
+                serialized_items.append(item.caption_text(doc))
+                serialized_labels.append(item.label.value)
+            elif self.is_table(item):
+                serialized_items.append(item.export_to_markdown(doc))
+                serialized_labels.append(item.label.value)
+            elif self.is_group(item):
+                # Recursively serialize group items
+                group_items, group_labels = self.serialize_ref_items(item.children, doc)
+                serialized_items.append(group_items)
+                serialized_labels.append(group_labels)
+            else:
+                serialized_items.append(item.text)
+                serialized_labels.append(item.label.value)
+        return serialized_items, serialized_labels
+
+
+    def serialize_doc(self, doc: DoclingDocument) -> dict:
+        """
+        Serialize the document to a string representation.
+        Leave furniture out. Bring text in order, to see if we can manually label the paragraphs.
+        """
+        full_text, full_text_labels = self.serialize_ref_items(
+            doc.body.children, doc
+        )
+        return {
+            "full_text": full_text,
+            "full_text_labels": full_text_labels,
+            "pictures": doc.pictures,
+            "tables": doc.tables,
+            "texts": doc.texts,
+        }
+
+    def extract(self, save_dir: str, filename: str) -> dict:
+        """
+        Extract and serialize the document, saving intermediate results.
+        """
+        self.logger.info(f"Saving parsed pdf as {filename}.json")
+
+        os.makedirs(save_dir, exist_ok=True)
+
+        self.docling_doc.document.save_as_json(
+            f"{save_dir}/{filename}_docling_doc.json"
+        )
+        self.logger.info(f"Start extracting {filename}...")
+        result = self.serialize_doc(self.doc)
+        with open(f"{save_dir}/{filename}_serialized_doc.pkl", "wb") as f:
+            pickle.dump(result, f)
+        return result
+
+
+def main():
+    # only for quick testing
+    source = "../pdf/ddg_pdf/ddg_praxisempfehlungen/DuS_2024_S02_Praxisempfehlungen_Aberle_Adipositas-und-Diabetes.pdf"
+    extractor = Extractor(source)
+    _ = extractor.extract("test", "test")
+
+if __name__ == "__main__":
+    main()

document2graph/utils/doc_to_tree_lib.py

@@ -0,0 +1,52 @@
+import numpy as np
+
+# schemas
+from ..models.TextSnippet import TextSnippet
+
+def get_heights(text_items: list[TextSnippet], skip_first_page: bool = False) -> np.ndarray:
+    if skip_first_page:
+        heights = [height for snippet in text_items for height in snippet.line_heights if snippet.text_item.prov[0].page_no > 1]
+    else:
+        heights = [height for snippet in text_items for height in snippet.line_heights]
+    return np.multiply(np.round(np.array(heights), 1), 10)
+
+def get_height_hist(heights=np.array, plot=False) -> dict[int, int]:
+        if not heights.size:
+            print("No heights found, returning empty histogram.")
+            return {}
+        unique_heights, counts = np.unique(heights, return_counts=True)
+        height_hist = dict(zip(unique_heights, counts))
+        if plot:
+            # matplotlib is only needed for this debug plot, not a package dependency
+            import matplotlib.pyplot as plt
+            plt.bar(height_hist.keys(), height_hist.values()) # type: ignore
+            xticks = unique_heights[::max(1, len(unique_heights)//20)]
+            plt.xticks(xticks, rotation=90, size=6)
+            plt.xlabel("Font size")
+            plt.ylabel("Count")
+            plt.show()
+            
+        return height_hist
+
+def get_text_body_height(heights: np.ndarray) -> tuple[float, float]:
+        assert heights.size > 0, "No heights provided to compute text body height."
+        # text body level range: between 25th and 75th percentile
+        text_body_height = np.percentile(heights, 50)  # median
+        iqr = np.percentile(heights, 75) - np.percentile(heights, 25)
+        return float(text_body_height), float(iqr)
+
+def line_height_to_key(heights: list[float]) -> np.ndarray:
+    return np.multiply(heights, 10).astype(int)  # multiply by 10 to avoid float precision issues
+
+def compute_snippet_level(snippet: TextSnippet, levels: dict[int, int]) -> int:
+    heights = line_height_to_key(snippet.line_heights)
+    if not heights.size:
+        return 1000  # unknown level
+    # get median height in case there are multiple lines with different heights
+    median_height = np.median(heights)
+    if median_height in levels:
+        return levels[int(median_height)]
+    else:
+        # find closest height in levels
+        closest_height = min(levels.keys(), key=lambda h: abs(h - median_height))
+        return levels[closest_height]

document2graph/utils/hybrid_chunker_extractor.py

@@ -0,0 +1,26 @@
+import os
+from .base_extractor import Extractor
+from docling.datamodel.pipeline_options import PdfPipelineOptions
+from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
+from docling_core.transforms.chunker.tokenizer.huggingface import HuggingFaceTokenizer
+
+class HybridChunkerExtractor(Extractor):
+    def __init__(self, source: str, merge_peers: bool, pipeline_options: PdfPipelineOptions = PdfPipelineOptions()):
+        super().__init__(source, pipeline_options)
+        self.merge_peers = merge_peers
+        self.chunker = HybridChunker(
+            tokenizer=HuggingFaceTokenizer.from_pretrained("google/gemma-3-4b-it", max_tokens=131072),
+            merge_peers=merge_peers)
+
+    def extract_and_chunk(self, save_dir: str, filename: str):
+        # First, we extract the document using the base extractor logic
+        docling_doc = self.converter.convert(self.source).document
+        # Then, we apply the hybrid chunker to the extracted document
+        self.logger.info(f"Saving parsed pdf as {filename}.json")
+
+        os.makedirs(save_dir, exist_ok=True)
+
+        docling_doc.save_as_json(
+            f"{save_dir}/{filename}_baseline_docling_doc.json"
+        )
+        yield from self.chunker.chunk(docling_doc)

document2graph/utils/log.py

@@ -0,0 +1,40 @@
+import sys
+import coloredlogs
+import logging
+from pprint import pprint
+
+logger = logging.getLogger()
+coloredlogs.install(level='INFO', logger=logger, isatty=True)
+
+
+class Log:
+    def __init__(self, name):
+        self.logger = logging.getLogger(name)
+
+    def log(self, s):
+        self.logger.info(s)
+        sys.stdout.flush()
+
+    def plog(self, s):
+        plog(s)
+
+    def info(self, s):
+        self.logger.info(s)
+        sys.stdout.flush()
+
+    def debug(self, s):
+        pprint(s)
+        sys.stdout.flush()
+
+    def error(self, s):
+        self.logger.error(s)
+        sys.stdout.flush()
+
+    def warning(self, s):
+        self.logger.warning(s)
+        sys.stdout.flush()
+
+
+def plog(s):
+    pprint(s)
+    sys.stdout.flush()

document2graph-0.1.0.dist-info/METADATA

@@ -0,0 +1,243 @@
+Metadata-Version: 2.4
+Name: document2graph
+Version: 0.1.0
+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"
+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
+    ),
+)
+```
+
+#### 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) |
+
+## Running tests
+
+```bash
+pytest
+```
+
+Skip slow tests that download models:
+
+```bash
+pytest -m "not slow"
+```