olmOCR

repository·main·Indexed 11 days ago

https://github.com/allenai/olmocr

A toolkit for converting PDFs and image-based documents into clean Markdown or plain text using a 7B parameter Vision Language Model (VLM). It handles complex layouts, equations, tables, and handwriting while automatically removing headers and footers. Supports local GPU inference, remote vLLM servers, Docker deployment, and large-scale processing via AWS S3.

Tokens
13K
Snippets
49
Records
58
Agent score
98%

What's inside olmOCR

  1. Reusable code components in olmOCR

    main

    The olmOCR repository contains several reusable scripts for advanced document processing, training, and data generation tasks:

    • Natural Text Parsing: olmocr/data/buildsilver.py provides a prompting strategy for high-quality parsing using GPT-4o.
    • Data Filtering: olmocr/filter/filter.py implements basic filtering by language and SEO spam removal.
    • Model Training:
      • olmocr/train/train.py: SFT Finetuning code for Qwen2.5-VL.
      • olmocr/train/grpo_train.py: GRPO RL Trainer.
    • Synthetic Data: olmocr/synth/mine_html_templates.py for synthetic data generation.
    • Large-scale Pipelines: olmocr/pipeline.py for processing millions of PDFs through a finetuned model using vLLM.
    • Visualization: olmocr/viewer/dolmaviewer.py for viewing Dolma docs created from PDFs.
  2. GRPO Training for olmOCR-2-7B-1025

    main

    The olmOCR-2-7B-1025-FP8 model utilizes an additional training step using GRPO Reinforcement Learning based on a synthetic version of olmOCR-bench called olmOCR-synthmix-1025.

    Download Synthetic Dataset

    hf download allenai/olmOCR-synthmix-1025 --repo-type dataset --local-dir olmOCR-synthmix-1025

    Running GRPO Training

    GRPO training is currently specialized for 8xH100 GPU nodes, where one GPU runs vLLM and seven GPUs run training. Use the grpotrainer-beaker-multi-gpu-augusta.sh script to initiate training.

    Example command:

    ./scripts/train/grpotrainer-beaker-multi-gpu-augusta.sh --num-gpus 8 --model_name [MODEL_NAME] --train_bench_data_folder [DATA_FOLDER] --reward_bench 1.0 --reward_front_matter 1.0 --reward_eos 1.0 --beta 0.01 --name [RUN_NAME] --seed 1 --gradient_accumulation_steps 28 --learning_rate 2e-6 --preemptible

    Checkpoint Souping

    To merge multiple seeds/runs into a final checkpoint (souping), pass multiple source checkpoint paths to the prepare_checkpoint script.

  3. Install system dependencies for olmOCR

    main

    Before installing the Python package, you must install poppler-utils and necessary fonts for rendering PDF images on 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
    sudo apt-get update
    sudo apt-get install poppler-utils ttf-mscorefonts-installer msttcorefonts fonts-crosextra-caladea fonts-crosextra-carlito gsfonts lcdf-typetools
  4. Generate boxplots of Elo ratings

    main

    To visualize the Elo ratings, use the draw_boxplots.py script. This script takes the text output generated by the Elo calculation (or the contents of results.txt) and produces a PNG image.

    Usage: Pass the path to the results text file and the desired output filename as arguments.

    python draw_boxplots.py results.txt boxplots.png
  5. Calculate Elo ratings for olmOCR vs other tools

    main

    Use the calculate_elo_ratings.py script to compute Elo ratings based on pairwise judgment data. The input data must be a CSV file (e.g., ratings.csv) containing win/loss counts for different methods.

    Input Data Format: The CSV must include the following columns:

    • MethodA: Name of the first method.
    • MethodB: Name of the second method.
    • A_wins: Number of times MethodA won.
    • B_wins: Number of times MethodB won.
    • A_rate(%): Win rate of MethodA.
    • B_rate(%): Win rate of MethodB.

    Note: pdelf is the identifier used for olmOCR in these datasets.

    CLI Arguments:

    • --num-bootstrap: Number of bootstrap iterations (e.g., 5000).
    • --num-elo-sims: Number of Elo simulations (e.g., 100).
    • --confidence-level: Confidence level for the interval (e.g., 95).
    • --seed: Random seed for reproducibility (e.g., 123).
    python calculate_elo_ratings.py ratings.csv --num-bootstrap 5000 --num-elo-sims 100 --confidence-level 95 --seed 123
  6. Use an external inference server with olmOCR

    main

    If you have a vLLM server or any OpenAI-compatible API running, you can point olmOCR to it using the --server flag. This allows for a lightweight installation without local GPU dependencies.

    Installation:

    pip install olmocr

    Usage:

    olmocr ./localworkspace --server http://remote-server:8000/v1 --model allenai/olmOCR-2-7B-1025-FP8 --markdown --pdfs *.pdf

    Key Arguments for Remote Providers:

    • --server: The OpenAI-compatible endpoint (e.g., https://api.deepinfra.com/v1/openai).
    • --api_key: Your API key (passed via Authorization Bearer header).
    • --max_concurrent_requests: Max concurrent requests in-flight to the provider.
    • --workers: Max number of page groups processed at once.
    • --model: The model identifier (must match the name served by the provider).
    olmocr ./localworkspace --server http://remote-server:8000/v1 --model allenai/olmOCR-2-7B-1025-FP8 --markdown --pdfs *.pdf
  7. Run olmOCR using Docker

    main

    You can use official Docker images to run olmOCR.

    Pull Images:

    • alleninstituteforai/olmocr:latest-with-model (~30GB, includes model)
    • alleninstituteforai/olmocr:latest (Base image, no model)

    Process a single PDF:

    docker run --gpus all \
      -v $(pwd):/workspace \
      alleninstituteforai/olmocr:latest-with-model \
      -c "olmocr /workspace/output --markdown --pdfs /workspace/sample.pdf"

    Process multiple PDFs:

    docker run --gpus all \
      -v /path/to/pdfs:/input \
      -v /path/to/output:/output \
      alleninstituteforai/olmocr:latest-with-model \
      -c "olmocr /output --markdown --pdfs /input/*.pdf"

    Interactive Mode:

    docker run -it --gpus all alleninstituteforai/olmocr:latest-with-model
    docker run --gpus all -v $(pwd):/workspace alleninstituteforai/olmocr:latest-with-model -c "olmocr /workspace/output --markdown --pdfs /workspace/sample.pdf"
  8. Scale olmOCR with AWS S3 and multiple nodes

    main

    To process millions of PDFs across multiple nodes, olmOCR supports reading from and writing to AWS S3 buckets to coordinate work.

    Start the first worker node: This initializes the workspace in S3 and starts converting PDFs.

    olmocr s3://my_s3_bucket/pdfworkspaces/exampleworkspace --pdfs s3://my_s3_bucket/jakep/gnarly_pdfs/*.pdf

    On subsequent worker nodes: Additional nodes will automatically join the existing workspace queue.

    olmocr s3://my_s3_bucket/pdfworkspaces/exampleworkspace
    olmocr s3://my_s3_bucket/pdfworkspaces/exampleworkspace --pdfs s3://my_s3_bucket/jakep/gnarly_pdfs/*.pdf
  9. Set up the olmOCR training environment

    main

    To prepare for training, first set up your standard Python/Conda environment as required for running olmocr. Then, install the training-specific dependencies using the following commands:

    pip install .[train]
    pip install transformers==4.52.4
    pip install flash-attn>=2.8.0.post2 --no-build-isolation
  10. Prepare and quantize model checkpoints

    main

    After training, follow these steps to prepare your model for inference (e.g., with vLLM).

    1. Prepare Checkpoints

    Use prepare_olmocr_checkpoint to convert saved checkpoints. If you trained a LoRA adapter, this script will merge the weights into a full model.

    python -m olmocr.train.prepare_olmocr_checkpoint [source dir]/checkpoint-xxxx [destination]

    While training occurs in BF16, it is recommended to perform FP8 quantization. This reduces memory usage and improves inference speed by ~12% with minimal performance loss.

    python -m olmocr.train.compress_checkpoint --config olmocr/train/quantization_configs/qwen2_5vl_w8a8_fp8.yaml [destination] [destination-FP8]
  11. Install olmOCR via pip

    main

    To install olmOCR, it is recommended to use a clean conda environment. Depending on your use case, choose one of the following installation options:

    Option 1: Remote Inference (Lightweight)

    Use this if you plan to use a remote vLLM server with the --server flag. This avoids installing heavy GPU dependencies like PyTorch.

    pip install olmocr

    Option 2: Local GPU Inference

    Requires a recent NVIDIA GPU (e.g., RTX 4090, L40S, A100, H100) with at least 12 GB of VRAM and 30GB of disk space.

    pip install olmocr[gpu] --extra-index-url https://download.pytorch.org/whl/cu128
    
    # Recommended: Install flash infer for faster inference
    pip install https://download.pytorch.org/whl/cu128/flashinfer/flashinfer_python-0.2.5%2Bcu128torch2.7-cp38-abi3-linux_x86_64.whl

    Option 3: Beaker Cluster Execution

    For submitting jobs to Beaker clusters using the --beaker flag:

    pip install olmocr[beaker]

    Option 4: Benchmark Suite

    To run the olmOCR benchmark suite:

    pip install olmocr[bench]

    Combined Installation Examples

    # GPU + Beaker support
    pip install olmocr[gpu,beaker] --extra-index-url https://download.pytorch.org/whl/cu128
    
    # GPU + Benchmark support
    pip install olmocr[gpu,bench] --extra-index-url https://download.pytorch.org/whl/cu128
    pip install olmocr[gpu] --extra-index-url https://download.pytorch.org/whl/cu128
  12. Launch olmOCR training jobs

    main

    To start a training job, use the olmocr.train.train module with a configuration YAML file.

    Reference Configs:

    • olmOCR-7B-0725-FP8: uses qwen25_vl_olmocrv2_2epoch.yaml
    • olmOCR-7B-0825-FP8: uses qwen25_vl_olmocrv3_rotation_1epoch.yaml
    • olmOCR-2-7B-1025-FP8: uses qwen25_vl_olmocrv4_rotation_1epoch_mix_1025_filtered.yaml
    • Finetuning (LoRA): Use qwen25_vl_olmocrv4_finetuning.yaml for smaller local datasets.

    Note: You must adjust the dataset paths and output_dir in your config file to match your local environment.

    python -m olmocr.train.train --config olmocr/train/configs/v0.4.0/qwen25_vl_olmocrv4_rotation_1epoch_mix_1025_filtered.yaml