olmOCR:用于大语言模型数据集/训练的PDF线性化工具包

olmOCR 是由 Allen Institute for Artificial Intelligence (AI2) 旗下的 AllenNLP 团队开发并维护的开源工具包,旨在帮助研究人员和开发人员构建能够更好地处理 “wild” PDF 文档的大语言模型 (LLM)。

项目地址: https://github.com/allenai/olmocr

主要功能:

  • 自然文本解析: 提供了一种 prompting 策略,可以利用 ChatGPT 4o 实现高质量的自然文本解析。 (通过 buildsilver.py 实现)
  • 评估工具: 提供了一个 side-by-side eval 工具包,用于比较不同pipeline版本的效果。(通过runeval.py实现)
  • 数据清洗: 具备基础的 语言过滤 和 SEO 垃圾信息移除功能。(通过filter.py实现)
  • 模型微调: 提供 Qwen2-VL 和 Molmo-O 模型微调代码。(通过train.py实现)
  • 大规模处理: 支持使用 Sglang 处理数百万个 PDF 文档,并提供使用 finetuned 模型的示例。(通过pipeline.py实现)
  • 结果查看: 提供 Dolma 文档查看器,方便浏览从 PDF 中提取的内容。(通过dolmaviewer.py实现)

安装与使用:

1. 环境要求:

  • NVIDIA GPU (已在 RTX 4090, L40S, A100, H100 上测试)
  • 30GB 可用磁盘空间

2. 依赖安装 (Ubuntu/Debian):

sudo apt-get update
sudo apt-get install poppler-utils ttf-mscorefonts-installer msttcorefonts fonts-crosextra-caladea fonts-crosextra-carlito gsfonts lcdf-typetools

3. 创建 conda 环境并安装 olmocr:

conda create -n olmocr python=3.11
conda activate olmocr
git clone https://github.com/allenai/olmocr.git
cd olmocr
pip install -e .

4. (可选) 安装 sglang (用于GPU 推理):

pip install sgl-kernel==0.0.3.post1 --force-reinstall --no-deps
pip install "sglang[all]==0.4.2" --find-links https://flashinfer.ai/whl/cu124/torch2.4/flashinfer/

5. 本地使用示例:

# 转换单个 PDF
python -m olmocr.pipeline ./localworkspace --pdfs tests/gnarly_pdfs/horribleocr.pdf

# 转换多个 PDF
python -m olmocr.pipeline ./localworkspace --pdfs tests/gnarly_pdfs/*.pdf

结果将以 JSON 格式存储在 ./localworkspace 目录中。

6. 查看结果:

提取的文本以 Dolma 风格 JSONL 格式存储在 ./localworkspace/results 目录中。

cat localworkspace/results/output_*.jsonl

使用 dolmaviewer 命令可以并排查看提取结果和原始 PDF:

python -m olmocr.viewer.dolmaviewer localworkspace/results/output_*.jsonl

然后在浏览器中打开 ./dolma_previews/tests_gnarly_pdfs_horribleocr_pdf.html 文件。

7. 多节点/集群使用:

olmOCR 支持从 AWS S3 读取 PDF 进行大规模处理,并使用 S3 bucket 协调多个 worker 节点的工作。

# 在第一个 worker 节点上运行:
python -m olmocr.pipeline s3://my_s3_bucket/pdfworkspaces/exampleworkspace --pdfs s3://my_s3_bucket/jakep/gnarly_pdfs/*.pdf

# 在后续节点上运行:
python -m olmocr.pipeline s3://my_s3_bucket/pdfworkspaces/exampleworkspace

如果在 Ai2 内部使用 Beaker,可以使用 --beaker flag 将作业提交到集群运行。

8. 查看 pipeline 完整文档:

python -m olmocr.pipeline --help

总结:

olmOCR 提供了一套完整的工具链,可以帮助开发者更轻松地处理 PDF 文档,并将其应用于大语言模型的训练和应用中。无论是进行小规模的测试,还是进行大规模的集群处理,olmOCR 都能提供强大的支持。