allenai--olmocr/olmocr/data/build_openai_batch_from_olmocrmix.py

#!/usr/bin/env python3
"""
Build OpenAI batch requests from OLMoCR-mix folder structure.

This script processes the folder structure created by prepare_olmocrmix.py
and generates OpenAI batch API requests for processing PDFs.
"""

import argparse
import json
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Dict, Generator, Optional, Tuple

from pypdf import PdfReader
from tqdm import tqdm

from olmocr.data.renderpdf import (
    get_png_dimensions_from_base64,
    render_pdf_to_base64png,
)
from olmocr.prompts.prompts import (
    build_openai_silver_data_prompt_v3_simple,
    openai_response_format_schema,
)

TARGET_IMAGE_DIM = 2048
MAX_FILE_SIZE = 99 * 1024 * 1024  # 99MB in bytes


def validate_single_page_pdf(pdf_path: Path) -> bool:
    """
    Validate that a PDF has exactly one page.

    Args:
        pdf_path: Path to the PDF file

    Returns:
        True if PDF has exactly one page, False otherwise
    """
    try:
        pdf = PdfReader(pdf_path)
        return len(pdf.pages) == 1
    except Exception as e:
        print(f"Error reading PDF {pdf_path}: {e}")
        return False


def build_custom_id(pdf_path: Path, base_dir: Path) -> str:
    """
    Build a custom ID for the request that can be used to recover the file later.

    The ID preserves the full path structure for easy recovery.
    Example: extracted/document_id.pdf becomes "extracted/document_id"

    Args:
        pdf_path: Full path to the PDF file
        base_dir: Base directory containing the processed folder

    Returns:
        Custom ID string that preserves path structure
    """
    # Get relative path from base directory
    rel_path = pdf_path.relative_to(base_dir)
    # Remove .pdf extension but keep directory structure
    path_without_ext = str(rel_path).replace(".pdf", "")
    return path_without_ext


def process_single_pdf(pdf_path: Path, base_dir: Path) -> Optional[Tuple[Dict[str, Any], Path]]:
    """
    Process a single PDF and return the batch request if valid.

    Args:
        pdf_path: Path to the PDF file
        base_dir: Base directory for building custom IDs

    Returns:
        Tuple of (request dict, pdf_path) if successful, None otherwise
    """
    # Validate PDF has single page
    try:
        pdf = PdfReader(pdf_path)
        if len(pdf.pages) != 1:
            return None
    except Exception as e:
        print(f"Error reading PDF {pdf_path}: {e}")
        return None

    try:
        # Render PDF to base64 image
        image_base64 = render_pdf_to_base64png(str(pdf_path), page_num=1, target_longest_image_dim=TARGET_IMAGE_DIM)

        # Get image dimensions for the prompt
        width, height = get_png_dimensions_from_base64(image_base64)

        # Build the prompt using v3 simple version
        prompt = build_openai_silver_data_prompt_v3_simple(width, height)

        # Build custom ID
        custom_id = build_custom_id(pdf_path, base_dir)

        # Build the request in OpenAI batch format
        request = {
            "custom_id": custom_id,
            "method": "POST",
            "url": "/v1/chat/completions",
            "body": {
                "model": "gpt-4.1",
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},
                            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
                        ],
                    }
                ],
                "temperature": 0.1,
                "max_completion_tokens": 12000,
                "response_format": openai_response_format_schema(),
            },
        }

        return (request, pdf_path)
    except Exception as e:
        print(f"Error processing {pdf_path}: {e}")
        return None


def find_pdf_files(input_dir: Path) -> Generator[Path, None, None]:
    """
    Find all PDF files in the processed folder structure.

    The structure is expected to be:
    processed_XX_subset_split/
        extracted/
            *.pdf

    Or for hugging_face downloads:
    hugging_face/
        pdf_tarballs/
            extracted/
                *.pdf

    Args:
        input_dir: Input directory path

    Yields:
        Path objects for each PDF file found
    """

    for pdf_path in input_dir.rglob("*.pdf"):
        yield pdf_path


