document2graph 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- document2graph-0.1.0/LICENSE +21 -0
- document2graph-0.1.0/PKG-INFO +243 -0
- document2graph-0.1.0/README.md +208 -0
- document2graph-0.1.0/document2graph/__init__.py +0 -0
- document2graph-0.1.0/document2graph/baseline_extractor/__init__.py +3 -0
- document2graph-0.1.0/document2graph/baseline_extractor/extractor.py +82 -0
- document2graph-0.1.0/document2graph/document2graph_extractor/__init__.py +3 -0
- document2graph-0.1.0/document2graph/document2graph_extractor/extractor.py +146 -0
- document2graph-0.1.0/document2graph/document2graph_extractor/snippet_graph_constructor.py +424 -0
- document2graph-0.1.0/document2graph/models/Chunk.py +10 -0
- document2graph-0.1.0/document2graph/models/Document.py +10 -0
- document2graph-0.1.0/document2graph/models/DocumentMetadata.py +23 -0
- document2graph-0.1.0/document2graph/models/EdgeWeightConfig.py +18 -0
- document2graph-0.1.0/document2graph/models/ExtractorConfig.py +14 -0
- document2graph-0.1.0/document2graph/models/ImageSnippetNode.py +21 -0
- document2graph-0.1.0/document2graph/models/Snippet.py +19 -0
- document2graph-0.1.0/document2graph/models/TableSnippetNode.py +6 -0
- document2graph-0.1.0/document2graph/models/TextSnippet.py +6 -0
- document2graph-0.1.0/document2graph/models/TextSnippetNode.py +19 -0
- document2graph-0.1.0/document2graph/models/__init__.py +25 -0
- document2graph-0.1.0/document2graph/utils/__init__.py +0 -0
- document2graph-0.1.0/document2graph/utils/base_extractor.py +127 -0
- document2graph-0.1.0/document2graph/utils/doc_to_tree_lib.py +52 -0
- document2graph-0.1.0/document2graph/utils/hybrid_chunker_extractor.py +26 -0
- document2graph-0.1.0/document2graph/utils/log.py +40 -0
- document2graph-0.1.0/document2graph.egg-info/PKG-INFO +243 -0
- document2graph-0.1.0/document2graph.egg-info/SOURCES.txt +32 -0
- document2graph-0.1.0/document2graph.egg-info/dependency_links.txt +1 -0
- document2graph-0.1.0/document2graph.egg-info/requires.txt +14 -0
- document2graph-0.1.0/document2graph.egg-info/top_level.txt +1 -0
- document2graph-0.1.0/pyproject.toml +53 -0
- document2graph-0.1.0/setup.cfg +4 -0
- document2graph-0.1.0/tests/test_document2graph_extractor.py +71 -0
- document2graph-0.1.0/tests/test_snippet_graph.py +119 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Marina Walther
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,243 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: document2graph
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Extract hierarchical document graphs and baseline chunks from PDF
|
|
5
|
+
Author-email: Marina Walther <marina.walther@gmail.com>
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/mwalther10/document2graph
|
|
8
|
+
Project-URL: Issues, https://github.com/mwalther10/document2graph/issues
|
|
9
|
+
Keywords: pdf,document2graph,docling,networkx,chunking,retrieval
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: Intended Audience :: Developers
|
|
12
|
+
Classifier: Intended Audience :: Science/Research
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
17
|
+
Classifier: Topic :: Text Processing :: Indexing
|
|
18
|
+
Requires-Python: >=3.11
|
|
19
|
+
Description-Content-Type: text/markdown
|
|
20
|
+
License-File: LICENSE
|
|
21
|
+
Requires-Dist: docling
|
|
22
|
+
Requires-Dist: docling-core
|
|
23
|
+
Requires-Dist: docling-parse
|
|
24
|
+
Requires-Dist: numpy
|
|
25
|
+
Requires-Dist: pydantic>=2
|
|
26
|
+
Requires-Dist: regex
|
|
27
|
+
Requires-Dist: networkx
|
|
28
|
+
Requires-Dist: coloredlogs
|
|
29
|
+
Requires-Dist: tqdm
|
|
30
|
+
Requires-Dist: transformers
|
|
31
|
+
Provides-Extra: dev
|
|
32
|
+
Requires-Dist: pytest; extra == "dev"
|
|
33
|
+
Requires-Dist: httpx; extra == "dev"
|
|
34
|
+
Dynamic: license-file
|
|
35
|
+
|
|
36
|
+
# document2graph
|
|
37
|
+
|
|
38
|
+
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:
|
|
39
|
+
|
|
40
|
+
- **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.
|
|
41
|
+
- **Baseline chunks** — flat, semantically merged chunks suitable for retrieval pipelines, with optional context enrichment.
|
|
42
|
+
|
|
43
|
+
## Installation
|
|
44
|
+
|
|
45
|
+
```bash
|
|
46
|
+
pip install document2graph
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
Or from source:
|
|
50
|
+
|
|
51
|
+
```bash
|
|
52
|
+
pip install -e .
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
For development dependencies (pytest):
|
|
56
|
+
|
|
57
|
+
```bash
|
|
58
|
+
pip install -e ".[dev]"
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
Requires Python 3.11+.
|
|
62
|
+
|
|
63
|
+
## Quick start
|
|
64
|
+
|
|
65
|
+
### Configuration
|
|
66
|
+
|
|
67
|
+
Both extractors share the same `ExtractorConfig`:
|
|
68
|
+
|
|
69
|
+
```python
|
|
70
|
+
from document2graph.models import ExtractorConfig
|
|
71
|
+
from docling.datamodel.pipeline_options import PdfPipelineOptions
|
|
72
|
+
|
|
73
|
+
config = ExtractorConfig(
|
|
74
|
+
pdf_path="data/pdfs/", # directory containing input PDFs
|
|
75
|
+
data_path="data/output/", # root directory for all output files
|
|
76
|
+
save_json=True, # write results to disk as JSON (default: True)
|
|
77
|
+
document_type="report", # arbitrary label stored in document metadata
|
|
78
|
+
pdfPipelineOptions=PdfPipelineOptions(), # optional Docling pipeline config
|
|
79
|
+
)
|
|
80
|
+
```
|
|
81
|
+
|
|
82
|
+
#### Metadata extraction
|
|
83
|
+
|
|
84
|
+
The `metadata_config` field controls how document metadata (title, authors, version, etc.) is
|
|
85
|
+
located inside each PDF. Each field is described by a search string and an inclusive page range
|
|
86
|
+
to search on. Fields set to `None` are skipped.
|
|
87
|
+
|
|
88
|
+
```python
|
|
89
|
+
from document2graph.models import ExtractorConfig, MetadataExtractionConfig, MetadataFieldConfig
|
|
90
|
+
|
|
91
|
+
config = ExtractorConfig(
|
|
92
|
+
pdf_path="data/pdfs/",
|
|
93
|
+
data_path="data/output/",
|
|
94
|
+
metadata_config=MetadataExtractionConfig(
|
|
95
|
+
title_page=1,
|
|
96
|
+
version=MetadataFieldConfig(label="edition", pages=(1, 1)),
|
|
97
|
+
authors=MetadataFieldConfig(label="authors", pages=(1, 2)),
|
|
98
|
+
institutions=MetadataFieldConfig(label="affiliations", pages=(1, 2)),
|
|
99
|
+
bibliography=None, # not present in this document type
|
|
100
|
+
correspondence=MetadataFieldConfig(label="contact", pages=(2, 3)),
|
|
101
|
+
),
|
|
102
|
+
)
|
|
103
|
+
```
|
|
104
|
+
|
|
105
|
+
The default `MetadataExtractionConfig()` targets German-language Praxisempfehlungen PDFs
|
|
106
|
+
(version on page 1, authors/institutions/bibliography/correspondence on page 2).
|
|
107
|
+
|
|
108
|
+
#### Edge weights
|
|
109
|
+
|
|
110
|
+
Every edge in the document graph carries a `weight` attribute reflecting the strength of the
|
|
111
|
+
structural connection. The defaults can be overridden per category via `edge_weights`:
|
|
112
|
+
|
|
113
|
+
```python
|
|
114
|
+
from document2graph.models import ExtractorConfig, EdgeWeightConfig
|
|
115
|
+
|
|
116
|
+
config = ExtractorConfig(
|
|
117
|
+
pdf_path="data/pdfs/",
|
|
118
|
+
data_path="data/output/",
|
|
119
|
+
edge_weights=EdgeWeightConfig(
|
|
120
|
+
section=1.0, # heading -> subheading
|
|
121
|
+
text=0.8, # heading/body -> body text, subtext, footnotes
|
|
122
|
+
list_item=0.6, # anchor text -> grouped list item (bullets)
|
|
123
|
+
media=0.9, # referencing text -> image/table (matched via caption)
|
|
124
|
+
unreferenced_media=0.3, # fallback heading -> image/table without a caption match
|
|
125
|
+
root=0.1, # document root -> otherwise disconnected node
|
|
126
|
+
),
|
|
127
|
+
)
|
|
128
|
+
```
|
|
129
|
+
|
|
130
|
+
#### Connectivity guarantee
|
|
131
|
+
|
|
132
|
+
Each graph contains a synthetic root node (`#/document-root`, labelled with the document
|
|
133
|
+
title). Any node whose parent cannot be resolved — as well as any component that would
|
|
134
|
+
otherwise be disconnected — is attached to this root with the `root` edge weight, so every
|
|
135
|
+
document graph is a single (weakly) connected component.
|
|
136
|
+
|
|
137
|
+
Extracted metadata is available on the `Document` object under the `metadata` sub-object:
|
|
138
|
+
|
|
139
|
+
```python
|
|
140
|
+
result = extractor.run("my_document.pdf")
|
|
141
|
+
doc = result["document_metadata"]
|
|
142
|
+
|
|
143
|
+
print(doc.title)
|
|
144
|
+
print(doc.metadata.authors)
|
|
145
|
+
print(doc.metadata.version)
|
|
146
|
+
```
|
|
147
|
+
|
|
148
|
+
### Document graph extraction
|
|
149
|
+
|
|
150
|
+
Processes every PDF in `pdf_path` and builds a NetworkX graph per document.
|
|
151
|
+
|
|
152
|
+
```python
|
|
153
|
+
from document2graph.document2graph_extractor import DocumentGraphExtractor
|
|
154
|
+
|
|
155
|
+
extractor = DocumentGraphExtractor(config)
|
|
156
|
+
|
|
157
|
+
# Process all PDFs and save graphs + snippets to disk
|
|
158
|
+
extractor.generate_snippets()
|
|
159
|
+
|
|
160
|
+
# Or collect Snippet objects in memory as well
|
|
161
|
+
snippets = extractor.generate_snippets(return_snippets=True)
|
|
162
|
+
```
|
|
163
|
+
|
|
164
|
+
To process a single file and get the raw graph components:
|
|
165
|
+
|
|
166
|
+
```python
|
|
167
|
+
result = extractor.run("my_document.pdf")
|
|
168
|
+
|
|
169
|
+
text_nodes = result["text_nodes"]
|
|
170
|
+
image_nodes = result["image_nodes"]
|
|
171
|
+
table_nodes = result["table_nodes"]
|
|
172
|
+
edges = result["edges"] # list of (parent_id, child_id, weight) tuples
|
|
173
|
+
metadata = result["document_metadata"]
|
|
174
|
+
```
|
|
175
|
+
|
|
176
|
+
Graphs are saved as `.gexf` files under `<data_path>/nx_graphs/` and can be loaded with NetworkX:
|
|
177
|
+
|
|
178
|
+
```python
|
|
179
|
+
import networkx as nx
|
|
180
|
+
G = nx.read_gexf("data/output/nx_graphs/my_document.gexf")
|
|
181
|
+
```
|
|
182
|
+
|
|
183
|
+
### Baseline chunk extraction
|
|
184
|
+
|
|
185
|
+
Produces flat, retrieval-ready chunks using Docling's hybrid chunker.
|
|
186
|
+
|
|
187
|
+
```python
|
|
188
|
+
from document2graph.baseline_extractor import BaselineExtractor
|
|
189
|
+
|
|
190
|
+
extractor = BaselineExtractor(config)
|
|
191
|
+
|
|
192
|
+
# Process all PDFs; saves JSON to disk
|
|
193
|
+
extractor.generate_baseline_chunks()
|
|
194
|
+
|
|
195
|
+
# Or return chunks in memory
|
|
196
|
+
all_chunks = extractor.generate_baseline_chunks(return_chunks=True)
|
|
197
|
+
# all_chunks["my_document.pdf"]["baseline"] -> List[Chunk]
|
|
198
|
+
# all_chunks["my_document.pdf"]["enriched"] -> List[Chunk] (context-enriched)
|
|
199
|
+
```
|
|
200
|
+
|
|
201
|
+
To process a single file:
|
|
202
|
+
|
|
203
|
+
```python
|
|
204
|
+
chunks = extractor.extract_baseline_chunks(
|
|
205
|
+
filename="my_document.pdf",
|
|
206
|
+
baseline_description="Q1 report chunks",
|
|
207
|
+
enrich=True, # also produce context-enriched variants
|
|
208
|
+
)
|
|
209
|
+
```
|
|
210
|
+
|
|
211
|
+
## Output structure
|
|
212
|
+
|
|
213
|
+
```
|
|
214
|
+
<data_path>/
|
|
215
|
+
├── raw_texts/ # intermediate Docling extraction output
|
|
216
|
+
├── nx_graphs/ # .gexf graph files (one per document)
|
|
217
|
+
├── snippets/ # <filename>_snippets.json (document graph output)
|
|
218
|
+
└── baseline_chunks/ # <filename>_baseline_chunks.json
|
|
219
|
+
```
|
|
220
|
+
|
|
221
|
+
## Data models
|
|
222
|
+
|
|
223
|
+
| Model | Key fields |
|
|
224
|
+
|---|---|
|
|
225
|
+
| `Snippet` | `snippet_id`, `type` (text/image/table), `document_id`, `sequence_no`, `label`, `level`, `page_no`, `bbox`, `text` |
|
|
226
|
+
| `Chunk` | `chunk_id`, `document_id`, `text`, `meta`, `baseline_description`, `embedding` |
|
|
227
|
+
| `Document` | `document_id`, `document_type`, `filename`, `title`, `metadata: DocumentMetadata` |
|
|
228
|
+
| `DocumentMetadata` | `version`, `authors`, `institutions`, `bibliography`, `correspondence` |
|
|
229
|
+
| `MetadataExtractionConfig` | `title_page`, `version`, `authors`, `institutions`, `bibliography`, `correspondence` (each a `MetadataFieldConfig`) |
|
|
230
|
+
| `MetadataFieldConfig` | `label` (search string), `pages` (inclusive 1-based page range, e.g. `(1, 3)`) |
|
|
231
|
+
| `EdgeWeightConfig` | `section`, `text`, `list_item`, `media`, `unreferenced_media`, `root` (edge weights by category) |
|
|
232
|
+
|
|
233
|
+
## Running tests
|
|
234
|
+
|
|
235
|
+
```bash
|
|
236
|
+
pytest
|
|
237
|
+
```
|
|
238
|
+
|
|
239
|
+
Skip slow tests that download models:
|
|
240
|
+
|
|
241
|
+
```bash
|
|
242
|
+
pytest -m "not slow"
|
|
243
|
+
```
|
|
@@ -0,0 +1,208 @@
|
|
|
1
|
+
# document2graph
|
|
2
|
+
|
|
3
|
+
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:
|
|
4
|
+
|
|
5
|
+
- **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.
|
|
6
|
+
- **Baseline chunks** — flat, semantically merged chunks suitable for retrieval pipelines, with optional context enrichment.
|
|
7
|
+
|
|
8
|
+
## Installation
|
|
9
|
+
|
|
10
|
+
```bash
|
|
11
|
+
pip install document2graph
|
|
12
|
+
```
|
|
13
|
+
|
|
14
|
+
Or from source:
|
|
15
|
+
|
|
16
|
+
```bash
|
|
17
|
+
pip install -e .
|
|
18
|
+
```
|
|
19
|
+
|
|
20
|
+
For development dependencies (pytest):
|
|
21
|
+
|
|
22
|
+
```bash
|
|
23
|
+
pip install -e ".[dev]"
|
|
24
|
+
```
|
|
25
|
+
|
|
26
|
+
Requires Python 3.11+.
|
|
27
|
+
|
|
28
|
+
## Quick start
|
|
29
|
+
|
|
30
|
+
### Configuration
|
|
31
|
+
|
|
32
|
+
Both extractors share the same `ExtractorConfig`:
|
|
33
|
+
|
|
34
|
+
```python
|
|
35
|
+
from document2graph.models import ExtractorConfig
|
|
36
|
+
from docling.datamodel.pipeline_options import PdfPipelineOptions
|
|
37
|
+
|
|
38
|
+
config = ExtractorConfig(
|
|
39
|
+
pdf_path="data/pdfs/", # directory containing input PDFs
|
|
40
|
+
data_path="data/output/", # root directory for all output files
|
|
41
|
+
save_json=True, # write results to disk as JSON (default: True)
|
|
42
|
+
document_type="report", # arbitrary label stored in document metadata
|
|
43
|
+
pdfPipelineOptions=PdfPipelineOptions(), # optional Docling pipeline config
|
|
44
|
+
)
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
#### Metadata extraction
|
|
48
|
+
|
|
49
|
+
The `metadata_config` field controls how document metadata (title, authors, version, etc.) is
|
|
50
|
+
located inside each PDF. Each field is described by a search string and an inclusive page range
|
|
51
|
+
to search on. Fields set to `None` are skipped.
|
|
52
|
+
|
|
53
|
+
```python
|
|
54
|
+
from document2graph.models import ExtractorConfig, MetadataExtractionConfig, MetadataFieldConfig
|
|
55
|
+
|
|
56
|
+
config = ExtractorConfig(
|
|
57
|
+
pdf_path="data/pdfs/",
|
|
58
|
+
data_path="data/output/",
|
|
59
|
+
metadata_config=MetadataExtractionConfig(
|
|
60
|
+
title_page=1,
|
|
61
|
+
version=MetadataFieldConfig(label="edition", pages=(1, 1)),
|
|
62
|
+
authors=MetadataFieldConfig(label="authors", pages=(1, 2)),
|
|
63
|
+
institutions=MetadataFieldConfig(label="affiliations", pages=(1, 2)),
|
|
64
|
+
bibliography=None, # not present in this document type
|
|
65
|
+
correspondence=MetadataFieldConfig(label="contact", pages=(2, 3)),
|
|
66
|
+
),
|
|
67
|
+
)
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
The default `MetadataExtractionConfig()` targets German-language Praxisempfehlungen PDFs
|
|
71
|
+
(version on page 1, authors/institutions/bibliography/correspondence on page 2).
|
|
72
|
+
|
|
73
|
+
#### Edge weights
|
|
74
|
+
|
|
75
|
+
Every edge in the document graph carries a `weight` attribute reflecting the strength of the
|
|
76
|
+
structural connection. The defaults can be overridden per category via `edge_weights`:
|
|
77
|
+
|
|
78
|
+
```python
|
|
79
|
+
from document2graph.models import ExtractorConfig, EdgeWeightConfig
|
|
80
|
+
|
|
81
|
+
config = ExtractorConfig(
|
|
82
|
+
pdf_path="data/pdfs/",
|
|
83
|
+
data_path="data/output/",
|
|
84
|
+
edge_weights=EdgeWeightConfig(
|
|
85
|
+
section=1.0, # heading -> subheading
|
|
86
|
+
text=0.8, # heading/body -> body text, subtext, footnotes
|
|
87
|
+
list_item=0.6, # anchor text -> grouped list item (bullets)
|
|
88
|
+
media=0.9, # referencing text -> image/table (matched via caption)
|
|
89
|
+
unreferenced_media=0.3, # fallback heading -> image/table without a caption match
|
|
90
|
+
root=0.1, # document root -> otherwise disconnected node
|
|
91
|
+
),
|
|
92
|
+
)
|
|
93
|
+
```
|
|
94
|
+
|
|
95
|
+
#### Connectivity guarantee
|
|
96
|
+
|
|
97
|
+
Each graph contains a synthetic root node (`#/document-root`, labelled with the document
|
|
98
|
+
title). Any node whose parent cannot be resolved — as well as any component that would
|
|
99
|
+
otherwise be disconnected — is attached to this root with the `root` edge weight, so every
|
|
100
|
+
document graph is a single (weakly) connected component.
|
|
101
|
+
|
|
102
|
+
Extracted metadata is available on the `Document` object under the `metadata` sub-object:
|
|
103
|
+
|
|
104
|
+
```python
|
|
105
|
+
result = extractor.run("my_document.pdf")
|
|
106
|
+
doc = result["document_metadata"]
|
|
107
|
+
|
|
108
|
+
print(doc.title)
|
|
109
|
+
print(doc.metadata.authors)
|
|
110
|
+
print(doc.metadata.version)
|
|
111
|
+
```
|
|
112
|
+
|
|
113
|
+
### Document graph extraction
|
|
114
|
+
|
|
115
|
+
Processes every PDF in `pdf_path` and builds a NetworkX graph per document.
|
|
116
|
+
|
|
117
|
+
```python
|
|
118
|
+
from document2graph.document2graph_extractor import DocumentGraphExtractor
|
|
119
|
+
|
|
120
|
+
extractor = DocumentGraphExtractor(config)
|
|
121
|
+
|
|
122
|
+
# Process all PDFs and save graphs + snippets to disk
|
|
123
|
+
extractor.generate_snippets()
|
|
124
|
+
|
|
125
|
+
# Or collect Snippet objects in memory as well
|
|
126
|
+
snippets = extractor.generate_snippets(return_snippets=True)
|
|
127
|
+
```
|
|
128
|
+
|
|
129
|
+
To process a single file and get the raw graph components:
|
|
130
|
+
|
|
131
|
+
```python
|
|
132
|
+
result = extractor.run("my_document.pdf")
|
|
133
|
+
|
|
134
|
+
text_nodes = result["text_nodes"]
|
|
135
|
+
image_nodes = result["image_nodes"]
|
|
136
|
+
table_nodes = result["table_nodes"]
|
|
137
|
+
edges = result["edges"] # list of (parent_id, child_id, weight) tuples
|
|
138
|
+
metadata = result["document_metadata"]
|
|
139
|
+
```
|
|
140
|
+
|
|
141
|
+
Graphs are saved as `.gexf` files under `<data_path>/nx_graphs/` and can be loaded with NetworkX:
|
|
142
|
+
|
|
143
|
+
```python
|
|
144
|
+
import networkx as nx
|
|
145
|
+
G = nx.read_gexf("data/output/nx_graphs/my_document.gexf")
|
|
146
|
+
```
|
|
147
|
+
|
|
148
|
+
### Baseline chunk extraction
|
|
149
|
+
|
|
150
|
+
Produces flat, retrieval-ready chunks using Docling's hybrid chunker.
|
|
151
|
+
|
|
152
|
+
```python
|
|
153
|
+
from document2graph.baseline_extractor import BaselineExtractor
|
|
154
|
+
|
|
155
|
+
extractor = BaselineExtractor(config)
|
|
156
|
+
|
|
157
|
+
# Process all PDFs; saves JSON to disk
|
|
158
|
+
extractor.generate_baseline_chunks()
|
|
159
|
+
|
|
160
|
+
# Or return chunks in memory
|
|
161
|
+
all_chunks = extractor.generate_baseline_chunks(return_chunks=True)
|
|
162
|
+
# all_chunks["my_document.pdf"]["baseline"] -> List[Chunk]
|
|
163
|
+
# all_chunks["my_document.pdf"]["enriched"] -> List[Chunk] (context-enriched)
|
|
164
|
+
```
|
|
165
|
+
|
|
166
|
+
To process a single file:
|
|
167
|
+
|
|
168
|
+
```python
|
|
169
|
+
chunks = extractor.extract_baseline_chunks(
|
|
170
|
+
filename="my_document.pdf",
|
|
171
|
+
baseline_description="Q1 report chunks",
|
|
172
|
+
enrich=True, # also produce context-enriched variants
|
|
173
|
+
)
|
|
174
|
+
```
|
|
175
|
+
|
|
176
|
+
## Output structure
|
|
177
|
+
|
|
178
|
+
```
|
|
179
|
+
<data_path>/
|
|
180
|
+
├── raw_texts/ # intermediate Docling extraction output
|
|
181
|
+
├── nx_graphs/ # .gexf graph files (one per document)
|
|
182
|
+
├── snippets/ # <filename>_snippets.json (document graph output)
|
|
183
|
+
└── baseline_chunks/ # <filename>_baseline_chunks.json
|
|
184
|
+
```
|
|
185
|
+
|
|
186
|
+
## Data models
|
|
187
|
+
|
|
188
|
+
| Model | Key fields |
|
|
189
|
+
|---|---|
|
|
190
|
+
| `Snippet` | `snippet_id`, `type` (text/image/table), `document_id`, `sequence_no`, `label`, `level`, `page_no`, `bbox`, `text` |
|
|
191
|
+
| `Chunk` | `chunk_id`, `document_id`, `text`, `meta`, `baseline_description`, `embedding` |
|
|
192
|
+
| `Document` | `document_id`, `document_type`, `filename`, `title`, `metadata: DocumentMetadata` |
|
|
193
|
+
| `DocumentMetadata` | `version`, `authors`, `institutions`, `bibliography`, `correspondence` |
|
|
194
|
+
| `MetadataExtractionConfig` | `title_page`, `version`, `authors`, `institutions`, `bibliography`, `correspondence` (each a `MetadataFieldConfig`) |
|
|
195
|
+
| `MetadataFieldConfig` | `label` (search string), `pages` (inclusive 1-based page range, e.g. `(1, 3)`) |
|
|
196
|
+
| `EdgeWeightConfig` | `section`, `text`, `list_item`, `media`, `unreferenced_media`, `root` (edge weights by category) |
|
|
197
|
+
|
|
198
|
+
## Running tests
|
|
199
|
+
|
|
200
|
+
```bash
|
|
201
|
+
pytest
|
|
202
|
+
```
|
|
203
|
+
|
|
204
|
+
Skip slow tests that download models:
|
|
205
|
+
|
|
206
|
+
```bash
|
|
207
|
+
pytest -m "not slow"
|
|
208
|
+
```
|
|
File without changes
|
|
@@ -0,0 +1,82 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import uuid
|
|
3
|
+
|
|
4
|
+
from ..utils.hybrid_chunker_extractor import HybridChunkerExtractor
|
|
5
|
+
from tqdm import tqdm
|
|
6
|
+
import json
|
|
7
|
+
from ..models.Chunk import Chunk
|
|
8
|
+
from ..models.ExtractorConfig import ExtractorConfig
|
|
9
|
+
from typing import List
|
|
10
|
+
|
|
11
|
+
from ..utils.log import Log
|
|
12
|
+
|
|
13
|
+
_log = Log("BaselineExtractor")
|
|
14
|
+
|
|
15
|
+
class BaselineExtractor:
|
|
16
|
+
def __init__(self, config: ExtractorConfig):
|
|
17
|
+
self.logger = _log.logger
|
|
18
|
+
self.pdf_path = config.pdf_path
|
|
19
|
+
self.data_path = config.data_path
|
|
20
|
+
self.raw_text_save_dir = os.path.join(self.data_path, "raw_texts/")
|
|
21
|
+
self.pdfPipelineOptions = config.pdfPipelineOptions
|
|
22
|
+
self.save_json = config.save_json
|
|
23
|
+
self.document_type = config.document_type
|
|
24
|
+
|
|
25
|
+
def _get_clean_filename(self, filename: str) -> str:
|
|
26
|
+
return os.path.splitext(os.path.basename(filename))[0]
|
|
27
|
+
|
|
28
|
+
def extract_baseline_chunks(self, filename: str, baseline_description: str, enrich: bool = False) -> dict[str, List[Chunk]]:
|
|
29
|
+
"""Extract baseline chunks from the raw text files for a given document."""
|
|
30
|
+
sample = os.path.join(self.pdf_path, filename)
|
|
31
|
+
clean_filename = self._get_clean_filename(sample)
|
|
32
|
+
baseline_extractor = HybridChunkerExtractor(
|
|
33
|
+
source=sample,
|
|
34
|
+
merge_peers=True,
|
|
35
|
+
pipeline_options=self.pdfPipelineOptions
|
|
36
|
+
)
|
|
37
|
+
chunks = []
|
|
38
|
+
enriched_chunks = []
|
|
39
|
+
for chunk in baseline_extractor.extract_and_chunk(save_dir=self.raw_text_save_dir, filename=clean_filename):
|
|
40
|
+
meta = chunk.meta.export_json_dict()
|
|
41
|
+
|
|
42
|
+
chunks.append(Chunk(
|
|
43
|
+
chunk_id=str(uuid.uuid4()),
|
|
44
|
+
document_id=filename,
|
|
45
|
+
text=chunk.text,
|
|
46
|
+
meta=meta,
|
|
47
|
+
baseline_description=baseline_description
|
|
48
|
+
))
|
|
49
|
+
if enrich:
|
|
50
|
+
enriched_text = baseline_extractor.chunker.contextualize(chunk)
|
|
51
|
+
enriched_chunks.append(Chunk(
|
|
52
|
+
chunk_id=str(uuid.uuid4()),
|
|
53
|
+
document_id=filename,
|
|
54
|
+
text=enriched_text,
|
|
55
|
+
meta=meta,
|
|
56
|
+
baseline_description=f"Enriched; {baseline_description}"
|
|
57
|
+
))
|
|
58
|
+
return {"baseline": chunks, "enriched": enriched_chunks}
|
|
59
|
+
|
|
60
|
+
def generate_baseline_chunks(self, return_chunks: bool = False) -> dict[str, List[Chunk]] | None:
|
|
61
|
+
"""Generate baseline chunks for all documents in the pdf_path."""
|
|
62
|
+
all_chunks = {}
|
|
63
|
+
for filename in tqdm(os.listdir(self.pdf_path)):
|
|
64
|
+
clean_filename = self._get_clean_filename(filename)
|
|
65
|
+
if filename.endswith(".pdf"):
|
|
66
|
+
self.logger.info(f"Extracting baseline chunks for {filename}...")
|
|
67
|
+
baseline_description = f"Baseline chunks for {filename}"
|
|
68
|
+
chunks = self.extract_baseline_chunks(filename, baseline_description, enrich=True)
|
|
69
|
+
all_chunks[filename] = chunks
|
|
70
|
+
|
|
71
|
+
# by default, the chunks are saved as json
|
|
72
|
+
if self.save_json:
|
|
73
|
+
baseline_json = os.path.join(self.data_path, "baseline_chunks/")
|
|
74
|
+
os.makedirs(baseline_json, exist_ok=True)
|
|
75
|
+
with open(f"{self.raw_text_save_dir}/{clean_filename}_baseline_chunks.json", "w") as f:
|
|
76
|
+
json.dump([chunk.model_dump() for chunk in chunks["baseline"]], f, indent=4)
|
|
77
|
+
if len(chunks["enriched"]) > 0:
|
|
78
|
+
with open(f"{self.raw_text_save_dir}/{clean_filename}_enriched_chunks.json", "w") as f:
|
|
79
|
+
json.dump([chunk.model_dump() for chunk in chunks["enriched"]], f, indent=4)
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
return all_chunks if return_chunks else None
|