Unlimited-OCR:一次解析40页PDF,无需GPU熔化
如果你曾经构建过文档解析流水线,你一定熟悉这套流程:把PDF拆成单页,让每一页通过OCR模型,再将输出拼接回去。然后还得写一堆胶水代码,去修复在页面边界被切断的表格、丢失了所属章节的标题、以及被遗弃在距引用三页之外的脚注。
百度的Unlimited-OCR 赌的就是你根本不需要做这些事。它在单次前向传播中就能解析数十页内容,并且采用MIT 许可证。
本文首先介绍这个模型究竟做了什么不同的事,然后给出三种具体的运行方式:一个快速的Transformers脚本、一个用于实际吞吐量的SGLang 服务器、以及一个Docker 镜像(如果你完全不想碰Python环境的话)。
它解决的问题
标准Transformer解码有一个大多数OCR基准测试悄悄隐藏的成本:KV缓存会随着你生成的每个token不断增长。
KV缓存是模型的短期记忆。模型生成的每个token都会被追加到其中,以便后续token能向前做注意力。对于一个写三段文字的聊天机器人来说,这没问题。但对于一个将40页技术手册转录成30000个Markdown token的OCR模型来说,问题就大了。内存占用攀升,注意力成本随之增加,生成过程运行越久越慢。大约到第八页时,你的吞吐量图就不再是一条直线,而是一道悬崖。
行业内的变通办法是分块处理:处理一页,丢弃缓存,处理下一页,并接受模型对先前内容一无所知。
Unlimited-OCR直接攻击缓存增长问题。团队将解码器中的每个注意力层替换成了他们称之为参考滑动窗口注意力(R-SWA) 的结构。
直观理解类似于人类打字员的工作方式:你并不会把每个已经打过的词都记在脑子里,你只记住最近一两句话,并不断回头参照源文档。R-SWA做的也是同样的事:解码器保留一个固定大小的最近生成token窗口,但保留对原始图像token的永久访问权。KV缓存实现为一个固定容量的队列,当新token到来时,窗口中最旧的token被驱逐。
结果是缓存大小恒定而不是不断增长。无论你在第500个token还是第30000个token上,内存和每token延迟都保持平稳。
你实际能得到什么
- 总参数量3B,实际激活500M。它是一个混合专家(MoE)模型,因此推理时的计算成本更接近500M模型而非3B模型。
- 最大输出长度32K,这使得多页单次解析成为可能。
- OmniDocBench v1.5上93.23分,比其继续训练的基线DeepSeek-OCR高出约6.2个百分点。值得注意,因为效率提升通常以精度为代价,但这里没有。
- MIT许可证。商用无限制。
架构是DeepSeek-OCR的DeepEncoder(SAM-ViT-B加CLIP-L)馈送给一个MoE解码器。编码器激进的视觉token压缩使得图像侧的缓存足够小,让整个方案得以工作。
论文中有一个易被忽略的要点:R-SWA并非OCR专用。它是一种通用的注意力机制,适用于任何长跨度转录任务,包括ASR。
开始之前
你需要一张NVIDIA GPU。官方仓库不支持CPU或Apple Silicon。
权重是bfloat16,仅加载模型就需要约6GB,尚未包括激活值和图像token。12GB显存卡是单图像工作的合理下限;对于32K上下文的多次运行,你需要更多余量。维护者测试环境为Python 3.12.3和CUDA 12.9。
如果你只想先看看输出效果,Hugging Face Space 上有一个你可以直接上传文件试用的空间。
路径1:Transformers(从这里开始)
这是最快得到结果的方式。设置环境:
python -m venv .venv
source .venv/bin/activate
pip install torch==2.10.0 torchvision==0.25.0 transformers==4.57.1
pip install Pillow==12.1.1 matplotlib==3.10.8 einops==0.8.2
pip install addict==2.4.0 easydict==1.13 pymupdf==1.27.2.2 psutil==7.2.2
请锁定这些版本。模型通过trust_remote_code携带自定义建模代码,而这段代码是针对这些特定版本编写的。版本漂移会引发令人困惑的导入错误,而非干净的失败提示。
现在处理单张图片:
import torch
from transformers import AutoModel, AutoTokenizer
model_name = 'baidu/Unlimited-OCR'
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(
model_name,
trust_remote_code=True,
use_safetensors=True,
torch_dtype=torch.bfloat16,
)
model = model.eval().cuda()
model.infer(
tokenizer,
prompt='<image>document parsing.',
image_file='your_image.jpg',
output_path='your/output/dir',
base_size=1024, image_size=640, crop_mode=True, # gundam mode
max_length=32768,
no_repeat_ngram_size=35, ngram_window=128,
save_results=True,
)
其中有两点比看起来更重要:
- 提示必须以
<image>开头。这不是装饰,而是视觉特征被拼接进去的占位token。去掉它,你会得到看起来像是模型在幻觉一份从未见过的文档的输出。 no_repeat_ngram_size和ngram_window是承重的。在重复布局(如目录或价目表)上的长跨度生成,可能使模型陷入无限循环,愉快地重复同一行。这些参数阻止滑动窗口内任何35-gram重复。不要因为它们看起来像调优噪声就删除。
gundam vs base
单图像推理提供两种配置,命名并非一目了然:
- gundam:
base_size=1024, image_size=640, crop_mode=True— 单张图像,尤其是密集图片。将图像裁剪成瓦片并分别处理,因此小文字得以保留。 - base:
base_size=1024, image_size=1024, crop_mode=False— 一次性处理整张图像。多页面时必须使用此模式。
多页面和PDF路径只支持base模式。这是硬性约束,而非默认选项。
多页面
model.infer_multi(
tokenizer,
prompt='<image>Multi page parsing.',
image_files=['page1.png', 'page2.png', 'page3.png'],
output_path='your/output/dir',
image_size=1024,
max_length=32768,
no_repeat_ngram_size=35, ngram_window=1024,
save_results=True,
)
注意这里ngram_window从128跳到了1024。跨越多个页面时,确实存在更多重复结构,因此重复检测窗口需要扩大以跟上。
没有直接的PDF入口点。你需要先栅格化,再将图像送入infer_multi:
import os, tempfile
import fitz # PyMuPDF
def pdf_to_images(pdf_path, dpi=300):
doc = fitz.open(pdf_path)
tmp_dir = tempfile.mkdtemp(prefix='pdf_ocr_')
mat = fitz.Matrix(dpi / 72, dpi / 72)
paths = []
for i, page in enumerate(doc):
out = os.path.join(tmp_dir, f'page_{i+1:04d}.png')
page.get_pixmap(matrix=mat).save(out)
paths.append(out)
doc.close()
return paths
model.infer_multi(
tokenizer,
prompt='<image>Multi page parsing.',
image_files=pdf_to_images('your_doc.pdf', dpi=300),
output_path='your/output/dir',
image_size=1024,
max_length=32768,
no_repeat_ngram_size=35, ngram_window=1024,
save_results=True,
)
300 DPI是推荐默认值。降低分辨率以节省内存会损害小文字和表格线,而这通常正是你关心的内容。
路径2:SGLang 服务器(用于实际负载)
Transformers 适合评估模型。对于处理并发请求的服务,请运行 SGLang 并通过兼容 OpenAI 的 API 与其通信。
配置环境。注意 SGLang 以本地 wheel 形式提供在仓库中,而非来自 PyPI:
uv venv --python 3.12
source .venv/bin/activate
uv pip install wheel/sglang-0.0.0.dev11416+g92e8bb79e-py3-none-any.whl
uv pip install kernels==0.11.7
uv pip install pymupdf==1.27.2.2
小提示:README 文本中提到固定 kernels==0.9.0,而紧随其后的命令块安装的是 0.11.7。请按照命令块执行。如果遇到与 kernel 相关的错误,首先对照当前仓库状态检查此版本不匹配问题。
启动服务器:
python -m sglang.launch_server
--model baidu/Unlimited-OCR
--served-model-name Unlimited-OCR
--attention-backend fa3
--page-size 1
--mem-fraction-static 0.8
--context-length 32768
--enable-custom-logit-processor
--disable-overlap-schedule
--skip-server-warmup
--host 0.0.0.0
--port 10000
以下标志是必须的:
--enable-custom-logit-processor—— no-repeat-ngram 处理器以自定义 logit processor 运行。没有此标志,引用它的请求会失败。--attention-backend fa3—— FlashAttention 3,需要 Hopper 级硬件。在较老的 GPU 上需要使用其他后端。--page-size 1和--disable-overlap-schedule—— R-SWA 的缓存淘汰策略与通常的 paged-attention 和重叠调度优化不兼容。
然后是客户端代码。图像以 base64 data URL 传入,采用标准 OpenAI 视觉格式,并附带一些模型特定参数:
import base64, json, os
import requests
from sglang.srt.sampling.custom_logit_processor import (
DeepseekOCRNoRepeatNGramLogitProcessor,
)
server_url = "http://127.0.0.1:10000"
session = requests.Session()
session.trust_env = False
def encode_image(image_path):
ext = os.path.splitext(image_path)[1].lower()
mime = "image/jpeg" if ext in (".jpg", ".jpeg") else f"image/{ext.lstrip('.')}"
with open(image_path, "rb") as f:
data = base64.b64encode(f.read()).decode("utf-8")
return {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{data}"}}
def generate(prompt, image_paths, image_mode, ngram_window):
content = [{"type": "text", "text": prompt}] + [
encode_image(p) for p in image_paths
]
payload = {
"model": "Unlimited-OCR",
"messages": [{"role": "user", "content": content}],
"temperature": 0,
"skip_special_tokens": False,
"images_config": {"image_mode": image_mode},
"custom_logit_processor": DeepseekOCRNoRepeatNGramLogitProcessor.to_str(),
"custom_params": {"ngram_size": 35, "window_size": ngram_window},
"stream": True,
}
response = session.post(
f"{server_url}/v1/chat/completions",
headers={"Content-Type": "application/json"},
data=json.dumps(payload),
timeout=1200,
stream=True,
)
response.raise_for_status()
chunks = []
for line in response.iter_lines(chunk_size=1, decode_unicode=True):
if not line or not line.startswith("data: "):
continue
data = line[len("data: "):]
if data == "[DONE]":
break
delta = json.loads(data)["choices"][0].get("delta", {}).get("content", "")
if delta:
print(delta, end="", flush=True)
chunks.append(delta)
return "".join(chunks)
generate("document parsing.", ["your_image.jpg"],
image_mode="gundam", ngram_window=128)
temperature: 0 和 skip_special_tokens: False 都是有意为之。这是转录任务而非生成任务,因此任何采样随机性都是纯粹的缺点。而特殊标记携带了你希望在输出中保留的布局结构。
1200 秒超时也不是多虑。长文档处理需要一定时间,这正是使用流式输出的原因。
批量处理
仓库提供了 infer.py,它会为你启动 SGLang 服务器并发送并发请求:
# 图片目录
python infer.py
--image_dir ./examples/images
--output_dir ./outputs
--concurrency 8
--image_mode gundam
# PDF 文件
python infer.py
--pdf ./examples/document.pdf
--output_dir ./outputs
--concurrency 8
--image_mode gundam
有用的额外参数:--model_dir 接受本地路径或 Hugging Face ID,--gpu 设置 CUDA_VISIBLE_DEVICES,--server_log 将服务器输出保存到可读的位置。
开始时并发数要设低一些。8 是文档中的示例值,但实际合适数值取决于你的显存和文档长度。
路径 3:通过 Docker 使用 vLLM
如果你想完全跳过环境管理:
# 默认,CUDA 13.0
docker pull vllm/vllm-openai:unlimited-ocr
# Hopper GPU,CUDA 12.9
docker pull vllm/vllm-openai:unlimited-ocr-cu129
docker run --rm --gpus all --network host --ipc host
vllm/vllm-openai:unlimited-ocr
baidu/Unlimited-OCR
--trust-remote-code
--logits_processors vllm.model_executor.models.unlimited_ocr:NGramPerReqLogitsProcessor
--no-enable-prefix-caching
--mm-processor-cache-gb 0
与 SGLang 类似,ngram 处理器必须在服务器启动时注册。前缀缓存和多模态处理器缓存均被禁用,因为 R-SWA 的固定窗口缓存使这些优化的假设失效。
完整细节请参考官方 vLLM 指南。
参数速查表
- mode:单图 —
gundam或base;多页/PDF — 仅base - image_size:单图 — 640 (gundam) / 1024 (base);多页/PDF — 1024
- crop_mode:单图 —
True(gundam) /False(base);多页/PDF —False - ngram_window:单图 — 128;多页/PDF — 1024
- no_repeat_ngram_size:单图 — 35;多页/PDF — 35
- max_length:单图 — 32768;多页/PDF — 32768
- prompt:单图 —
<image>document parsing.;多页/PDF —<image>Multi page parsing.
何时不适合使用此工具
如果你需要边界框和逐词置信度分数,请选择其他工具,因为这是端到端模型,直接输出 Markdown,而非检测+识别流水线。同样,如果只是 OCR 短收据或单行文本,使用 3B 模型配合 GPU 相比 PaddleOCR 或 Tesseract 显得大材小用。此外,如果没有 NVIDIA GPU,目前也没有支持路径。
总结
这里有趣的主张并非基准测试分数本身,而在于固定大小的 KV 缓存使模型在长文档上既更快又更准确——而性能优化通常会在某些方面牺牲质量。
如果 R-SWA 如作者所暗示的那样具有通用性,那么同样的技巧可以应用于任何需要模型将长输入转录为长输出并同时保留源信息的任务。长音频 ASR 显然是下一个候选场景。