uv-scripts/ocr

1# /// script2# requires-python = ">=3.11"3# dependencies = [4#     "datasets>=4.0.0",5#     "huggingface-hub",6#     "pillow",7#     "vllm",8#     "tqdm",9#     "toolz",10#     "torch",11#     "triton-kernels @ git+https://github.com/triton-lang/triton.git@v3.5.0#subdirectory=python/triton_kernels",12# ]13#14# [[tool.uv.index]]15# url = "https://wheels.vllm.ai/nightly/cu129"16#17# [tool.uv]18# prerelease = "allow"19# ///20 21"""22Convert document images to markdown using LightOnOCR with vLLM.23 24LightOnOCR is a compact 1B multilingual OCR model optimized for production speed.25Combines Pixtral ViT encoder with Qwen3 language model for efficient document parsing.26 27NOTE: Requires vLLM nightly wheels for LightOnOCR support. First run may take28a few minutes to download and install dependencies.29 30Features:31- ⚡ Fastest: 5.71 pages/sec on H100 GPU32- 🎯 Compact: Only 1B parameters33- 🌍 Multilingual with European language optimization34- 📐 LaTeX formula recognition35- 📊 Table extraction (markdown format)36- 📝 Document structure preservation37- 🔤 3 vocabulary sizes (151k/32k/16k tokens)38 39Model: lightonai/LightOnOCR-1B-102540vLLM: Requires nightly build from main branch41Performance: 76.1% overall benchmark score42"""43 44import argparse45import base6446import io47import json48import logging49import os50import sys51from typing import Any, Dict, List, Union52from datetime import datetime53 54import torch55from datasets import load_dataset56from huggingface_hub import DatasetCard, login57from PIL import Image58from toolz import partition_all59from tqdm.auto import tqdm60# Disable vLLM's FlashInfer sampler: it JIT-compiles a CUDA kernel needing nvcc, which the61# default uv-script image lacks (engine init then crashes). Greedy OCR doesn't use it; this62# lets the plain default-image command work. On the vllm/vllm-openai image it's a harmless no-op.63os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")64from vllm import LLM, SamplingParams65 66logging.basicConfig(level=logging.INFO)67logger = logging.getLogger(__name__)68 69 70# Model variants with different vocabulary sizes71MODEL_VARIANTS = {72    "151k": "lightonai/LightOnOCR-1B-1025",  # Full vocabulary (default)73    "32k": "lightonai/LightOnOCR-0.9B-32k-1025",  # European languages optimized74    "16k": "lightonai/LightOnOCR-0.9B-16k-1025",  # European languages optimized75}76 77 78def check_cuda_availability():79    """Check if CUDA is available and exit if not."""80    if not torch.cuda.is_available():81        logger.error("CUDA is not available. This script requires a GPU.")82        logger.error("Please run on a machine with a CUDA-capable GPU.")83        sys.exit(1)84    else:85        logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")86 87 88def ensure_output_columns_free(dataset, columns, overwrite=False):89    """Fail fast if an output column would collide with an existing input column.90 91    Adding a column that already exists silently overwrites it (e.g. a ground-truth92    `text`/`markdown` column) or crashes on push with a duplicate-column error only93    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing94    column(s) here instead (logged) so the later add_column is clean.95    """96    clash = [c for c in columns if c in dataset.column_names]97    if not clash:98        return dataset99    if overwrite:100        logger.warning(f"--overwrite: replacing existing column(s) {clash}")101        return dataset.remove_columns(clash)102    logger.error(103        f"Output column(s) {clash} already exist in the input dataset "104        f"(columns: {dataset.column_names})."105    )106    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")107    sys.exit(1)108 109 110def resize_image_to_target(image: Image.Image, target_size: int = 1540) -> Image.Image:111    """112    Resize image so longest dimension is target_size while maintaining aspect ratio.113 114    LightOnOCR was trained with images at 1540px max resolution and 200 DPI.115    """116    width, height = image.size117 118    # If image is already smaller, don't upscale119    if max(width, height) <= target_size:120        return image121 122    # Calculate new dimensions maintaining aspect ratio123    if width > height:124        new_width = target_size125        new_height = int(height * (target_size / width))126    else:127        new_height = target_size128        new_width = int(width * (target_size / height))129 130    return image.resize((new_width, new_height), Image.Resampling.LANCZOS)131 132 133def make_ocr_message(134    image: Union[Image.Image, Dict[str, Any], str],135    resize: bool = True,136    target_size: int = 1540,137) -> List[Dict]:138    """139    Create chat message for OCR processing.140 141    LightOnOCR was trained with 1540px max resolution at 200 DPI for optimal results.142    """143    # Convert to PIL Image if needed144    if isinstance(image, Image.Image):145        pil_img = image146    elif isinstance(image, dict) and "bytes" in image:147        pil_img = Image.open(io.BytesIO(image["bytes"]))148    elif isinstance(image, str):149        pil_img = Image.open(image)150    else:151        raise ValueError(f"Unsupported image type: {type(image)}")152 153    # Convert to RGB154    pil_img = pil_img.convert("RGB")155 156    # Resize to optimal dimensions for LightOnOCR157    if resize:158        pil_img = resize_image_to_target(pil_img, target_size)159        logger.debug(f"Resized image to {pil_img.size}")160 161    # Convert to base64 data URI162    buf = io.BytesIO()163    pil_img.save(buf, format="PNG")164    data_uri = f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode()}"165 166    # LightOnOCR uses message format with empty text prompt before image167    # (matching official demo: text first, then image)168    return [169        {170            "role": "user",171            "content": [172                {"type": "text", "text": ""},173                {"type": "image_url", "image_url": {"url": data_uri}},174            ],175        }176    ]177 178 179def create_dataset_card(180    source_dataset: str,181    model: str,182    vocab_size: str,183    num_samples: int,184    processing_time: str,185    batch_size: int,186    max_model_len: int,187    max_tokens: int,188    gpu_memory_utilization: float,189    temperature: float,190    top_p: float,191    target_size: int,192    image_column: str = "image",193    split: str = "train",194) -> str:195    """Create a dataset card documenting the OCR process."""196    model_name = model.split("/")[-1]197 198    return f"""---199tags:200- ocr201- document-processing202- lighton-ocr203- markdown204- uv-script205- generated206---207 208# Document OCR using {model_name}209 210This dataset contains OCR results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using LightOnOCR, a fast and compact 1B OCR model.211 212## Processing Details213 214- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})215- **Model**: [{model}](https://huggingface.co/{model})216- **Vocabulary Size**: {vocab_size} tokens217- **Number of Samples**: {num_samples:,}218- **Processing Time**: {processing_time}219- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}220 221### Configuration222 223- **Image Column**: `{image_column}`224- **Output Column**: `markdown`225- **Dataset Split**: `{split}`226- **Batch Size**: {batch_size}227- **Target Image Size**: {target_size}px (longest dimension)228- **Max Model Length**: {max_model_len:,} tokens229- **Max Output Tokens**: {max_tokens:,}230- **Temperature**: {temperature}231- **Top P**: {top_p}232- **GPU Memory Utilization**: {gpu_memory_utilization:.1%}233 234## Model Information235 236LightOnOCR is a fast, compact OCR model that excels at:237- ⚡ **Production Speed** - 5.71 pages/second on H100 GPU238- 🎯 **Compact Size** - Only 1B parameters239- 📐 **LaTeX formulas** - Mathematical notation in LaTeX format240- 📊 **Tables** - Extracted and formatted as markdown241- 📝 **Document structure** - Hierarchy and layout preservation242- 🌍 **Multilingual** - Optimized for European languages243- 🔤 **Flexible vocabulary** - 151k/32k/16k token variants244 245### Vocabulary Variants246 247- **151k tokens**: Full vocabulary, supports all languages248- **32k tokens**: European languages optimized (~12% faster decoding)249- **16k tokens**: European languages optimized (~12% faster decoding)250 251## Dataset Structure252 253The dataset contains all original columns plus:254- `markdown`: The extracted text in markdown format with LaTeX formulas255- `inference_info`: JSON list tracking all OCR models applied to this dataset256 257## Usage258 259```python260from datasets import load_dataset261import json262 263# Load the dataset264dataset = load_dataset("{{output_dataset_id}}", split="{split}")265 266# Access the markdown text267for example in dataset:268    print(example["markdown"])269    break270 271# View all OCR models applied to this dataset272inference_info = json.loads(dataset[0]["inference_info"])273for info in inference_info:274    print(f"Column: {{info['column_name']}} - Model: {{info['model_id']}}")275```276 277## Reproduction278 279This dataset was generated using the [uv-scripts/ocr](https://huggingface.co/datasets/uv-scripts/ocr) LightOnOCR script:280 281```bash282uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr.py \\283    {source_dataset} \\284    <output-dataset> \\285    --vocab-size {vocab_size} \\286    --image-column {image_column} \\287    --batch-size {batch_size}288```289 290## Performance291 292- **Processing Speed**: ~{num_samples / (float(processing_time.split()[0]) * 60):.2f} images/second293- **Benchmark Score**: 76.1% overall (across diverse document types)294- **Optimization**: Native resolution ViT + lightweight decoder295 296Generated with 🤖 [UV Scripts](https://huggingface.co/uv-scripts)297"""298 299 300def main(301    input_dataset: str,302    output_dataset: str,303    image_column: str = "image",304    batch_size: int = 16,305    vocab_size: str = "151k",306    max_model_len: int = 8192,307    max_tokens: int = 6500,308    temperature: float = 0.2,309    top_p: float = 0.9,310    gpu_memory_utilization: float = 0.8,311    target_size: int = 1540,312    no_resize: bool = False,313    hf_token: str = None,314    split: str = "train",315    max_samples: int = None,316    private: bool = False,317    shuffle: bool = False,318    seed: int = 42,319    output_column: str = "markdown",320    overwrite: bool = False,321):322    """Process images from HF dataset through LightOnOCR model."""323 324    # Check CUDA availability first325    check_cuda_availability()326 327    # Track processing start time328    start_time = datetime.now()329 330    # Login to HF if token provided331    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")332    if HF_TOKEN:333        login(token=HF_TOKEN)334 335    # Get model ID from vocabulary size336    if vocab_size not in MODEL_VARIANTS:337        raise ValueError(338            f"Invalid vocab_size '{vocab_size}'. Choose from: {list(MODEL_VARIANTS.keys())}"339        )340    model = MODEL_VARIANTS[vocab_size]341    logger.info(f"Using model: {model} ({vocab_size} vocabulary)")342 343    # Load dataset344    logger.info(f"Loading dataset: {input_dataset}")345    dataset = load_dataset(input_dataset, split=split)346 347    # Validate image column348    if image_column not in dataset.column_names:349        raise ValueError(350            f"Column '{image_column}' not found. Available: {dataset.column_names}"351        )352 353    # Fail fast if the output column would collide with an existing input column354    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)355 356    # Shuffle if requested357    if shuffle:358        logger.info(f"Shuffling dataset with seed {seed}")359        dataset = dataset.shuffle(seed=seed)360 361    # Limit samples if requested362    if max_samples:363        dataset = dataset.select(range(min(max_samples, len(dataset))))364        logger.info(f"Limited to {len(dataset)} samples")365 366    # Initialize vLLM model367    logger.info("Initializing vLLM with LightOnOCR")368    logger.info("This may take a few minutes on first run...")369    llm = LLM(370        model=model,371        trust_remote_code=True,372        max_model_len=max_model_len,373        gpu_memory_utilization=gpu_memory_utilization,374        limit_mm_per_prompt={"image": 1},  # One image per prompt375        enforce_eager=False,  # Use torch.compile for better performance376    )377 378    # LightOnOCR recommended sampling parameters379    sampling_params = SamplingParams(380        temperature=temperature,381        top_p=top_p,382        max_tokens=max_tokens,383    )384 385    logger.info(f"Processing {len(dataset)} images in batches of {batch_size}")386    logger.info(f"Output will be written to column: {output_column}")387    if not no_resize:388        logger.info(f"Images will be resized to {target_size}px (longest dimension)")389 390    # Process images in batches391    all_outputs = []392 393    for batch_indices in tqdm(394        partition_all(batch_size, range(len(dataset))),395        total=(len(dataset) + batch_size - 1) // batch_size,396        desc="LightOnOCR processing",397    ):398        batch_indices = list(batch_indices)399        batch_images = [dataset[i][image_column] for i in batch_indices]400 401        try:402            # Create messages for batch403            batch_messages = [404                make_ocr_message(img, resize=not no_resize, target_size=target_size)405                for img in batch_images406            ]407 408            # Process with vLLM409            outputs = llm.chat(batch_messages, sampling_params)410 411            # Extract outputs412            for output in outputs:413                text = output.outputs[0].text.strip()414                all_outputs.append(text)415 416        except Exception as e:417            logger.error(f"Error processing batch: {e}")418            # Add error placeholders for failed batch419            all_outputs.extend(["[OCR ERROR]"] * len(batch_images))420 421    # Calculate processing time422    processing_duration = datetime.now() - start_time423    processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"424 425    # Add output column to dataset426    logger.info(f"Adding '{output_column}' column to dataset")427    dataset = dataset.add_column(output_column, all_outputs)428 429    # Handle inference_info tracking (for multi-model comparisons)430    inference_entry = {431        "model_id": model,432        "model_name": "LightOnOCR",433        "vocab_size": vocab_size,434        "column_name": output_column,435        "timestamp": datetime.now().isoformat(),436        "temperature": temperature,437        "top_p": top_p,438        "max_tokens": max_tokens,439        "target_size": target_size if not no_resize else "original",440    }441 442    if "inference_info" in dataset.column_names:443        # Append to existing inference info444        logger.info("Updating existing inference_info column")445 446        def update_inference_info(example):447            try:448                existing_info = (449                    json.loads(example["inference_info"])450                    if example["inference_info"]451                    else []452                )453            except (json.JSONDecodeError, TypeError):454                existing_info = []455 456            existing_info.append(inference_entry)457            return {"inference_info": json.dumps(existing_info)}458 459        dataset = dataset.map(update_inference_info)460    else:461        # Create new inference_info column462        logger.info("Creating new inference_info column")463        inference_list = [json.dumps([inference_entry])] * len(dataset)464        dataset = dataset.add_column("inference_info", inference_list)465 466    # Push to hub467    logger.info(f"Pushing to {output_dataset}")468    dataset.push_to_hub(output_dataset, private=private, token=HF_TOKEN)469 470    # Create and push dataset card471    logger.info("Creating dataset card")472    card_content = create_dataset_card(473        source_dataset=input_dataset,474        model=model,475        vocab_size=vocab_size,476        num_samples=len(dataset),477        processing_time=processing_time_str,478        batch_size=batch_size,479        max_model_len=max_model_len,480        max_tokens=max_tokens,481        gpu_memory_utilization=gpu_memory_utilization,482        temperature=temperature,483        top_p=top_p,484        target_size=target_size,485        image_column=image_column,486        split=split,487    )488 489    card = DatasetCard(card_content)490    card.push_to_hub(output_dataset, token=HF_TOKEN)491 492    logger.info("✅ LightOnOCR processing complete!")493    logger.info(494        f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"495    )496    logger.info(f"Processing time: {processing_time_str}")497    logger.info(498        f"Processing speed: {len(dataset) / processing_duration.total_seconds():.2f} images/sec"499    )500 501 502if __name__ == "__main__":503    # Show example usage if no arguments504    if len(sys.argv) == 1:505        print("=" * 80)506        print("LightOnOCR Document Processing")507        print("=" * 80)508        print("\nFast, compact 1B OCR model for production workloads")509        print("\nFeatures:")510        print("- ⚡ Fastest processing: 5.71 pages/sec on H100")511        print("- 🎯 Compact: Only 1B parameters")512        print("- 🌍 Multilingual with European language optimization")513        print("- 📐 LaTeX formula recognition")514        print("- 📊 Table extraction (markdown format)")515        print("- 🔤 3 vocabulary sizes for speed/quality tradeoffs")516        print("\nExample usage:")517        print("\n1. Basic OCR (full vocabulary):")518        print("   uv run lighton-ocr.py input-dataset output-dataset")519        print("\n2. European languages optimized (faster):")520        print("   uv run lighton-ocr.py docs results --vocab-size 32k")521        print("\n3. Custom batch size for performance:")522        print("   uv run lighton-ocr.py docs results --batch-size 32")523        print("\n4. Test with small sample:")524        print("   uv run lighton-ocr.py large-dataset test --max-samples 50 --shuffle")525        print("\n5. Original image size (no resize):")526        print("   uv run lighton-ocr.py docs output --no-resize")527        print("\n6. Running on HF Jobs:")528        print("   hf jobs uv run --flavor l4x1 \\")529        print(530            '     -e HF_TOKEN=$(python3 -c "from huggingface_hub import get_token; print(get_token())") \\'531        )532        print(533            "     https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr.py \\"534        )535        print("       input-dataset output-dataset --vocab-size 32k")536        print("\n" + "=" * 80)537        print("\nVocabulary Size Options:")538        print("  151k - Full vocabulary (all languages)")539        print("  32k  - European languages (~12% faster)")540        print("  16k  - European languages (~12% faster)")541        print("\nFor full help, run: uv run lighton-ocr.py --help")542        sys.exit(0)543 544    parser = argparse.ArgumentParser(545        description="Document OCR using LightOnOCR (fast 1B model)",546        formatter_class=argparse.RawDescriptionHelpFormatter,547        epilog="""548Vocabulary Size Options:549  151k    Full vocabulary supporting all languages (default)550  32k     European languages optimized (~12% faster decoding)551  16k     European languages optimized (~12% faster decoding)552 553Examples:554  # Basic text OCR with full vocabulary555  uv run lighton-ocr.py my-docs analyzed-docs556 557  # Fast processing for European languages558  uv run lighton-ocr.py papers results --vocab-size 32k559 560  # Test with random sampling561  uv run lighton-ocr.py large-dataset test --max-samples 50 --shuffle562 563  # Custom batch size for GPU optimization564  uv run lighton-ocr.py dataset output --batch-size 32 --gpu-memory-utilization 0.9565        """,566    )567 568    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")569    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")570    parser.add_argument(571        "--image-column",572        default="image",573        help="Column containing images (default: image)",574    )575    parser.add_argument(576        "--batch-size",577        type=int,578        default=16,579        help="Batch size for processing (default: 16)",580    )581    parser.add_argument(582        "--vocab-size",583        default="151k",584        choices=list(MODEL_VARIANTS.keys()),585        help="Vocabulary size variant (default: 151k)",586    )587    parser.add_argument(588        "--max-model-len",589        type=int,590        default=8192,591        help="Maximum model context length (default: 8192)",592    )593    parser.add_argument(594        "--max-tokens",595        type=int,596        default=6500,597        help="Maximum tokens to generate (default: 6500)",598    )599    parser.add_argument(600        "--temperature",601        type=float,602        default=0.2,603        help="Sampling temperature (default: 0.2)",604    )605    parser.add_argument(606        "--top-p",607        type=float,608        default=0.9,609        help="Top-p sampling parameter (default: 0.9)",610    )611    parser.add_argument(612        "--gpu-memory-utilization",613        type=float,614        default=0.8,615        help="GPU memory utilization (default: 0.8)",616    )617    parser.add_argument(618        "--target-size",619        type=int,620        default=1540,621        help="Target size for longest image dimension in pixels (default: 1540, matching training)",622    )623    parser.add_argument(624        "--no-resize",625        action="store_true",626        help="Don't resize images (use original size)",627    )628    parser.add_argument("--hf-token", help="Hugging Face API token")629    parser.add_argument(630        "--split", default="train", help="Dataset split to use (default: train)"631    )632    parser.add_argument(633        "--max-samples",634        type=int,635        help="Maximum number of samples to process (for testing)",636    )637    parser.add_argument(638        "--private", action="store_true", help="Make output dataset private"639    )640    parser.add_argument(641        "--shuffle", action="store_true", help="Shuffle dataset before processing"642    )643    parser.add_argument(644        "--seed",645        type=int,646        default=42,647        help="Random seed for shuffling (default: 42)",648    )649    parser.add_argument(650        "--output-column",651        default="markdown",652        help="Column name for output text (default: markdown)",653    )654    parser.add_argument(655        "--overwrite",656        action="store_true",657        help="Replace the output column if it already exists in the input dataset "658        "(default: error out to avoid clobbering an existing column).",659    )660 661    args = parser.parse_args()662 663    main(664        input_dataset=args.input_dataset,665        output_dataset=args.output_dataset,666        image_column=args.image_column,667        batch_size=args.batch_size,668        vocab_size=args.vocab_size,669        max_model_len=args.max_model_len,670        max_tokens=args.max_tokens,671        temperature=args.temperature,672        top_p=args.top_p,673        gpu_memory_utilization=args.gpu_memory_utilization,674        target_size=args.target_size,675        no_resize=args.no_resize,676        hf_token=args.hf_token,677        split=args.split,678        max_samples=args.max_samples,679        private=args.private,680        shuffle=args.shuffle,681        seed=args.seed,682        output_column=args.output_column,683        overwrite=args.overwrite,684    )685