def process_pdfs_to_batch_requests(input_dir: Path, output_dir: Path, max_pdfs: int = None, num_workers: int = 8) -> int:
    """
    Process PDFs and create batch request files using parallel processing.

    Args:
        input_dir: Directory containing the processed folder structure
        output_dir: Directory to save batch request files
        max_pdfs: Maximum number of PDFs to process (None for all)
        num_workers: Number of parallel workers for processing

    Returns:
        Number of PDFs processed
    """
    # Ensure output directory exists
    output_dir.mkdir(parents=True, exist_ok=True)

    # Initialize file management
    file_num = 0
    current_file_size = 0
    current_file_path = output_dir / f"batch_requests_{file_num:04d}.jsonl"
    current_file = open(current_file_path, "w")

    pdfs_processed = 0
    pdfs_skipped = 0

    # Find PDF files
    pdf_files = list(find_pdf_files(input_dir))

    # Limit files if max_pdfs is specified
    if max_pdfs:
        pdf_files = pdf_files[:max_pdfs]

    total_pdfs = len(pdf_files)

    print(f"Found {total_pdfs} PDF files to process")
    print(f"Using {num_workers} parallel workers")

    # Process PDFs in parallel using ThreadPoolExecutor
    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        # Submit all PDF processing tasks
        future_to_pdf = {executor.submit(process_single_pdf, pdf_path, input_dir): pdf_path for pdf_path in pdf_files}

        # Process results as they complete
        with tqdm(total=total_pdfs, desc="Processing PDFs") as pbar:
            for future in as_completed(future_to_pdf):
                pdf_path = future_to_pdf[future]

                try:
                    result = future.result()

                    if result is None:
                        # PDF was skipped (multi-page or error)
                        pdfs_skipped += 1
                    else:
                        request, _ = result
                        request_json = json.dumps(request)
                        request_size = len(request_json.encode("utf-8"))

                        # Check if we need to start a new file
                        if current_file_size + request_size > MAX_FILE_SIZE:
                            current_file.close()
                            file_num += 1
                            current_file_path = output_dir / f"batch_requests_{file_num:04d}.jsonl"
                            current_file = open(current_file_path, "w")
                            current_file_size = 0
                            print(f"\nStarting new batch file: {current_file_path.name}")

                        # Write the request (only in main thread)
                        current_file.write(request_json)
                        current_file.write("\n")
                        current_file_size += request_size

                        pdfs_processed += 1

                except Exception as e:
                    print(f"\nError with {pdf_path}: {e}")
                    pdfs_skipped += 1

                pbar.update(1)

    # Close the last file
    current_file.close()

    print(f"\nProcessing complete:")
    print(f"  - PDFs processed: {pdfs_processed}")
    print(f"  - PDFs skipped: {pdfs_skipped}")
    print(f"  - Batch files created: {file_num + 1}")
    print(f"  - Output directory: {output_dir}")

    return pdfs_processed


def main():
    parser = argparse.ArgumentParser(description="Build OpenAI batch requests from OLMoCR-mix folder structure")
    parser.add_argument("--output_dir", type=str, default=None, help="Output directory for batch request files (default: input_dir/batch_requests)")
    parser.add_argument("--max_pdfs", type=int, default=None, help="Maximum number of PDFs to process (default: all)")
    parser.add_argument("--num_workers", type=int, default=8, help="Number of parallel workers for processing (default: 8)")
    parser.add_argument(
        "input_dir",
        type=str,
        help="Input directory containing processed folder structure (e.g., ~/olmOCR-mix-0225/processed_00_documents_eval_s2pdf or ~/olmOCR-mix-0225)",
    )

    args = parser.parse_args()

    # Convert paths to Path objects
    input_dir = Path(args.input_dir).expanduser().resolve()

    if not input_dir.exists():
        print(f"Error: Input directory does not exist: {input_dir}")
        return 1

    # Set default output directory if not specified
    if args.output_dir:
        output_dir = Path(args.output_dir).expanduser().resolve()
    else:
        output_dir = input_dir / "batch_requests"

    print(f"Input directory: {input_dir}")
    print(f"Output directory: {output_dir}")

    # Process PDFs
    process_pdfs_to_batch_requests(input_dir=input_dir, output_dir=output_dir, max_pdfs=args.max_pdfs, num_workers=args.num_workers)

    return 0


if __name__ == "__main__":
    exit(main())