uv-scripts/ocr

1# /// script2# requires-python = ">=3.11"3# dependencies = [4#     "datasets>=4.0.0",5#     "huggingface-hub",6#     "pillow",7#     "vllm>=0.15.1",8#     "tqdm",9#     "toolz",10#     "torch",11# ]12#13# ///14 15"""16Convert document images to markdown using FireRed-OCR with vLLM.17 18FireRed-OCR is a 2.1B document OCR model fine-tuned from Qwen3-VL-2B-Instruct.19It converts PDF/document images to structured Markdown with LaTeX formulas and20HTML tables. Apache 2.0 licensed.21 22Model: FireRedTeam/FireRed-OCR23vLLM: Uses stable Qwen3-VL support (>=0.15.1)24"""25 26import argparse27import base6428import io29import json30import logging31import os32import sys33from datetime import datetime34from typing import Any, Dict, List, Union35 36import torch37from datasets import load_dataset38from huggingface_hub import DatasetCard, login39from PIL import Image40from toolz import partition_all41from tqdm.auto import tqdm42# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the43# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this44# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.45os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")46from vllm import LLM, SamplingParams47 48logging.basicConfig(level=logging.INFO)49logger = logging.getLogger(__name__)50 51 52# ────────────────────────────────────────────────────────────────53# FireRed-OCR Prompt (from official conv_for_infer.py)54# Source: https://github.com/FireRedTeam/FireRed-OCR/blob/main/conv_for_infer.py55# ────────────────────────────────────────────────────────────────56 57FIRERED_OCR_PROMPT = """You are an AI assistant specialized in converting PDF images to Markdown format. Please follow these instructions for the conversion:58 591. Text Processing:60- Accurately recognize all text content in the PDF image without guessing or inferring.61- Convert the recognized text into Markdown format.62- Maintain the original document structure, including headings, paragraphs, lists, etc.63 642. Mathematical Formula Processing:65- Convert all mathematical formulas to LaTeX format.66- Enclose inline formulas with \\( \\). For example: This is an inline formula \\( E = mc^2 \\)67- Enclose block formulas with \\[ \\]. For example: \\[ \\frac{-b \\pm \\sqrt{b^2 - 4ac}}{2a} \\]68 693. Table Processing:70- Convert tables to HTML format.71- Wrap the entire table with <table> and </table>.72 734. Figure Handling:74- Ignore figures content in the PDF image. Do not attempt to describe or convert images.75 765. Output Format:77- Ensure the output Markdown document has a clear structure with appropriate line breaks between elements.78- For complex layouts, try to maintain the original document's structure and format as closely as possible.79 80Please strictly follow these guidelines to ensure accuracy and consistency in the conversion. Your task is to accurately convert the content of the PDF image into Markdown format without adding any extra explanations or comments."""81 82 83def check_cuda_availability():84    """Check if CUDA is available and exit if not."""85    if not torch.cuda.is_available():86        logger.error("CUDA is not available. This script requires a GPU.")87        logger.error("Please run on a machine with a CUDA-capable GPU.")88        sys.exit(1)89    else:90        logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")91 92 93def ensure_output_columns_free(dataset, columns, overwrite=False):94    """Fail fast if an output column would collide with an existing input column.95 96    Adding a column that already exists silently overwrites it (e.g. a ground-truth97    `text`/`markdown` column) or crashes on push with a duplicate-column error only98    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing99    column(s) here instead (logged) so the later add_column is clean.100    """101    clash = [c for c in columns if c in dataset.column_names]102    if not clash:103        return dataset104    if overwrite:105        logger.warning(f"--overwrite: replacing existing column(s) {clash}")106        return dataset.remove_columns(clash)107    logger.error(108        f"Output column(s) {clash} already exist in the input dataset "109        f"(columns: {dataset.column_names})."110    )111    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")112    sys.exit(1)113 114 115def make_ocr_message(116    image: Union[Image.Image, Dict[str, Any], str],117    prompt: str = FIRERED_OCR_PROMPT,118) -> List[Dict]:119    """Create chat message for OCR processing."""120    # Convert to PIL Image if needed121    if isinstance(image, Image.Image):122        pil_img = image123    elif isinstance(image, dict) and "bytes" in image:124        pil_img = Image.open(io.BytesIO(image["bytes"]))125    elif isinstance(image, str):126        pil_img = Image.open(image)127    else:128        raise ValueError(f"Unsupported image type: {type(image)}")129 130    # Convert to RGB131    pil_img = pil_img.convert("RGB")132 133    # Convert to base64 data URI (JPEG is faster than PNG for encoding)134    buf = io.BytesIO()135    pil_img.save(buf, format="JPEG", quality=95)136    data_uri = f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"137 138    # Return message in vLLM format139    return [140        {141            "role": "user",142            "content": [143                {"type": "image_url", "image_url": {"url": data_uri}},144                {"type": "text", "text": prompt},145            ],146        }147    ]148 149 150def create_dataset_card(151    source_dataset: str,152    model: str,153    num_samples: int,154    processing_time: str,155    batch_size: int,156    max_model_len: int,157    max_tokens: int,158    gpu_memory_utilization: float,159    image_column: str = "image",160    split: str = "train",161) -> str:162    """Create a dataset card documenting the OCR process."""163    model_name = model.split("/")[-1]164 165    return f"""---166tags:167- ocr168- document-processing169- firered-ocr170- markdown171- uv-script172- generated173---174 175# Document OCR using {model_name}176 177This dataset contains OCR results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using FireRed-OCR, a 2.1B model fine-tuned from Qwen3-VL-2B-Instruct.178 179## Processing Details180 181- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})182- **Model**: [{model}](https://huggingface.co/{model})183- **Number of Samples**: {num_samples:,}184- **Processing Time**: {processing_time}185- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}186 187### Configuration188 189- **Image Column**: `{image_column}`190- **Output Column**: `markdown`191- **Dataset Split**: `{split}`192- **Batch Size**: {batch_size}193- **Max Model Length**: {max_model_len:,} tokens194- **Max Output Tokens**: {max_tokens:,}195- **GPU Memory Utilization**: {gpu_memory_utilization:.1%}196 197## Model Information198 199FireRed-OCR is a document OCR model that converts images to structured Markdown:200- Fine-tuned from Qwen3-VL-2B-Instruct (2.1B parameters)201- LaTeX formula support (inline and block)202- HTML table extraction203- Layout-aware text extraction204- Apache 2.0 licensed205 206## Dataset Structure207 208The dataset contains all original columns plus:209- `markdown`: The extracted text in markdown format210- `inference_info`: JSON list tracking all OCR models applied to this dataset211 212## Usage213 214```python215from datasets import load_dataset216import json217 218# Load the dataset219dataset = load_dataset("{{output_dataset_id}}", split="{split}")220 221# Access the markdown text222for example in dataset:223    print(example["markdown"])224    break225 226# View all OCR models applied to this dataset227inference_info = json.loads(dataset[0]["inference_info"])228for info in inference_info:229    print(f"Column: {{info['column_name']}} - Model: {{info['model_id']}}")230```231 232## Reproduction233 234This dataset was generated using the [uv-scripts/ocr](https://huggingface.co/datasets/uv-scripts/ocr) FireRed-OCR script:235 236```bash237uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/firered-ocr.py \\238    {source_dataset} \\239    <output-dataset> \\240    --image-column {image_column} \\241    --batch-size {batch_size} \\242    --max-model-len {max_model_len} \\243    --max-tokens {max_tokens} \\244    --gpu-memory-utilization {gpu_memory_utilization}245```246 247Generated with [UV Scripts](https://huggingface.co/uv-scripts)248"""249 250 251def main(252    input_dataset: str,253    output_dataset: str,254    image_column: str = "image",255    batch_size: int = 16,256    model: str = "FireRedTeam/FireRed-OCR",257    max_model_len: int = 32768,258    max_tokens: int = 8192,259    gpu_memory_utilization: float = 0.8,260    hf_token: str = None,261    split: str = "train",262    max_samples: int = None,263    private: bool = False,264    shuffle: bool = False,265    seed: int = 42,266    output_column: str = "markdown",267    overwrite: bool = False,268    config: str = None,269    create_pr: bool = False,270):271    """Process images from HF dataset through FireRed-OCR model."""272 273    # Check CUDA availability first274    check_cuda_availability()275 276    # Track processing start time277    start_time = datetime.now()278 279    # Login to HF if token provided280    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")281    if HF_TOKEN:282        login(token=HF_TOKEN)283 284    # Load dataset285    logger.info(f"Loading dataset: {input_dataset}")286    dataset = load_dataset(input_dataset, split=split)287 288    # Validate image column289    if image_column not in dataset.column_names:290        raise ValueError(291            f"Column '{image_column}' not found. Available: {dataset.column_names}"292        )293 294    # Fail fast if the output column would collide with an existing input column295    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)296 297    # Shuffle if requested298    if shuffle:299        logger.info(f"Shuffling dataset with seed {seed}")300        dataset = dataset.shuffle(seed=seed)301 302    # Limit samples if requested303    if max_samples:304        dataset = dataset.select(range(min(max_samples, len(dataset))))305        logger.info(f"Limited to {len(dataset)} samples")306 307    # Initialize vLLM model308    logger.info(f"Initializing vLLM with model: {model}")309    logger.info("This may take a few minutes on first run...")310    llm = LLM(311        model=model,312        trust_remote_code=True,313        max_model_len=max_model_len,314        gpu_memory_utilization=gpu_memory_utilization,315        limit_mm_per_prompt={"image": 1},316    )317 318    sampling_params = SamplingParams(319        temperature=0.0,  # Deterministic for OCR320        max_tokens=max_tokens,321    )322 323    logger.info(f"Processing {len(dataset)} images in batches of {batch_size}")324    logger.info(f"Output will be written to column: {output_column}")325 326    # Process images in batches327    all_outputs = []328 329    for batch_indices in tqdm(330        partition_all(batch_size, range(len(dataset))),331        total=(len(dataset) + batch_size - 1) // batch_size,332        desc="FireRed-OCR processing",333    ):334        batch_indices = list(batch_indices)335        batch_images = [dataset[i][image_column] for i in batch_indices]336 337        try:338            # Create messages for batch339            batch_messages = [340                make_ocr_message(img, FIRERED_OCR_PROMPT) for img in batch_images341            ]342 343            # Process with vLLM344            outputs = llm.chat(batch_messages, sampling_params)345 346            # Extract outputs347            for output in outputs:348                text = output.outputs[0].text.strip()349                all_outputs.append(text)350 351        except Exception as e:352            logger.warning(353                f"Batch failed ({len(batch_images)} images), retrying individually: {e}"354            )355            for img in batch_images:356                try:357                    msg = make_ocr_message(img, FIRERED_OCR_PROMPT)358                    out = llm.chat([msg], sampling_params)359                    all_outputs.append(out[0].outputs[0].text.strip())360                except Exception as img_e:361                    logger.error(f"Image failed: {img_e}")362                    all_outputs.append("[OCR ERROR]")363 364    # Calculate processing time365    processing_duration = datetime.now() - start_time366    processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"367 368    # Add output column to dataset (remove existing column if present)369    if output_column in dataset.column_names:370        logger.info(f"Removing existing '{output_column}' column before adding new results")371        dataset = dataset.remove_columns([output_column])372    logger.info(f"Adding '{output_column}' column to dataset")373    dataset = dataset.add_column(output_column, all_outputs)374 375    # Handle inference_info tracking (for multi-model comparisons)376    inference_entry = {377        "model_id": model,378        "column_name": output_column,379        "timestamp": datetime.now().isoformat(),380    }381 382    if "inference_info" in dataset.column_names:383        # Append to existing inference info384        logger.info("Updating existing inference_info column")385 386        def update_inference_info(example):387            try:388                existing_info = (389                    json.loads(example["inference_info"])390                    if example["inference_info"]391                    else []392                )393            except (json.JSONDecodeError, TypeError):394                existing_info = []395 396            existing_info.append(inference_entry)397            return {"inference_info": json.dumps(existing_info)}398 399        dataset = dataset.map(update_inference_info)400    else:401        # Create new inference_info column402        logger.info("Creating new inference_info column")403        inference_list = [json.dumps([inference_entry])] * len(dataset)404        dataset = dataset.add_column("inference_info", inference_list)405 406    # Push to hub407    logger.info(f"Pushing to {output_dataset}")408    dataset.push_to_hub(409        output_dataset,410        private=private,411        token=HF_TOKEN,412        **({"config_name": config} if config else {}),413        create_pr=create_pr,414        commit_message=f"Add {model} OCR results ({len(dataset)} samples)"415        + (f" [{config}]" if config else ""),416    )417 418    # Create and push dataset card (skip for PR-based benchmark runs)419    if not create_pr:420        logger.info("Creating dataset card")421        card_content = create_dataset_card(422            source_dataset=input_dataset,423            model=model,424            num_samples=len(dataset),425            processing_time=processing_time_str,426            batch_size=batch_size,427            max_model_len=max_model_len,428            max_tokens=max_tokens,429            gpu_memory_utilization=gpu_memory_utilization,430            image_column=image_column,431            split=split,432        )433 434        card = DatasetCard(card_content)435        card.push_to_hub(output_dataset, token=HF_TOKEN)436 437    logger.info("FireRed-OCR processing complete!")438    logger.info(439        f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"440    )441    logger.info(f"Processing time: {processing_time_str}")442 443 444if __name__ == "__main__":445    # Show example usage if no arguments446    if len(sys.argv) == 1:447        print("=" * 80)448        print("FireRed-OCR Document Processing")449        print("=" * 80)450        print("\n2.1B document OCR model (Qwen3-VL-2B fine-tune, Apache 2.0)")451        print("\nFeatures:")452        print("- Structured Markdown output")453        print("- LaTeX formula support (inline and block)")454        print("- HTML table extraction")455        print("- Layout-aware text extraction")456        print("\nExample usage:")457        print("\n1. Basic OCR:")458        print("   uv run firered-ocr.py input-dataset output-dataset")459        print("\n2. With custom settings:")460        print(461            "   uv run firered-ocr.py docs analyzed-docs --batch-size 20 --max-samples 100"462        )463        print("\n3. Running on HF Jobs:")464        print("   hf jobs uv run --flavor l4x1 \\")465        print("     -s HF_TOKEN \\")466        print(467            "     https://huggingface.co/datasets/uv-scripts/ocr/raw/main/firered-ocr.py \\"468        )469        print("       input-dataset output-dataset")470        print("\n" + "=" * 80)471        print("\nFor full help, run: uv run firered-ocr.py --help")472        sys.exit(0)473 474    parser = argparse.ArgumentParser(475        description="Document OCR using FireRed-OCR (2.1B, Qwen3-VL fine-tune)",476        formatter_class=argparse.RawDescriptionHelpFormatter,477        epilog="""478Examples:479  # Basic text OCR480  uv run firered-ocr.py my-docs analyzed-docs481 482  # Random sampling for testing483  uv run firered-ocr.py large-dataset test --max-samples 50 --shuffle484 485  # Benchmark mode (push as config with PR)486  uv run firered-ocr.py source-data bench-repo --config firered-ocr --create-pr487 488  # HF Jobs489  hf jobs uv run --flavor l4x1 -s HF_TOKEN \\490      https://huggingface.co/datasets/uv-scripts/ocr/raw/main/firered-ocr.py \\491      input-dataset output-dataset --max-samples 50492        """,493    )494 495    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")496    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")497    parser.add_argument(498        "--image-column",499        default="image",500        help="Column containing images (default: image)",501    )502    parser.add_argument(503        "--batch-size",504        type=int,505        default=16,506        help="Batch size for processing (default: 16)",507    )508    parser.add_argument(509        "--model",510        default="FireRedTeam/FireRed-OCR",511        help="Model to use (default: FireRedTeam/FireRed-OCR)",512    )513    parser.add_argument(514        "--max-model-len",515        type=int,516        default=32768,517        help="Maximum model context length (default: 32768)",518    )519    parser.add_argument(520        "--max-tokens",521        type=int,522        default=8192,523        help="Maximum tokens to generate (default: 8192)",524    )525    parser.add_argument(526        "--gpu-memory-utilization",527        type=float,528        default=0.8,529        help="GPU memory utilization (default: 0.8)",530    )531    parser.add_argument("--hf-token", help="Hugging Face API token")532    parser.add_argument(533        "--split", default="train", help="Dataset split to use (default: train)"534    )535    parser.add_argument(536        "--max-samples",537        type=int,538        help="Maximum number of samples to process (for testing)",539    )540    parser.add_argument(541        "--private", action="store_true", help="Make output dataset private"542    )543    parser.add_argument(544        "--shuffle", action="store_true", help="Shuffle dataset before processing"545    )546    parser.add_argument(547        "--seed",548        type=int,549        default=42,550        help="Random seed for shuffling (default: 42)",551    )552    parser.add_argument(553        "--output-column",554        default="markdown",555        help="Column name for output text (default: markdown)",556    )557    parser.add_argument(558        "--overwrite",559        action="store_true",560        help="Replace the output column if it already exists in the input dataset "561        "(default: error out to avoid clobbering an existing column).",562    )563    parser.add_argument(564        "--config",565        help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",566    )567    parser.add_argument(568        "--create-pr",569        action="store_true",570        help="Create a pull request instead of pushing directly (for parallel benchmarking)",571    )572 573    args = parser.parse_args()574 575    main(576        input_dataset=args.input_dataset,577        output_dataset=args.output_dataset,578        image_column=args.image_column,579        batch_size=args.batch_size,580        model=args.model,581        max_model_len=args.max_model_len,582        max_tokens=args.max_tokens,583        gpu_memory_utilization=args.gpu_memory_utilization,584        hf_token=args.hf_token,585        split=args.split,586        max_samples=args.max_samples,587        private=args.private,588        shuffle=args.shuffle,589        seed=args.seed,590        output_column=args.output_column,591        overwrite=args.overwrite,592        config=args.config,593        create_pr=args.create_pr,594    )595