OCRFlux

OCRFlux provides several reusable scripts for high-scale document processing, inference, and evaluation. You can leverage these components for your own projects:

• Large-scale Processing: Use pipeline.py to process millions of PDFs using the released model via VLLM.
• Data Conversion: Use jsonl_to_markdown.py to generate final Markdown files from .jsonl data files.
• Inference:
  • inferencer.py for running offline inference using VLLM.
  • server.sh to launch a VLLM server.
  • client.py for running online inference using VLLM.
• Evaluation:
  • eval_page_to_markdown.py: Evaluates single-page parsing.
  • eval_table_to_html.py: Evaluates table parsing.
  • eval_element_merge_detect.py: Evaluates detection of paragraph/table merging.
  • eval_html_table_merge.py: Evaluates table merging tasks.

After running the pipeline, the results are stored in JSONL format. To generate the final human-readable Markdown files, run the ocrflux.jsonl_to_markdown module pointing to your workspace.

python -m ocrflux.jsonl_to_markdown ./localworkspace

You can deploy OCRFlux as a server using the provided shell script, which starts a vLLM server. Once running, you can use the ocrflux.client.request function to interact with it asynchronously.

# Start the server
bash ocrflux/server.sh ChatDOC/OCRFlux-3B 30024

import asyncio
from argparse import Namespace
from ocrflux.client import request

args = Namespace(
    model="/path/to/OCRFlux-3B",
    skip_cross_page_merge=False,
    max_page_retries=1,
    url="http://localhost",
    port=30024,
)

file_path = 'test.pdf'
result = asyncio.run(request(args, file_path))
if result != None:
    document_markdown = result['document_text']
    print(document_markdown)

The pipeline implements two distinct layers of retries to ensure robustness during massive batch jobs:

1. Model/Content Retries (max_page_retries): If a model produces invalid JSON or poor quality results, the process_task function retries the specific page. It uses an increasing temperature strategy (TEMPERATURE_BY_ATTEMPT = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]) to help the model overcome repetitive generation issues.
2. Server/Connection Retries: If a ConnectionError or OSError occurs (e.g., the vLLM server is restarting), the worker performs an exponential backoff (10 * (2**exponential_backoffs)) to allow the server to recover without failing the entire document.
3. Document Discarding: If the number of failed pages exceeds --max_page_error_rate, the entire document is discarded to prevent corrupt data from entering the final dataset.

The OCRFlux pipeline is managed through a command-line interface that supports three primary tasks: converting PDFs to Markdown, merging text elements across pages, and merging HTML tables. The pipeline orchestrates a vLLM server to perform batch inference on documents.

Available Tasks:

• pdf2markdown: Converts PDF or image files into structured Markdown text.
• merge_pages: Uses JSON input to detect and merge text elements across page boundaries.
• merge_tables: Uses JSON input to merge HTML tables across page boundaries.

Basic Usage Pattern:

python -m ocrflux.pipeline <workspace_path> --task <task_name> --data <file_paths>

You can integrate OCRFlux directly into your Python code using the ocrflux.inference.parse function without running a separate vLLM server. This requires initializing a vllm.LLM instance first.

If parsing fails or certain pages are skipped, use the max_page_retries argument in parse() to improve results at the cost of higher inference time.

from vllm import LLM
from ocrflux.inference import parse

file_path = 'test.pdf'
# Initialize vLLM
llm = LLM(model="model_dir/OCRFlux-3B", gpu_memory_utilization=0.8, max_model_len=8192)

# Perform parsing
result = parse(llm, file_path)
if result != None:
    document_markdown = result['document_text']
    print(document_markdown)
    with open('test.md', 'w') as f:
        f.write(document_markdown)
else:
    print("Parse failed.")

When calling request, you must provide an args object (typically a Namespace) with the following configuration keys:

The ocrflux.pipeline module accepts the following arguments for batch processing:

Each line in the JSONL files generated by the pipeline is a JSON object with the following structure:

{
    "orig_path": "str",  // path to the raw file
    "num_pages": "int",  // number of pages
    "document_text": "str", // full Markdown text
    "page_texts": {"int": "str"}, // key: page index, value: page Markdown
    "fallback_pages": ["int"], // indexes of failed pages
}