Thread picture chunk merging off the asyncio event loop
Extracts build_picture_chunks + iterate_items position mapping into a sync helper and wraps it with asyncio.to_thread so image-heavy documents don't block the event loop during chunking. Fixes #456.
This commit is contained in:
parent
0fcf91fcf7
commit
ccab012747
1 changed files with 38 additions and 22 deletions
|
|
@ -1,3 +1,4 @@
|
||||||
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
import tempfile
|
import tempfile
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
@ -167,6 +168,38 @@ async def convert(
|
||||||
return doc
|
return doc
|
||||||
|
|
||||||
|
|
||||||
|
def _merge_picture_chunks(
|
||||||
|
docling_document,
|
||||||
|
text_chunks: list,
|
||||||
|
document_id: str | None,
|
||||||
|
existing_picture_data: dict | None,
|
||||||
|
) -> list:
|
||||||
|
picture_chunks = build_picture_chunks(
|
||||||
|
docling_document,
|
||||||
|
document_id=document_id,
|
||||||
|
existing_picture_data=existing_picture_data,
|
||||||
|
)
|
||||||
|
|
||||||
|
if not picture_chunks:
|
||||||
|
for i, c in enumerate(text_chunks):
|
||||||
|
c.order = i
|
||||||
|
return text_chunks
|
||||||
|
|
||||||
|
positions = {
|
||||||
|
item.self_ref: pos
|
||||||
|
for pos, (item, _level) in enumerate(docling_document.iterate_items())
|
||||||
|
}
|
||||||
|
|
||||||
|
def first_pos(c):
|
||||||
|
refs = (c.metadata or {}).get("doc_item_refs") or []
|
||||||
|
return positions.get(refs[0], len(positions)) if refs else len(positions)
|
||||||
|
|
||||||
|
merged = sorted(text_chunks + picture_chunks, key=first_pos)
|
||||||
|
for i, c in enumerate(merged):
|
||||||
|
c.order = i
|
||||||
|
return merged
|
||||||
|
|
||||||
|
|
||||||
async def chunk(
|
async def chunk(
|
||||||
config: AppConfig,
|
config: AppConfig,
|
||||||
docling_document: "DoclingDocument",
|
docling_document: "DoclingDocument",
|
||||||
|
|
@ -196,31 +229,14 @@ async def chunk(
|
||||||
c.order = i
|
c.order = i
|
||||||
return text_chunks
|
return text_chunks
|
||||||
|
|
||||||
picture_chunks = build_picture_chunks(
|
return await asyncio.to_thread(
|
||||||
|
_merge_picture_chunks,
|
||||||
docling_document,
|
docling_document,
|
||||||
document_id=document_id,
|
text_chunks,
|
||||||
existing_picture_data=existing_picture_data,
|
document_id,
|
||||||
|
existing_picture_data,
|
||||||
)
|
)
|
||||||
|
|
||||||
if not picture_chunks:
|
|
||||||
for i, c in enumerate(text_chunks):
|
|
||||||
c.order = i
|
|
||||||
return text_chunks
|
|
||||||
|
|
||||||
positions = {
|
|
||||||
item.self_ref: pos
|
|
||||||
for pos, (item, _level) in enumerate(docling_document.iterate_items())
|
|
||||||
}
|
|
||||||
|
|
||||||
def first_pos(c: Chunk) -> int:
|
|
||||||
refs = (c.metadata or {}).get("doc_item_refs") or []
|
|
||||||
return positions.get(refs[0], len(positions)) if refs else len(positions)
|
|
||||||
|
|
||||||
merged = sorted(text_chunks + picture_chunks, key=first_pos)
|
|
||||||
for i, c in enumerate(merged):
|
|
||||||
c.order = i
|
|
||||||
return merged
|
|
||||||
|
|
||||||
|
|
||||||
def build_picture_chunks(
|
def build_picture_chunks(
|
||||||
docling_document: "DoclingDocument",
|
docling_document: "DoclingDocument",
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue