Coverage for src / quber / core / parsers / docling_parser.py: 77%
83 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""
2DoclingParser — docling pipeline machinery behind the `Parser` ABC.
4`DoclingParser` is abstract: it owns the converter construction, the run
5fingerprint, and the `parse` template, but defers the pipeline-option bundle
6to a concrete subclass. The configuration is a *type*, not a runtime flag:
8- `TunedFinancialParser` — the canonical configuration selected by QUE-217 and
9 adopted by QUE-218 (TableFormerMode.ACCURATE, DoclingParse backend, OCR,
10 picture classification, CUDA, page_batch_size=32). The default everywhere.
11- `LegacyDoclingParser` — the pre-QUE-218 behaviour (PdfPipelineOptions with
12 only do_ocr / do_table_structure / do_cell_matching set). Kept for
13 regression comparison and as a fallback; select it with `--preset legacy`.
15Use `parser_for_preset(preset)` to map a preset string to the right class.
17The header-consolidation post-processor that builds `TableMetadata` lives in
18`core.consolidation` and operates on the returned DoclingDocument. That keeps
19these classes focused on a single job: source -> DoclingDocument.
20"""
22import os
23from abc import abstractmethod
24from pathlib import Path
25from typing import Any, ClassVar
27import logfire
28from docling.backend.docling_parse_backend import DoclingParseDocumentBackend
29from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions
30from docling.datamodel.base_models import InputFormat
31from docling.datamodel.pipeline_options import (
32 PdfPipelineOptions,
33 RapidOcrOptions,
34 TableFormerMode,
35)
36from docling.datamodel.settings import settings
37from docling.document_converter import DocumentConverter, PdfFormatOption
38from docling.utils.accelerator_utils import decide_device
40from quber.core.parsers.base import Parser
41from quber.core.parsers.result import ParseResult
43CANONICAL_PAGE_BATCH_SIZE = 32
46class DoclingParser(Parser):
47 """Abstract docling parser. Subclasses define the pipeline-option bundle.
49 Subclasses set the `preset` and `backend_label` class attributes and
50 implement `build_pipeline_options`. Everything else (converter, perf
51 settings, fingerprint, the `parse` template) is shared.
52 """
54 #: Stable name of the configuration, surfaced in the run fingerprint.
55 preset: ClassVar[str]
56 #: How the PDF backend is reported in the fingerprint.
57 backend_label: ClassVar[str]
59 def __init__(
60 self,
61 do_ocr: bool = False,
62 do_table_structure: bool = True,
63 do_cell_matching: bool = True,
64 page_batch_size: int = CANONICAL_PAGE_BATCH_SIZE,
65 accelerator_device: AcceleratorDevice = AcceleratorDevice.AUTO,
66 num_threads: int | None = None,
67 ) -> None:
68 self.do_ocr = do_ocr
69 self.do_table_structure = do_table_structure
70 self.do_cell_matching = do_cell_matching
71 self.page_batch_size = page_batch_size
72 self.accelerator_device = accelerator_device
73 self.num_threads = num_threads if num_threads is not None else (os.cpu_count() or 4)
75 @abstractmethod
76 def build_pipeline_options(self) -> PdfPipelineOptions:
77 """Return the docling pipeline options for this configuration."""
78 ...
80 def build_converter(self, pipeline_options: PdfPipelineOptions) -> DocumentConverter:
81 """Default converter (no explicit backend). Overridden by presets that pin one."""
82 return DocumentConverter(
83 format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)}
84 )
86 def apply_perf_settings(self) -> None:
87 """Hook for process-wide perf tweaks. No-op by default."""
89 def fingerprint(self, pipeline_options: PdfPipelineOptions) -> dict[str, Any]:
90 device: Any = pipeline_options.accelerator_options.device
91 device_name = device.name if hasattr(device, "name") else str(device)
92 mode: Any = pipeline_options.table_structure_options.mode # type: ignore[union-attr]
93 mode_name = mode.name if hasattr(mode, "name") else str(mode)
94 ocr_engine = pipeline_options.ocr_options.__class__.__name__
95 ocr_cuda = isinstance(pipeline_options.ocr_options, RapidOcrOptions) and bool(
96 (pipeline_options.ocr_options.rapidocr_params or {}).get(
97 "EngineConfig.onnxruntime.use_cuda", False
98 )
99 )
100 resolved_device = decide_device(self.accelerator_device.value)
101 return {
102 "preset": self.preset,
103 "table_mode": mode_name,
104 "do_cell_matching": pipeline_options.table_structure_options.do_cell_matching, # type: ignore[union-attr]
105 "do_ocr": pipeline_options.do_ocr,
106 "do_picture_classification": pipeline_options.do_picture_classification,
107 "generate_picture_images": pipeline_options.generate_picture_images,
108 "backend": self.backend_label,
109 "accelerator_device_requested": device_name,
110 "accelerator_device_resolved": resolved_device,
111 "num_threads": pipeline_options.accelerator_options.num_threads,
112 "page_batch_size": self.page_batch_size,
113 "ocr_engine": ocr_engine,
114 "ocr_cuda": ocr_cuda,
115 }
117 def parse(self, source: Path) -> ParseResult:
118 pipeline_options = self.build_pipeline_options()
119 # Keep docling's parsed-page cells (text, box, from_ocr, confidence) alive
120 # past document assembly. Off by default, so docling frees them; they are
121 # the OCR-vs-native signal the reconciliation needs, so capture them for
122 # every preset.
123 pipeline_options.generate_parsed_pages = True
124 converter = self.build_converter(pipeline_options)
125 self.apply_perf_settings()
127 fingerprint = self.fingerprint(pipeline_options)
128 with logfire.span("docling_parser.parse", source=str(source), **fingerprint):
129 result = converter.convert(str(source))
130 return ParseResult.from_conversion(result)
133class LegacyDoclingParser(DoclingParser):
134 """Pre-QUE-218 docling defaults: bare PdfPipelineOptions, default backend."""
136 preset = "legacy"
137 backend_label = "default"
139 def build_pipeline_options(self) -> PdfPipelineOptions:
140 pipeline_options = PdfPipelineOptions()
141 pipeline_options.do_table_structure = self.do_table_structure
142 pipeline_options.table_structure_options.do_cell_matching = self.do_cell_matching # type: ignore[misc]
143 pipeline_options.do_ocr = self.do_ocr
144 return pipeline_options
147class TunedFinancialParser(DoclingParser):
148 """Canonical QUE-218 configuration for dense SEC financial documents."""
150 preset = "tuned-financial"
151 backend_label = DoclingParseDocumentBackend.__name__
153 def build_pipeline_options(self) -> PdfPipelineOptions:
154 pipeline_options = PdfPipelineOptions()
155 pipeline_options.do_table_structure = self.do_table_structure
156 pipeline_options.table_structure_options.do_cell_matching = self.do_cell_matching # type: ignore[misc]
157 pipeline_options.do_ocr = True
158 pipeline_options.table_structure_options.mode = TableFormerMode.ACCURATE # type: ignore[misc]
159 pipeline_options.do_picture_classification = True
160 pipeline_options.generate_picture_images = True
161 pipeline_options.accelerator_options = AcceleratorOptions(
162 device=self.accelerator_device,
163 num_threads=self.num_threads,
164 )
165 # QUE-219: docling propagates accelerator_options.device=CUDA to RapidOCR's
166 # paddle/torch engines but not the onnxruntime engine (which is the default).
167 # Flip EngineConfig.onnxruntime.use_cuda explicitly so the OCR pass actually
168 # runs on GPU. CUDA 12 runtime libs preloaded in quber/__init__.py.
169 # Gate on the *resolved* device (decide_device handles AUTO -> cuda:0/cpu/mps)
170 # so non-GPU hosts (e.g., Intel MacBook) skip the override and run OCR on CPU.
171 resolved_device = decide_device(self.accelerator_device.value)
172 if resolved_device.startswith("cuda"):
173 pipeline_options.ocr_options = RapidOcrOptions(
174 rapidocr_params={
175 "EngineConfig.onnxruntime.use_cuda": True,
176 "EngineConfig.onnxruntime.cuda_ep_cfg.device_id": 0,
177 },
178 )
179 return pipeline_options
181 def build_converter(self, pipeline_options: PdfPipelineOptions) -> DocumentConverter:
182 return DocumentConverter(
183 format_options={
184 InputFormat.PDF: PdfFormatOption(
185 pipeline_options=pipeline_options,
186 backend=DoclingParseDocumentBackend,
187 )
188 }
189 )
191 def apply_perf_settings(self) -> None:
192 settings.perf.page_batch_size = self.page_batch_size
195#: Preset name -> concrete parser class. The CLI `--preset` choices mirror these keys.
196PARSERS_BY_PRESET: dict[str, type[DoclingParser]] = {
197 "tuned-financial": TunedFinancialParser,
198 "legacy": LegacyDoclingParser,
199}
202def parser_for_preset(preset: str, **kwargs: Any) -> DoclingParser:
203 """Instantiate the parser for a preset name (forwards kwargs to the class)."""
204 try:
205 cls = PARSERS_BY_PRESET[preset]
206 except KeyError:
207 valid = ", ".join(sorted(PARSERS_BY_PRESET))
208 raise ValueError(f"Unknown parser preset {preset!r}; valid presets: {valid}") from None
209 return cls(**kwargs)