Reuse one MarkdownDocSerializer per document in item extraction
This commit is contained in:
parent
54ac3450b5
commit
9f4dda9254
3 changed files with 151 additions and 12 deletions
|
|
@ -5,6 +5,7 @@
|
||||||
|
|
||||||
- zstd compression uses a fresh `zstandard` compressor/decompressor per call instead of shared module-level singletons, fixing process segfaults when ingester workers compress documents concurrently (Python < 3.14).
|
- zstd compression uses a fresh `zstandard` compressor/decompressor per call instead of shared module-level singletons, fixing process segfaults when ingester workers compress documents concurrently (Python < 3.14).
|
||||||
- Ingestion compresses the DoclingDocument structure directly from the in-memory dict, removing one full-size serialized copy from per-document peak memory.
|
- Ingestion compresses the DoclingDocument structure directly from the in-memory dict, removing one full-size serialized copy from per-document peak memory.
|
||||||
|
- Document item extraction builds one `MarkdownDocSerializer` per document (reused across tables) and derives description-less picture text from captions, instead of constructing a serializer per picture and table.
|
||||||
|
|
||||||
## [0.61.0] - 2026-06-23
|
## [0.61.0] - 2026-06-23
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,6 @@
|
||||||
import base64
|
import base64
|
||||||
from typing import TYPE_CHECKING
|
from collections.abc import Callable
|
||||||
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
|
@ -53,18 +54,33 @@ def _decode_picture_bytes(item: "PictureItem") -> bytes | None:
|
||||||
return base64.b64decode(encoded, validate=False)
|
return base64.b64decode(encoded, validate=False)
|
||||||
|
|
||||||
|
|
||||||
def extract_item_text(item: "NodeItem", docling_doc: "DoclingDocument") -> str | None:
|
def _picture_caption_text(item: "PictureItem", docling_doc: "DoclingDocument") -> str:
|
||||||
|
"""Join a picture's caption texts with spaces, preserving word boundaries."""
|
||||||
|
return " ".join(
|
||||||
|
text
|
||||||
|
for caption in item.captions
|
||||||
|
if (text := caption.resolve(docling_doc).text.strip())
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def extract_item_text(
|
||||||
|
item: "NodeItem",
|
||||||
|
docling_doc: "DoclingDocument",
|
||||||
|
*,
|
||||||
|
get_serializer: Callable[[], Any] | None = None,
|
||||||
|
) -> str | None:
|
||||||
"""Extract text content from a DocItem.
|
"""Extract text content from a DocItem.
|
||||||
|
|
||||||
Handles different item types:
|
Handles different item types:
|
||||||
- TextItem, SectionHeaderItem, etc.: Use .text attribute
|
- TextItem, SectionHeaderItem, etc.: Use .text attribute
|
||||||
- TableItem: Use export_to_markdown() for table content
|
- TableItem: serialize to markdown. ``get_serializer`` supplies a reused
|
||||||
|
``MarkdownDocSerializer`` (see ``extract_items``); when absent a one-off
|
||||||
|
serializer is built so direct calls keep working.
|
||||||
- PictureItem: Prefer the VLM description (when picture_description is on)
|
- PictureItem: Prefer the VLM description (when picture_description is on)
|
||||||
so pictures carry meaningful prose into chunk text and survive
|
so pictures carry meaningful prose into chunk text and survive
|
||||||
``expand_with_items``' ``if item.text:`` filter; otherwise fall back to
|
``expand_with_items``' ``if item.text:`` filter; otherwise fall back to
|
||||||
a placeholder markdown export (no base64).
|
the picture's caption text.
|
||||||
"""
|
"""
|
||||||
from docling_core.types.doc.base import ImageRefMode
|
|
||||||
from docling_core.types.doc.document import PictureItem, TableItem
|
from docling_core.types.doc.document import PictureItem, TableItem
|
||||||
|
|
||||||
if text := getattr(item, "text", None):
|
if text := getattr(item, "text", None):
|
||||||
|
|
@ -73,15 +89,19 @@ def extract_item_text(item: "NodeItem", docling_doc: "DoclingDocument") -> str |
|
||||||
if isinstance(item, PictureItem):
|
if isinstance(item, PictureItem):
|
||||||
if description := _picture_description_text(item):
|
if description := _picture_description_text(item):
|
||||||
return description
|
return description
|
||||||
return item.export_to_markdown(
|
return _picture_caption_text(item, docling_doc)
|
||||||
docling_doc,
|
|
||||||
image_mode=ImageRefMode.PLACEHOLDER,
|
|
||||||
image_placeholder="",
|
|
||||||
)
|
|
||||||
|
|
||||||
if isinstance(item, TableItem):
|
if isinstance(item, TableItem):
|
||||||
try:
|
try:
|
||||||
return item.export_to_markdown(docling_doc)
|
if get_serializer is None:
|
||||||
|
from docling_core.transforms.serializer.markdown import (
|
||||||
|
MarkdownDocSerializer,
|
||||||
|
)
|
||||||
|
|
||||||
|
serializer = MarkdownDocSerializer(doc=docling_doc)
|
||||||
|
else:
|
||||||
|
serializer = get_serializer()
|
||||||
|
return serializer.serialize(item=item).text
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
@ -116,11 +136,23 @@ def extract_items(
|
||||||
existing = existing_picture_data or {}
|
existing = existing_picture_data or {}
|
||||||
items: list[DocumentItem] = []
|
items: list[DocumentItem] = []
|
||||||
|
|
||||||
|
serializer: Any = None
|
||||||
|
|
||||||
|
def get_serializer() -> Any:
|
||||||
|
nonlocal serializer
|
||||||
|
if serializer is None:
|
||||||
|
from docling_core.transforms.serializer.markdown import (
|
||||||
|
MarkdownDocSerializer,
|
||||||
|
)
|
||||||
|
|
||||||
|
serializer = MarkdownDocSerializer(doc=docling_doc)
|
||||||
|
return serializer
|
||||||
|
|
||||||
for position, (item, level) in enumerate(docling_doc.iterate_items()):
|
for position, (item, level) in enumerate(docling_doc.iterate_items()):
|
||||||
label = getattr(item, "label", None)
|
label = getattr(item, "label", None)
|
||||||
label_str = str(label.value) if hasattr(label, "value") else str(label or "")
|
label_str = str(label.value) if hasattr(label, "value") else str(label or "")
|
||||||
|
|
||||||
text = extract_item_text(item, docling_doc) or ""
|
text = extract_item_text(item, docling_doc, get_serializer=get_serializer) or ""
|
||||||
|
|
||||||
page_numbers: list[int] = []
|
page_numbers: list[int] = []
|
||||||
if prov := getattr(item, "prov", None):
|
if prov := getattr(item, "prov", None):
|
||||||
|
|
|
||||||
|
|
@ -140,6 +140,112 @@ class TestExtractItemText:
|
||||||
assert items == []
|
assert items == []
|
||||||
|
|
||||||
|
|
||||||
|
def _doc_with_captioned_picture(*captions: str):
|
||||||
|
from docling_core.types.doc.document import DoclingDocument, ImageRef
|
||||||
|
from docling_core.types.doc.labels import DocItemLabel
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
doc = DoclingDocument(name="pics")
|
||||||
|
caption_items = [doc.add_text(label=DocItemLabel.CAPTION, text=c) for c in captions]
|
||||||
|
img = ImageRef.from_pil(Image.new("RGB", (8, 8), "red"), dpi=72)
|
||||||
|
doc.add_picture(image=img, caption=caption_items[0] if caption_items else None)
|
||||||
|
pic = doc.pictures[0]
|
||||||
|
for extra in caption_items[1:]:
|
||||||
|
pic.captions.append(extra.get_ref())
|
||||||
|
return doc, pic
|
||||||
|
|
||||||
|
|
||||||
|
def _doc_with_tables(n: int):
|
||||||
|
from docling_core.types.doc.document import DoclingDocument, TableCell, TableData
|
||||||
|
|
||||||
|
doc = DoclingDocument(name="tables")
|
||||||
|
for _ in range(n):
|
||||||
|
cells = [
|
||||||
|
TableCell(
|
||||||
|
text=f"r{r}c{c}",
|
||||||
|
row_span=1,
|
||||||
|
col_span=1,
|
||||||
|
start_row_offset_idx=r,
|
||||||
|
end_row_offset_idx=r + 1,
|
||||||
|
start_col_offset_idx=c,
|
||||||
|
end_col_offset_idx=c + 1,
|
||||||
|
)
|
||||||
|
for r in range(2)
|
||||||
|
for c in range(2)
|
||||||
|
]
|
||||||
|
doc.add_table(data=TableData(num_rows=2, num_cols=2, table_cells=cells))
|
||||||
|
return doc
|
||||||
|
|
||||||
|
|
||||||
|
class TestExtractItemTextPictures:
|
||||||
|
"""Description-less pictures derive text from captions, without a serializer."""
|
||||||
|
|
||||||
|
def test_multiple_captions_space_joined(self):
|
||||||
|
doc, pic = _doc_with_captioned_picture("First", "Second")
|
||||||
|
assert extract_item_text(pic, doc) == "First Second"
|
||||||
|
|
||||||
|
def test_description_wins_over_captions(self):
|
||||||
|
from docling_core.types.doc.document import DescriptionMetaField, PictureMeta
|
||||||
|
|
||||||
|
doc, pic = _doc_with_captioned_picture("A caption")
|
||||||
|
pic.meta = PictureMeta(description=DescriptionMetaField(text="A red square."))
|
||||||
|
assert extract_item_text(pic, doc) == "A red square."
|
||||||
|
|
||||||
|
def test_picture_path_does_not_export_markdown(self, monkeypatch):
|
||||||
|
from docling_core.types.doc.document import PictureItem
|
||||||
|
|
||||||
|
doc, pic = _doc_with_captioned_picture("Only caption")
|
||||||
|
|
||||||
|
def _boom(*args, **kwargs):
|
||||||
|
raise AssertionError("picture path must not build a serializer")
|
||||||
|
|
||||||
|
monkeypatch.setattr(PictureItem, "export_to_markdown", _boom)
|
||||||
|
assert extract_item_text(pic, doc) == "Only caption"
|
||||||
|
|
||||||
|
def test_picture_path_never_requests_serializer(self):
|
||||||
|
doc, pic = _doc_with_captioned_picture("Only caption")
|
||||||
|
|
||||||
|
def _explode():
|
||||||
|
raise AssertionError("picture path must not request a serializer")
|
||||||
|
|
||||||
|
assert extract_item_text(pic, doc, get_serializer=_explode) == "Only caption"
|
||||||
|
|
||||||
|
|
||||||
|
class TestExtractItemsTableSerializer:
|
||||||
|
"""Table text is unchanged; one serializer is reused across the whole pass."""
|
||||||
|
|
||||||
|
def test_table_text_matches_export_to_markdown(self):
|
||||||
|
doc = _doc_with_tables(1)
|
||||||
|
expected = doc.tables[0].export_to_markdown(doc)
|
||||||
|
items = extract_items("doc-1", doc)
|
||||||
|
table_items = [i for i in items if i.label == "table"]
|
||||||
|
assert len(table_items) == 1
|
||||||
|
assert table_items[0].text == expected
|
||||||
|
|
||||||
|
def test_one_serializer_built_for_multiple_tables(self, monkeypatch):
|
||||||
|
import docling_core.transforms.serializer.markdown as md
|
||||||
|
|
||||||
|
count = {"n": 0}
|
||||||
|
base = md.MarkdownDocSerializer
|
||||||
|
|
||||||
|
class Counting(base):
|
||||||
|
def __init__(self, *args, **kwargs):
|
||||||
|
count["n"] += 1
|
||||||
|
super().__init__(*args, **kwargs)
|
||||||
|
|
||||||
|
doc = _doc_with_tables(3)
|
||||||
|
monkeypatch.setattr(md, "MarkdownDocSerializer", Counting)
|
||||||
|
|
||||||
|
items = extract_items("doc-1", doc)
|
||||||
|
assert sum(1 for i in items if i.label == "table") == 3
|
||||||
|
assert count["n"] == 1
|
||||||
|
|
||||||
|
def test_direct_table_call_builds_one_off_serializer(self):
|
||||||
|
doc = _doc_with_tables(1)
|
||||||
|
expected = doc.tables[0].export_to_markdown(doc)
|
||||||
|
assert extract_item_text(doc.tables[0], doc) == expected
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
class TestDocumentItemRepository:
|
class TestDocumentItemRepository:
|
||||||
async def test_create_and_get_range(self, temp_db_path):
|
async def test_create_and_get_range(self, temp_db_path):
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue