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 tempfile
|
||||
from pathlib import Path
|
||||
|
|
@ -167,6 +168,38 @@ async def convert(
|
|||
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(
|
||||
config: AppConfig,
|
||||
docling_document: "DoclingDocument",
|
||||
|
|
@ -196,31 +229,14 @@ async def chunk(
|
|||
c.order = i
|
||||
return text_chunks
|
||||
|
||||
picture_chunks = build_picture_chunks(
|
||||
return await asyncio.to_thread(
|
||||
_merge_picture_chunks,
|
||||
docling_document,
|
||||
document_id=document_id,
|
||||
existing_picture_data=existing_picture_data,
|
||||
text_chunks,
|
||||
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(
|
||||
docling_document: "DoclingDocument",
|
||||
|
|
|
|||
Loading…
Reference in a new issue