Skip to content

Commit 77449dc

Browse files
fenixc9fenixc9
committed
feat(api): expose task page progress
1 parent 3e60291 commit 77449dc

5 files changed

Lines changed: 286 additions & 2 deletions

File tree

‎mineru/backend/hybrid/hybrid_analyze.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -886,6 +886,15 @@ def _close_images(images_list):
886886
pass
887887

888888

889+
def _emit_progress(progress_callback, **payload):
890+
if progress_callback is None:
891+
return
892+
try:
893+
progress_callback(**payload)
894+
except Exception:
895+
logger.exception("Progress callback failed")
896+
897+
889898
def doc_analyze(
890899
pdf_bytes,
891900
image_writer: DataWriter | None,
@@ -901,6 +910,7 @@ def doc_analyze(
901910
):
902911
effort = _validate_parse_effort(effort)
903912
effective_image_analysis = _resolve_effective_image_analysis(effort, image_analysis)
913+
progress_callback = kwargs.pop("progress_callback", None)
904914
client_side_output_generation = bool(
905915
kwargs.pop("client_side_output_generation", False)
906916
)
@@ -932,6 +942,14 @@ def doc_analyze(
932942
f'Hybrid processing-window run. page_count={page_count}, '
933943
f'window_size={configured_window_size}, total_windows={total_windows}'
934944
)
945+
_emit_progress(
946+
progress_callback,
947+
page_total=page_count,
948+
page_current=0,
949+
window_index=0,
950+
window_total=total_windows,
951+
text=f"0/{page_count} pages",
952+
)
935953

936954
batch_ratio = get_batch_ratio(device) if not _ocr_enable else 1
937955

@@ -956,6 +974,16 @@ def doc_analyze(
956974
f'pages {window_start + 1}-{window_end + 1}/{page_count} '
957975
f'({len(images_pil_list)} pages)'
958976
)
977+
_emit_progress(
978+
progress_callback,
979+
page_total=page_count,
980+
page_current=window_start,
981+
window_index=window_index + 1,
982+
window_total=total_windows,
983+
window_start=window_start + 1,
984+
window_end=window_end + 1,
985+
text=f"pages {window_start + 1}-{window_end + 1}/{page_count}",
986+
)
959987
images_layout_res, hybrid_pipeline_model = _predict_layout_for_window(
960988
images_pil_list,
961989
inline_formula_enable,
@@ -1058,6 +1086,16 @@ def doc_analyze(
10581086
_ocr_enable=_ocr_enable,
10591087
progress_bar=progress_bar,
10601088
)
1089+
_emit_progress(
1090+
progress_callback,
1091+
page_total=page_count,
1092+
page_current=window_end + 1,
1093+
window_index=window_index + 1,
1094+
window_total=total_windows,
1095+
window_start=window_start + 1,
1096+
window_end=window_end + 1,
1097+
text=f"pages {window_start + 1}-{window_end + 1}/{page_count}",
1098+
)
10611099
last_append_end_time = time.time()
10621100
finally:
10631101
_close_images(images_list)
@@ -1109,6 +1147,7 @@ async def aio_doc_analyze(
11091147
):
11101148
effort = _validate_parse_effort(effort)
11111149
effective_image_analysis = _resolve_effective_image_analysis(effort, image_analysis)
1150+
progress_callback = kwargs.pop("progress_callback", None)
11121151
client_side_output_generation = bool(
11131152
kwargs.pop("client_side_output_generation", False)
11141153
)
@@ -1140,6 +1179,14 @@ async def aio_doc_analyze(
11401179
f'Hybrid processing-window run. page_count={page_count}, '
11411180
f'window_size={configured_window_size}, total_windows={total_windows}'
11421181
)
1182+
_emit_progress(
1183+
progress_callback,
1184+
page_total=page_count,
1185+
page_current=0,
1186+
window_index=0,
1187+
window_total=total_windows,
1188+
text=f"0/{page_count} pages",
1189+
)
11431190

11441191
batch_ratio = get_batch_ratio(device) if not _ocr_enable else 1
11451192

@@ -1163,6 +1210,16 @@ async def aio_doc_analyze(
11631210
f'pages {window_start + 1}-{window_end + 1}/{page_count} '
11641211
f'({len(images_pil_list)} pages)'
11651212
)
1213+
_emit_progress(
1214+
progress_callback,
1215+
page_total=page_count,
1216+
page_current=window_start,
1217+
window_index=window_index + 1,
1218+
window_total=total_windows,
1219+
window_start=window_start + 1,
1220+
window_end=window_end + 1,
1221+
text=f"pages {window_start + 1}-{window_end + 1}/{page_count}",
1222+
)
11661223
images_layout_res, hybrid_pipeline_model = await asyncio.to_thread(
11671224
_predict_layout_for_window,
11681225
images_pil_list,
@@ -1272,6 +1329,16 @@ async def aio_doc_analyze(
12721329
_ocr_enable=_ocr_enable,
12731330
progress_bar=progress_bar,
12741331
)
1332+
_emit_progress(
1333+
progress_callback,
1334+
page_total=page_count,
1335+
page_current=window_end + 1,
1336+
window_index=window_index + 1,
1337+
window_total=total_windows,
1338+
window_start=window_start + 1,
1339+
window_end=window_end + 1,
1340+
text=f"pages {window_start + 1}-{window_end + 1}/{page_count}",
1341+
)
12751342
last_append_end_time = time.time()
12761343
finally:
12771344
_close_images(images_list)

‎mineru/backend/pipeline/pipeline_analyze.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,16 @@
3030

3131
os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1' # 让mps可以fallback
3232

33+
34+
def _emit_progress(progress_callback, **payload):
35+
if progress_callback is None:
36+
return
37+
try:
38+
progress_callback(**payload)
39+
except Exception:
40+
logger.exception("Progress callback failed")
41+
42+
3343
class ModelSingleton:
3444
_instance = None
3545
_models = {}
@@ -163,6 +173,7 @@ def doc_analyze_streaming(
163173
formula_enable=True,
164174
table_enable=True,
165175
client_side_output_generation=False,
176+
progress_callback=None,
166177
):
167178
if not (len(pdf_bytes_list) == len(image_writer_list) == len(lang_list)):
168179
raise ValueError("pdf_bytes_list, image_writer_list, and lang_list must have the same length")
@@ -210,6 +221,15 @@ def doc_analyze_streaming(
210221
f'Pipeline processing-window multi-file run. doc_count={len(doc_contexts)}, '
211222
f'total_pages={total_pages}, window_size={window_size}, total_batches={total_batches}'
212223
)
224+
_emit_progress(
225+
progress_callback,
226+
file_index=0,
227+
page_total=total_pages,
228+
page_current=0,
229+
window_index=0,
230+
window_total=total_batches,
231+
text=f"0/{total_pages} pages",
232+
)
213233

214234
_emit_zero_page_contexts(
215235
doc_contexts,
@@ -266,6 +286,15 @@ def doc_analyze_streaming(
266286
f'{processed_pages + len(batch_images)}/{total_pages} pages, '
267287
f'batch_pages={len(batch_images)}, doc_slices={_format_doc_slices(batch_slices)}'
268288
)
289+
_emit_progress(
290+
progress_callback,
291+
file_index=0,
292+
page_total=total_pages,
293+
page_current=processed_pages,
294+
window_index=batch_index,
295+
window_total=total_batches,
296+
text=f"{processed_pages}/{total_pages} pages",
297+
)
269298

270299
try:
271300
batch_results = batch_image_analyze(
@@ -311,6 +340,15 @@ def doc_analyze_streaming(
311340

312341
last_append_end_time = time.time()
313342
processed_pages += len(batch_images)
343+
_emit_progress(
344+
progress_callback,
345+
file_index=0,
346+
page_total=total_pages,
347+
page_current=processed_pages,
348+
window_index=batch_index,
349+
window_total=total_batches,
350+
text=f"{processed_pages}/{total_pages} pages",
351+
)
314352
finally:
315353
if progress_bar is not None:
316354
progress_bar.close()

‎mineru/backend/vlm/vlm_analyze.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -420,6 +420,15 @@ def _close_images(images_list):
420420
pass
421421

422422

423+
def _emit_progress(progress_callback, **payload):
424+
if progress_callback is None:
425+
return
426+
try:
427+
progress_callback(**payload)
428+
except Exception:
429+
logger.exception("Progress callback failed")
430+
431+
423432
def doc_analyze(
424433
pdf_bytes,
425434
image_writer: DataWriter | None,
@@ -430,6 +439,7 @@ def doc_analyze(
430439
image_analysis: bool = True,
431440
**kwargs,
432441
):
442+
progress_callback = kwargs.pop("progress_callback", None)
433443
client_side_output_generation = bool(
434444
kwargs.pop("client_side_output_generation", False)
435445
)
@@ -454,6 +464,14 @@ def doc_analyze(
454464
f'VLM processing-window run. page_count={page_count}, '
455465
f'window_size={configured_window_size}, total_windows={total_windows}'
456466
)
467+
_emit_progress(
468+
progress_callback,
469+
page_total=page_count,
470+
page_current=0,
471+
window_index=0,
472+
window_total=total_windows,
473+
text=f"0/{page_count} pages",
474+
)
457475

458476
infer_start = time.time()
459477
progress_bar = None
@@ -475,6 +493,16 @@ def doc_analyze(
475493
f'pages {window_start + 1}-{window_end + 1}/{page_count} '
476494
f'({len(images_pil_list)} pages)'
477495
)
496+
_emit_progress(
497+
progress_callback,
498+
page_total=page_count,
499+
page_current=window_start,
500+
window_index=window_index + 1,
501+
window_total=total_windows,
502+
window_start=window_start + 1,
503+
window_end=window_end + 1,
504+
text=f"pages {window_start + 1}-{window_end + 1}/{page_count}",
505+
)
478506
with predictor_execution_guard(predictor):
479507
window_results = predictor.batch_two_step_extract(
480508
images=images_pil_list,
@@ -498,6 +526,16 @@ def doc_analyze(
498526
page_start_index=window_start,
499527
progress_bar=progress_bar,
500528
)
529+
_emit_progress(
530+
progress_callback,
531+
page_total=page_count,
532+
page_current=window_end + 1,
533+
window_index=window_index + 1,
534+
window_total=total_windows,
535+
window_start=window_start + 1,
536+
window_end=window_end + 1,
537+
text=f"pages {window_start + 1}-{window_end + 1}/{page_count}",
538+
)
501539
last_append_end_time = time.time()
502540
finally:
503541
_close_images(images_list)
@@ -530,6 +568,7 @@ async def aio_doc_analyze(
530568
image_analysis: bool = True,
531569
**kwargs,
532570
):
571+
progress_callback = kwargs.pop("progress_callback", None)
533572
client_side_output_generation = bool(
534573
kwargs.pop("client_side_output_generation", False)
535574
)
@@ -554,6 +593,14 @@ async def aio_doc_analyze(
554593
f'VLM processing-window run. page_count={page_count}, '
555594
f'window_size={configured_window_size}, total_windows={total_windows}'
556595
)
596+
_emit_progress(
597+
progress_callback,
598+
page_total=page_count,
599+
page_current=0,
600+
window_index=0,
601+
window_total=total_windows,
602+
text=f"0/{page_count} pages",
603+
)
557604

558605
infer_start = time.time()
559606
progress_bar = None
@@ -574,6 +621,16 @@ async def aio_doc_analyze(
574621
f'pages {window_start + 1}-{window_end + 1}/{page_count} '
575622
f'({len(images_pil_list)} pages)'
576623
)
624+
_emit_progress(
625+
progress_callback,
626+
page_total=page_count,
627+
page_current=window_start,
628+
window_index=window_index + 1,
629+
window_total=total_windows,
630+
window_start=window_start + 1,
631+
window_end=window_end + 1,
632+
text=f"pages {window_start + 1}-{window_end + 1}/{page_count}",
633+
)
577634
async with aio_predictor_execution_guard(predictor):
578635
window_results = await predictor.aio_batch_two_step_extract(
579636
images=images_pil_list,
@@ -597,6 +654,16 @@ async def aio_doc_analyze(
597654
page_start_index=window_start,
598655
progress_bar=progress_bar,
599656
)
657+
_emit_progress(
658+
progress_callback,
659+
page_total=page_count,
660+
page_current=window_end + 1,
661+
window_index=window_index + 1,
662+
window_total=total_windows,
663+
window_start=window_start + 1,
664+
window_end=window_end + 1,
665+
text=f"pages {window_start + 1}-{window_end + 1}/{page_count}",
666+
)
600667
last_append_end_time = time.time()
601668
finally:
602669
_close_images(images_list)

0 commit comments

Comments
 (0)