add optimized path for image text that avoids re-parsing whole document
This commit is contained in:
parent
65f1339aa0
commit
3b3d097c46
6 changed files with 96 additions and 9 deletions
|
|
@ -118,7 +118,12 @@ async def _store_document_with_chunks(
|
||||||
Handles versioning/rollback on failure.
|
Handles versioning/rollback on failure.
|
||||||
"""
|
"""
|
||||||
chunks = await ensure_chunks_embedded(client._config, chunks, client.embedder)
|
chunks = await ensure_chunks_embedded(client._config, chunks, client.embedder)
|
||||||
items = await asyncio.to_thread(extract_items, "", docling_document)
|
items = await asyncio.to_thread(
|
||||||
|
extract_items,
|
||||||
|
"",
|
||||||
|
docling_document,
|
||||||
|
fast_picture_text=client._config.processing.fast_picture_text,
|
||||||
|
)
|
||||||
|
|
||||||
async with client.store._write_lock:
|
async with client.store._write_lock:
|
||||||
versions = await client.store.current_table_versions()
|
versions = await client.store.current_table_versions()
|
||||||
|
|
@ -175,7 +180,11 @@ async def _update_document_with_chunks(
|
||||||
items: list[DocumentItem] | None = None
|
items: list[DocumentItem] | None = None
|
||||||
if docling_document is not None:
|
if docling_document is not None:
|
||||||
items = await asyncio.to_thread(
|
items = await asyncio.to_thread(
|
||||||
extract_items, document.id, docling_document, existing_picture_data
|
extract_items,
|
||||||
|
document.id,
|
||||||
|
docling_document,
|
||||||
|
existing_picture_data,
|
||||||
|
fast_picture_text=client._config.processing.fast_picture_text,
|
||||||
)
|
)
|
||||||
|
|
||||||
async with client.store._write_lock:
|
async with client.store._write_lock:
|
||||||
|
|
@ -282,8 +291,13 @@ async def _store_documents_with_chunks(
|
||||||
for _, chunks, _ in prepared
|
for _, chunks, _ in prepared
|
||||||
]
|
]
|
||||||
|
|
||||||
|
fast_picture_text = client._config.processing.fast_picture_text
|
||||||
|
|
||||||
def _extract_all_items():
|
def _extract_all_items():
|
||||||
return [extract_items("", d) for _, _, d in prepared]
|
return [
|
||||||
|
extract_items("", d, fast_picture_text=fast_picture_text)
|
||||||
|
for _, _, d in prepared
|
||||||
|
]
|
||||||
|
|
||||||
all_item_lists = await asyncio.to_thread(_extract_all_items)
|
all_item_lists = await asyncio.to_thread(_extract_all_items)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -173,11 +173,13 @@ def _merge_picture_chunks(
|
||||||
text_chunks: list[Chunk],
|
text_chunks: list[Chunk],
|
||||||
document_id: str | None,
|
document_id: str | None,
|
||||||
existing_picture_data: dict[str, bytes] | None,
|
existing_picture_data: dict[str, bytes] | None,
|
||||||
|
fast_picture_text: bool,
|
||||||
) -> list[Chunk]:
|
) -> list[Chunk]:
|
||||||
picture_chunks = build_picture_chunks(
|
picture_chunks = build_picture_chunks(
|
||||||
docling_document,
|
docling_document,
|
||||||
document_id=document_id,
|
document_id=document_id,
|
||||||
existing_picture_data=existing_picture_data,
|
existing_picture_data=existing_picture_data,
|
||||||
|
fast_picture_text=fast_picture_text,
|
||||||
)
|
)
|
||||||
|
|
||||||
if not picture_chunks:
|
if not picture_chunks:
|
||||||
|
|
@ -235,6 +237,7 @@ async def chunk(
|
||||||
text_chunks,
|
text_chunks,
|
||||||
document_id,
|
document_id,
|
||||||
existing_picture_data,
|
existing_picture_data,
|
||||||
|
config.processing.fast_picture_text,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -243,6 +246,7 @@ def build_picture_chunks(
|
||||||
*,
|
*,
|
||||||
document_id: str | None = None,
|
document_id: str | None = None,
|
||||||
existing_picture_data: dict[str, bytes] | None = None,
|
existing_picture_data: dict[str, bytes] | None = None,
|
||||||
|
fast_picture_text: bool = True,
|
||||||
) -> list[Chunk]:
|
) -> list[Chunk]:
|
||||||
"""Emit one synthetic ``Chunk`` per ``PictureItem`` with available bytes.
|
"""Emit one synthetic ``Chunk`` per ``PictureItem`` with available bytes.
|
||||||
|
|
||||||
|
|
@ -271,7 +275,12 @@ def build_picture_chunks(
|
||||||
if picture_data is None:
|
if picture_data is None:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
text = extract_item_text(picture, docling_document) or ""
|
text = (
|
||||||
|
extract_item_text(
|
||||||
|
picture, docling_document, fast_picture_text=fast_picture_text
|
||||||
|
)
|
||||||
|
or ""
|
||||||
|
)
|
||||||
|
|
||||||
page_numbers: list[int] = []
|
page_numbers: list[int] = []
|
||||||
for p in picture.prov:
|
for p in picture.prov:
|
||||||
|
|
|
||||||
|
|
@ -191,6 +191,14 @@ class ProcessingConfig(BaseModel):
|
||||||
"""When a PDF carries `/EmbeddedFiles`, ingest each attachment as a separate
|
"""When a PDF carries `/EmbeddedFiles`, ingest each attachment as a separate
|
||||||
Document linked back to the wrapper via ``metadata.parent_uri``. Cap depth
|
Document linked back to the wrapper via ``metadata.parent_uri``. Cap depth
|
||||||
at 3 to bound nested-attachment recursion."""
|
at 3 to bound nested-attachment recursion."""
|
||||||
|
fast_picture_text: bool = True
|
||||||
|
"""Picture chunk/item text uses the cheap caption accessor
|
||||||
|
(``FloatingItem.caption_text``) instead of rendering each picture through a
|
||||||
|
full-document Markdown serializer. The serializer re-validates the entire
|
||||||
|
DoclingDocument per picture (bbox clamping over every table cell), which is
|
||||||
|
O(pictures x document) at ingest. Set False to restore exact
|
||||||
|
``export_to_markdown`` text. Only governs ingestion; rebuild and store
|
||||||
|
migrations always use the fast path."""
|
||||||
auto_title: bool = False
|
auto_title: bool = False
|
||||||
title_model: ModelConfig = Field(
|
title_model: ModelConfig = Field(
|
||||||
default_factory=lambda: ModelConfig(
|
default_factory=lambda: ModelConfig(
|
||||||
|
|
|
||||||
|
|
@ -53,7 +53,12 @@ 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 extract_item_text(
|
||||||
|
item: "NodeItem",
|
||||||
|
docling_doc: "DoclingDocument",
|
||||||
|
*,
|
||||||
|
fast_picture_text: bool = True,
|
||||||
|
) -> str | None:
|
||||||
"""Extract text content from a DocItem.
|
"""Extract text content from a DocItem.
|
||||||
|
|
||||||
Handles different item types:
|
Handles different item types:
|
||||||
|
|
@ -62,9 +67,15 @@ def extract_item_text(item: "NodeItem", docling_doc: "DoclingDocument") -> str |
|
||||||
- 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.
|
||||||
|
|
||||||
|
With ``fast_picture_text`` (the default), a description-less picture's text
|
||||||
|
comes from ``caption_text`` — O(captions). When False, it is rendered via
|
||||||
|
``export_to_markdown`` with an empty image placeholder, which builds a
|
||||||
|
full-document serializer (re-validating the whole DoclingDocument) per
|
||||||
|
picture; the resulting text is still just the caption for an
|
||||||
|
annotation-less picture.
|
||||||
"""
|
"""
|
||||||
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,6 +84,10 @@ 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
|
||||||
|
if fast_picture_text:
|
||||||
|
return item.caption_text(docling_doc)
|
||||||
|
from docling_core.types.doc.base import ImageRefMode
|
||||||
|
|
||||||
return item.export_to_markdown(
|
return item.export_to_markdown(
|
||||||
docling_doc,
|
docling_doc,
|
||||||
image_mode=ImageRefMode.PLACEHOLDER,
|
image_mode=ImageRefMode.PLACEHOLDER,
|
||||||
|
|
@ -96,6 +111,8 @@ def extract_items(
|
||||||
document_id: str,
|
document_id: str,
|
||||||
docling_doc: "DoclingDocument",
|
docling_doc: "DoclingDocument",
|
||||||
existing_picture_data: dict[str, bytes] | None = None,
|
existing_picture_data: dict[str, bytes] | None = None,
|
||||||
|
*,
|
||||||
|
fast_picture_text: bool = True,
|
||||||
) -> list[DocumentItem]:
|
) -> list[DocumentItem]:
|
||||||
"""Extract document items from a DoclingDocument for the items table.
|
"""Extract document items from a DoclingDocument for the items table.
|
||||||
|
|
||||||
|
|
@ -120,7 +137,10 @@ def extract_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, fast_picture_text=fast_picture_text)
|
||||||
|
or ""
|
||||||
|
)
|
||||||
|
|
||||||
page_numbers: list[int] = []
|
page_numbers: list[int] = []
|
||||||
if prov := getattr(item, "prov", None):
|
if prov := getattr(item, "prov", None):
|
||||||
|
|
|
||||||
|
|
@ -35,7 +35,7 @@ logging.getLogger("vcr.cassette").setLevel(logging.WARNING)
|
||||||
@pytest.fixture(scope="session")
|
@pytest.fixture(scope="session")
|
||||||
def qa_corpus() -> list[dict[str, str]]:
|
def qa_corpus() -> list[dict[str, str]]:
|
||||||
corpus_path = Path(__file__).parent / "data" / "qa_corpus.json"
|
corpus_path = Path(__file__).parent / "data" / "qa_corpus.json"
|
||||||
with open(corpus_path) as f:
|
with open(corpus_path, encoding="utf-8") as f:
|
||||||
return json.load(f)
|
return json.load(f)
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -281,6 +281,42 @@ def test_set_docling_with_page_images():
|
||||||
assert "1" in pages
|
assert "1" in pages
|
||||||
|
|
||||||
|
|
||||||
|
def test_extract_item_text_picture_modes(monkeypatch):
|
||||||
|
"""fast_picture_text (default) returns the caption via caption_text and
|
||||||
|
never builds a Markdown serializer; fast_picture_text=False matches the
|
||||||
|
placeholder export."""
|
||||||
|
from docling_core.types.doc.base import ImageRefMode
|
||||||
|
from docling_core.types.doc.document import (
|
||||||
|
DoclingDocument,
|
||||||
|
ImageRef,
|
||||||
|
PictureItem,
|
||||||
|
)
|
||||||
|
from docling_core.types.doc.labels import DocItemLabel
|
||||||
|
from PIL import Image as PILImageModule
|
||||||
|
|
||||||
|
from haiku.rag.store.models.document_item import extract_item_text
|
||||||
|
|
||||||
|
img = PILImageModule.new("RGB", (8, 8), "blue")
|
||||||
|
doc = DoclingDocument(name="caption_test")
|
||||||
|
caption = doc.add_text(label=DocItemLabel.CAPTION, text="A figure caption")
|
||||||
|
doc.add_picture(image=ImageRef.from_pil(img, dpi=72), caption=caption)
|
||||||
|
picture = doc.pictures[0]
|
||||||
|
|
||||||
|
# Slow path equals the placeholder markdown export.
|
||||||
|
expected_slow = picture.export_to_markdown(
|
||||||
|
doc, image_mode=ImageRefMode.PLACEHOLDER, image_placeholder=""
|
||||||
|
)
|
||||||
|
assert extract_item_text(picture, doc, fast_picture_text=False) == expected_slow
|
||||||
|
|
||||||
|
# Fast path returns the caption and must not construct a serializer.
|
||||||
|
def _boom(*args, **kwargs):
|
||||||
|
raise AssertionError("export_to_markdown should not run on the fast path")
|
||||||
|
|
||||||
|
monkeypatch.setattr(PictureItem, "export_to_markdown", _boom)
|
||||||
|
assert extract_item_text(picture, doc) == "A figure caption"
|
||||||
|
assert extract_item_text(picture, doc, fast_picture_text=True) == "A figure caption"
|
||||||
|
|
||||||
|
|
||||||
def test_get_page_images():
|
def test_get_page_images():
|
||||||
"""get_page_images returns requested pages from docling_pages blob."""
|
"""get_page_images returns requested pages from docling_pages blob."""
|
||||||
import json
|
import json
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue