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.
Files changed (34) hide show
  1. document2graph-0.1.0/LICENSE +21 -0
  2. document2graph-0.1.0/PKG-INFO +243 -0
  3. document2graph-0.1.0/README.md +208 -0
  4. document2graph-0.1.0/document2graph/__init__.py +0 -0
  5. document2graph-0.1.0/document2graph/baseline_extractor/__init__.py +3 -0
  6. document2graph-0.1.0/document2graph/baseline_extractor/extractor.py +82 -0
  7. document2graph-0.1.0/document2graph/document2graph_extractor/__init__.py +3 -0
  8. document2graph-0.1.0/document2graph/document2graph_extractor/extractor.py +146 -0
  9. document2graph-0.1.0/document2graph/document2graph_extractor/snippet_graph_constructor.py +424 -0
  10. document2graph-0.1.0/document2graph/models/Chunk.py +10 -0
  11. document2graph-0.1.0/document2graph/models/Document.py +10 -0
  12. document2graph-0.1.0/document2graph/models/DocumentMetadata.py +23 -0
  13. document2graph-0.1.0/document2graph/models/EdgeWeightConfig.py +18 -0
  14. document2graph-0.1.0/document2graph/models/ExtractorConfig.py +14 -0
  15. document2graph-0.1.0/document2graph/models/ImageSnippetNode.py +21 -0
  16. document2graph-0.1.0/document2graph/models/Snippet.py +19 -0
  17. document2graph-0.1.0/document2graph/models/TableSnippetNode.py +6 -0
  18. document2graph-0.1.0/document2graph/models/TextSnippet.py +6 -0
  19. document2graph-0.1.0/document2graph/models/TextSnippetNode.py +19 -0
  20. document2graph-0.1.0/document2graph/models/__init__.py +25 -0
  21. document2graph-0.1.0/document2graph/utils/__init__.py +0 -0
  22. document2graph-0.1.0/document2graph/utils/base_extractor.py +127 -0
  23. document2graph-0.1.0/document2graph/utils/doc_to_tree_lib.py +52 -0
  24. document2graph-0.1.0/document2graph/utils/hybrid_chunker_extractor.py +26 -0
  25. document2graph-0.1.0/document2graph/utils/log.py +40 -0
  26. document2graph-0.1.0/document2graph.egg-info/PKG-INFO +243 -0
  27. document2graph-0.1.0/document2graph.egg-info/SOURCES.txt +32 -0
  28. document2graph-0.1.0/document2graph.egg-info/dependency_links.txt +1 -0
  29. document2graph-0.1.0/document2graph.egg-info/requires.txt +14 -0
  30. document2graph-0.1.0/document2graph.egg-info/top_level.txt +1 -0
  31. document2graph-0.1.0/pyproject.toml +53 -0
  32. document2graph-0.1.0/setup.cfg +4 -0
  33. document2graph-0.1.0/tests/test_document2graph_extractor.py +71 -0
  34. 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,3 @@
1
+ from .extractor import BaselineExtractor
2
+
3
+ __all__ = ["BaselineExtractor"]
@@ -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
@@ -0,0 +1,3 @@
1
+ from .extractor import DocumentGraphExtractor
2
+
3
+ __all__ = ["DocumentGraphExtractor"]