Skip to content

feat: add PDF batch processing and merge results to prevent OOM - #86

Open
tiankuolu wants to merge 1 commit into
baidu:mainfrom
tiankuolu:pdf-batch-and-merge
Open

feat: add PDF batch processing and merge results to prevent OOM#86
tiankuolu wants to merge 1 commit into
baidu:mainfrom
tiankuolu:pdf-batch-and-merge

Conversation

@tiankuolu

@tiankuolu tiankuolu commented Jul 29, 2026

Copy link
Copy Markdown

中文说明

问题

当前 Transformers PDF 示例会先渲染全部页面,并将所有页面一起传给模型。对于页数较多的 PDF,页面图片和 OCR 输出可能占用大量内存,最终导致 OOM。

改动

  • 新增 infer_transformers.py
    • 按页或按批次渲染、处理 PDF
    • 通过 --pages-per-batch 控制批次大小,默认每批 1 页
    • 每批处理完成后立即删除临时页面图片
    • 不在内存中累计所有批次的模型返回结果
    • 批次结果完成后再原子写入最终目录
    • 支持使用 --resume 跳过已经完成的批次
  • 新增 merge_results.py
    • 逐批读取并合并 Markdown,不一次性加载所有 OCR 结果
    • 为提取图片生成唯一文件名
    • 自动更新合并后 Markdown 中的图片引用
    • 原子生成 merged.md
  • 更新 README,补充分批推理、断点续跑和结果合并用法
  • 新增 5 个无需 GPU 和真实模型的 CPU 单元测试
  • 将本地输出目录 ocr_output/ 加入 .gitignore

内存行为

内存占用由 --pages-per-batch 和最大单批输出决定,不再随 PDF 总页数持续积累。

单个批次仍然需要能够放入可用内存。较小的批次也可能减少批次边界处的跨页上下文。

用法

python infer_transformers.py your_doc.pdf \
    --output-dir ./ocr_output \
    --pages-per-batch 1

中断后继续:

python infer_transformers.py your_doc.pdf \
    --output-dir ./ocr_output \
    --pages-per-batch 1 \
    --resume

重新合并已有结果:

python merge_results.py ./ocr_output

测试

python -m unittest discover -s tests -v

Ran 5 tests
OK

测试覆盖:

  • PDF 按指定批次处理
  • 临时页面图片清理
  • 推理失败后的临时目录清理
  • 已完成批次的断点续跑
  • Markdown 和图片引用合并
  • 大量结果合并时的内存回归

English

Problem

The current Transformers PDF example renders all pages before passing them to
the model. For PDFs with many pages, rendered images and OCR outputs may consume
significant memory and eventually cause an out-of-memory error.

Changes

  • Add infer_transformers.py
    • Render and process PDFs page by page or in configurable batches
    • Control batch size with --pages-per-batch, defaulting to one page
    • Remove temporary page images after each batch
    • Avoid retaining model outputs from all batches in memory
    • Publish batch output directories only after inference completes
    • Support resuming completed batches with --resume
  • Add merge_results.py
    • Merge Markdown incrementally without loading all OCR results into memory
    • Generate unique names for extracted images
    • Rewrite image references in the merged Markdown
    • Write merged.md atomically
  • Update the README with batching, resume, and result-merging instructions
  • Add five CPU-only unit tests that do not require a GPU or the real model
  • Ignore the local ocr_output/ directory

Memory behavior

Memory usage is bounded by --pages-per-batch and the largest batch output
instead of growing with the total number of PDF pages.

A single batch must still fit in available memory. Smaller batches may also
reduce cross-page context across batch boundaries.

Usage

python infer_transformers.py your_doc.pdf \
    --output-dir ./ocr_output \
    --pages-per-batch 1

Resume an interrupted run:

python infer_transformers.py your_doc.pdf \
    --output-dir ./ocr_output \
    --pages-per-batch 1 \
    --resume

Merge existing results again:

python merge_results.py ./ocr_output

Tests

python -m unittest discover -s tests -v

Ran 5 tests
OK

The tests cover:

  • Configurable PDF page batches
  • Temporary rendered-page cleanup
  • Partial-output cleanup after inference failures
  • Resume behavior for completed batches
  • Incremental Markdown and image-reference merging
  • Merge-memory regression with many result files

Related work

PR #36 focuses on resumability and resource cleanup for the SGLang batch
runner. This PR focuses on the Transformers backend and incremental result
merging.

@tiankuolu
tiankuolu force-pushed the pdf-batch-and-merge branch 3 times, most recently from 5f6c264 to 11e6128 Compare July 29, 2026 15:56
@tiankuolu
tiankuolu force-pushed the pdf-batch-and-merge branch from 11e6128 to 6c4aa13 Compare July 29, 2026 16:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant