uv-scripts/ocr

1# /// script2# requires-python = ">=3.11"3# dependencies = [4#     "datasets>=4.0.0",5#     "huggingface-hub",6#     "pillow",7#     "tqdm",8#     "toolz",9# ]10#11# ///12 13"""14Convert document images to markdown using Nanonets-OCR2-3B with vLLM.15 16This script processes images through the Nanonets-OCR2-3B model (3.75B params)17to extract text and structure as markdown, ideal for document understanding tasks.18 19Features:20- LaTeX equation recognition21- Table extraction and formatting (HTML)22- Document structure preservation23- Image descriptions and captions24- Signature and watermark detection25- Checkbox recognition26- Multilingual support27 28⚠️ Pinned vLLM image required. Nanonets-OCR2-3B is a Qwen2.5-VL model, and vLLM29>=0.11 regressed Qwen2.5-VL decoding to a stream of "!" (vllm-project/vllm#27775,30#14126); 0.10.2 is the last known-good. So `vllm` and `torch` are intentionally31omitted from the deps above and come from the `vllm/vllm-openai:v0.10.2` image at32runtime via PYTHONPATH. Run it like:33 34    hf jobs uv run --flavor a10g-small -s HF_TOKEN \35        --image vllm/vllm-openai:v0.10.2 --python /usr/bin/python3 \36        -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \37        nanonets-ocr2.py INPUT_DATASET OUTPUT_DATASET --max-samples 10 --shuffle38 39Re-test the pin when a newer vLLM ships a Qwen2.5-VL decode fix; it can then move40back to a plain `vllm>=...` dep on the default image.41"""42 43import argparse44import base6445import io46import json47import logging48import os49import sys50from typing import Any, Dict, List, Union51from datetime import datetime52 53# Preflight: torch/vllm are deliberately NOT PEP 723 deps (see docstring) — they come from54# the pinned vllm/vllm-openai:v0.10.2 image. On the bare uv image the imports below would55# die with a raw ModuleNotFoundError; fail fast naming the exact flags instead.56import importlib.util57 58if importlib.util.find_spec("vllm") is None or importlib.util.find_spec("torch") is None:59    sys.stderr.write(60        "ERROR: vllm/torch not importable — this recipe cannot run on the bare uv image.\n"61        "Nanonets-OCR2-3B needs vLLM 0.10.2 (>=0.11 regresses Qwen2.5-VL decoding to '!').\n"62        "Re-run with the pinned image + interpreter flags:\n"63        "    hf jobs uv run --flavor a10g-small -s HF_TOKEN \\\n"64        "        --image vllm/vllm-openai:v0.10.2 --python /usr/bin/python3 \\\n"65        "        -e PYTHONPATH=/usr/local/lib/python3.12/dist-packages \\\n"66        "        nanonets-ocr2.py INPUT_DATASET OUTPUT_DATASET ...\n"67    )68    sys.exit(1)69 70import torch71from datasets import load_dataset72from huggingface_hub import DatasetCard, login73from PIL import Image74from toolz import partition_all75from tqdm.auto import tqdm76# Disable vLLM's FlashInfer sampler (JIT-compiles a CUDA kernel; greedy OCR doesn't use it).77# Harmless no-op on the pinned vllm/vllm-openai image; kept in case the pin is ever relaxed78# back to an environment without nvcc.79os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")80from vllm import LLM, SamplingParams81 82logging.basicConfig(level=logging.INFO)83logger = logging.getLogger(__name__)84 85 86def check_cuda_availability():87    """Check if CUDA is available and exit if not."""88    if not torch.cuda.is_available():89        logger.error("CUDA is not available. This script requires a GPU.")90        logger.error("Please run on a machine with a CUDA-capable GPU.")91        sys.exit(1)92    else:93        logger.info(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")94 95 96def ensure_output_columns_free(dataset, columns, overwrite=False):97    """Fail fast if an output column would collide with an existing input column.98 99    Adding a column that already exists silently overwrites it (e.g. a ground-truth100    `text`/`markdown` column) or crashes on push with a duplicate-column error only101    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing102    column(s) here instead (logged) so the later add_column is clean.103    """104    clash = [c for c in columns if c in dataset.column_names]105    if not clash:106        return dataset107    if overwrite:108        logger.warning(f"--overwrite: replacing existing column(s) {clash}")109        return dataset.remove_columns(clash)110    logger.error(111        f"Output column(s) {clash} already exist in the input dataset "112        f"(columns: {dataset.column_names})."113    )114    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")115    sys.exit(1)116 117 118def make_ocr_message(119    image: Union[Image.Image, Dict[str, Any], str],120    prompt: str = "Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the <img></img> tag; otherwise, add the image caption inside <img></img>. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes.",121) -> List[Dict]:122    """Create chat message for OCR processing."""123    # Convert to PIL Image if needed124    if isinstance(image, Image.Image):125        pil_img = image126    elif isinstance(image, dict) and "bytes" in image:127        pil_img = Image.open(io.BytesIO(image["bytes"]))128    elif isinstance(image, str):129        pil_img = Image.open(image)130    else:131        raise ValueError(f"Unsupported image type: {type(image)}")132 133    # Convert to base64 data URI134    buf = io.BytesIO()135    pil_img.save(buf, format="PNG")136    data_uri = f"data:image/png;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- nanonets170- nanonets-ocr2171- markdown172- uv-script173- generated174---175 176# Document OCR using {model_name}177 178This dataset contains markdown-formatted OCR results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using Nanonets-OCR2-3B.179 180## Processing Details181 182- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})183- **Model**: [{model}](https://huggingface.co/{model})184- **Model Size**: 3.75B parameters185- **Number of Samples**: {num_samples:,}186- **Processing Time**: {processing_time}187- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}188 189### Configuration190 191- **Image Column**: `{image_column}`192- **Output Column**: `markdown`193- **Dataset Split**: `{split}`194- **Batch Size**: {batch_size}195- **Max Model Length**: {max_model_len:,} tokens196- **Max Output Tokens**: {max_tokens:,}197- **GPU Memory Utilization**: {gpu_memory_utilization:.1%}198 199## Model Information200 201Nanonets-OCR2-3B is a state-of-the-art document OCR model that excels at:202- 📐 **LaTeX equations** - Mathematical formulas preserved in LaTeX format203- 📊 **Tables** - Extracted and formatted as HTML204- 📝 **Document structure** - Headers, lists, and formatting maintained205- 🖼️ **Images** - Captions and descriptions included in `<img>` tags206- ☑️ **Forms** - Checkboxes rendered as ☐/☑207- 🔖 **Watermarks** - Wrapped in `<watermark>` tags208- 📄 **Page numbers** - Wrapped in `<page_number>` tags209- 🌍 **Multilingual** - Supports multiple languages210 211## Dataset Structure212 213The dataset contains all original columns plus:214- `markdown`: The extracted text in markdown format with preserved structure215- `inference_info`: JSON list tracking all OCR models applied to this dataset216 217## Usage218 219```python220from datasets import load_dataset221import json222 223# Load the dataset224dataset = load_dataset("{{{{output_dataset_id}}}}", split="{split}")225 226# Access the markdown text227for example in dataset:228    print(example["markdown"])229    break230 231# View all OCR models applied to this dataset232inference_info = json.loads(dataset[0]["inference_info"])233for info in inference_info:234    print(f"Column: {{{{info['column_name']}}}} - Model: {{{{info['model_id']}}}}")235```236 237## Reproduction238 239This dataset was generated using the [uv-scripts/ocr](https://huggingface.co/datasets/uv-scripts/ocr) Nanonets OCR2 script:240 241```bash242uv run https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nanonets-ocr2.py \\243    {source_dataset} \\244    <output-dataset> \\245    --model {model} \\246    --image-column {image_column} \\247    --batch-size {batch_size} \\248    --max-model-len {max_model_len} \\249    --max-tokens {max_tokens} \\250    --gpu-memory-utilization {gpu_memory_utilization}251```252 253## Performance254 255- **Processing Speed**: ~{num_samples / (float(processing_time.split()[0]) * 60):.1f} images/second256- **GPU Configuration**: vLLM with {gpu_memory_utilization:.0%} GPU memory utilization257 258Generated with 🤖 [UV Scripts](https://huggingface.co/uv-scripts)259"""260 261 262def main(263    input_dataset: str,264    output_dataset: str,265    image_column: str = "image",266    batch_size: int = 16,267    model: str = "nanonets/Nanonets-OCR2-3B",268    max_model_len: int = 32768,269    max_tokens: int = 15000,270    gpu_memory_utilization: float = 0.8,271    hf_token: str = None,272    split: str = "train",273    max_samples: int = None,274    private: bool = False,275    shuffle: bool = False,276    seed: int = 42,277    output_column: str = "markdown",278    overwrite: bool = False,279    verbose: bool = False,280    config: str = None,281    create_pr: bool = False,282):283    """Process images from HF dataset through Nanonets-OCR2-3B model."""284 285    # Check CUDA availability first286    check_cuda_availability()287 288    # Track processing start time289    start_time = datetime.now()290 291    # Enable high-performance Xet downloads292    os.environ["HF_XET_HIGH_PERFORMANCE"] = "1"293 294    # Login to HF if token provided295    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")296    if HF_TOKEN:297        login(token=HF_TOKEN)298 299    # Load dataset300    logger.info(f"Loading dataset: {input_dataset}")301    dataset = load_dataset(input_dataset, split=split)302 303    # Validate image column304    if image_column not in dataset.column_names:305        raise ValueError(306            f"Column '{image_column}' not found. Available: {dataset.column_names}"307        )308 309    # Fail fast if the output column would collide with an existing input column310    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)311 312    # Shuffle if requested313    if shuffle:314        logger.info(f"Shuffling dataset with seed {seed}")315        dataset = dataset.shuffle(seed=seed)316 317    # Limit samples if requested318    if max_samples:319        dataset = dataset.select(range(min(max_samples, len(dataset))))320        logger.info(f"Limited to {len(dataset)} samples")321 322    # Initialize vLLM323    logger.info(f"Initializing vLLM with model: {model}")324    llm = LLM(325        model=model,326        trust_remote_code=True,327        max_model_len=max_model_len,328        gpu_memory_utilization=gpu_memory_utilization,329        limit_mm_per_prompt={"image": 1},330    )331 332    sampling_params = SamplingParams(333        temperature=0.0,  # Deterministic for OCR334        max_tokens=max_tokens,335    )336 337    # Process images in batches338    all_markdown = []339 340    logger.info(f"Processing {len(dataset)} images in batches of {batch_size}")341 342    # Process in batches to avoid memory issues343    for batch_indices in tqdm(344        partition_all(batch_size, range(len(dataset))),345        total=(len(dataset) + batch_size - 1) // batch_size,346        desc="OCR processing",347    ):348        batch_indices = list(batch_indices)349        batch_images = [dataset[i][image_column] for i in batch_indices]350 351        try:352            # Create messages for batch353            batch_messages = [make_ocr_message(img) for img in batch_images]354 355            # Process with vLLM356            outputs = llm.chat(batch_messages, sampling_params)357 358            # Extract markdown from outputs359            for output in outputs:360                markdown_text = output.outputs[0].text.strip()361                all_markdown.append(markdown_text)362 363        except Exception as e:364            logger.error(f"Error processing batch: {e}")365            # Add error placeholders for failed batch366            all_markdown.extend(["[OCR FAILED]"] * len(batch_images))367 368    # Add output column to dataset369    logger.info(f"Adding '{output_column}' column to dataset")370    dataset = dataset.add_column(output_column, all_markdown)371 372    # Handle inference_info tracking373    logger.info("Updating inference_info...")374 375    inference_entry = {376        "model_id": model,377        "model_name": "Nanonets-OCR2-3B",378        "column_name": output_column,379        "timestamp": datetime.now().isoformat(),380        "batch_size": batch_size,381        "max_tokens": max_tokens,382        "gpu_memory_utilization": gpu_memory_utilization,383        "max_model_len": max_model_len,384        "script": "nanonets-ocr2.py",385        "script_url": "https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nanonets-ocr2.py",386    }387 388    if "inference_info" in dataset.column_names:389        logger.info("Updating existing inference_info column")390 391        def update_inference_info(example):392            try:393                existing_info = (394                    json.loads(example["inference_info"])395                    if example["inference_info"]396                    else []397                )398            except (json.JSONDecodeError, TypeError):399                existing_info = []400            existing_info.append(inference_entry)401            return {"inference_info": json.dumps(existing_info)}402 403        dataset = dataset.map(update_inference_info)404    else:405        logger.info("Creating new inference_info column")406        inference_list = [json.dumps([inference_entry])] * len(dataset)407        dataset = dataset.add_column("inference_info", inference_list)408 409    # Push to hub410    logger.info(f"Pushing to {output_dataset}")411    commit_info = dataset.push_to_hub(412        output_dataset,413        private=private,414        token=HF_TOKEN,415        **({"config_name": config} if config else {}),416        create_pr=create_pr,417        commit_message=f"Add {model} OCR results ({len(dataset)} samples)"418        + (f" [{config}]" if config else ""),419    )420 421    # Calculate processing time422    end_time = datetime.now()423    processing_duration = end_time - start_time424    processing_time = f"{processing_duration.total_seconds() / 60:.1f} minutes"425 426    # Create and push dataset card427    logger.info("Creating dataset card...")428    card_content = create_dataset_card(429        source_dataset=input_dataset,430        model=model,431        num_samples=len(dataset),432        processing_time=processing_time,433        batch_size=batch_size,434        max_model_len=max_model_len,435        max_tokens=max_tokens,436        gpu_memory_utilization=gpu_memory_utilization,437        image_column=image_column,438        split=split,439    )440 441    card = DatasetCard(card_content)442    card.push_to_hub(output_dataset, token=HF_TOKEN)443    logger.info("✅ Dataset card created and pushed!")444 445    logger.info("✅ OCR conversion complete!")446    logger.info(447        f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"448    )449    if create_pr and getattr(commit_info, "pr_url", None):450        logger.info(f"Pull request created: {commit_info.pr_url}")451 452    if verbose:453        import importlib.metadata454 455        logger.info("--- Resolved package versions ---")456        for pkg in ["vllm", "transformers", "torch", "datasets", "pyarrow", "pillow"]:457            try:458                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")459            except importlib.metadata.PackageNotFoundError:460                logger.info(f"  {pkg}: not installed")461        logger.info("--- End versions ---")462 463 464if __name__ == "__main__":465    # Show example usage if no arguments466    if len(sys.argv) == 1:467        print("=" * 80)468        print("Nanonets OCR2-3B to Markdown Converter")469        print("=" * 80)470        print("\nThis script converts document images to structured markdown using")471        print("the Nanonets-OCR2-3B model (3.75B params) with vLLM acceleration.")472        print("\nFeatures:")473        print("- LaTeX equation recognition")474        print("- Table extraction and formatting (HTML)")475        print("- Document structure preservation")476        print("- Image descriptions and captions")477        print("- Signature and watermark detection")478        print("- Checkbox recognition (☐/☑)")479        print("- Multilingual support")480        print("\nExample usage:")481        print("\n1. Basic OCR conversion:")482        print("   uv run nanonets-ocr2.py document-images markdown-docs")483        print("\n2. With custom settings:")484        print("   uv run nanonets-ocr2.py scanned-pdfs extracted-text \\")485        print("       --image-column page \\")486        print("       --batch-size 32 \\")487        print("       --gpu-memory-utilization 0.8")488        print("\n3. Process a subset for testing:")489        print("   uv run nanonets-ocr2.py large-dataset test-output --max-samples 10")490        print("\n4. Random sample from ordered dataset:")491        print("   uv run nanonets-ocr2.py ordered-dataset random-test \\")492        print("       --max-samples 50 --shuffle")493        print("\n5. Running on HF Jobs:")494        print("   hf jobs uv run --flavor l4x1 \\")495        print(496            '     -e HF_TOKEN=$(python3 -c "from huggingface_hub import get_token; print(get_token())") \\'497        )498        print(499            "     https://huggingface.co/datasets/uv-scripts/ocr/raw/main/nanonets-ocr2.py \\"500        )501        print("       your-document-dataset \\")502        print("       your-markdown-output")503        print("\n" + "=" * 80)504        print("\nFor full help, run: uv run nanonets-ocr2.py --help")505        sys.exit(0)506 507    parser = argparse.ArgumentParser(508        description="OCR images to markdown using Nanonets-OCR2-3B",509        formatter_class=argparse.RawDescriptionHelpFormatter,510        epilog="""511Examples:512  # Basic usage513  uv run nanonets-ocr2.py my-images-dataset ocr-results514 515  # With specific image column516  uv run nanonets-ocr2.py documents extracted-text --image-column scan517 518  # Process subset for testing519  uv run nanonets-ocr2.py large-dataset test-output --max-samples 100520 521  # Random sample from ordered dataset522  uv run nanonets-ocr2.py ordered-dataset random-sample --max-samples 50 --shuffle523        """,524    )525 526    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")527    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")528    parser.add_argument(529        "--image-column",530        default="image",531        help="Column containing images (default: image)",532    )533    parser.add_argument(534        "--batch-size",535        type=int,536        default=16,537        help="Batch size for processing (default: 16)",538    )539    parser.add_argument(540        "--model",541        default="nanonets/Nanonets-OCR2-3B",542        help="Model to use (default: nanonets/Nanonets-OCR2-3B)",543    )544    parser.add_argument(545        "--max-model-len",546        type=int,547        default=32768,548        help="Maximum model context length (default: 32768)",549    )550    parser.add_argument(551        "--max-tokens",552        type=int,553        default=15000,554        help="Maximum tokens to generate (default: 15000, per model card recommendation)",555    )556    parser.add_argument(557        "--gpu-memory-utilization",558        type=float,559        default=0.8,560        help="GPU memory utilization (default: 0.8)",561    )562    parser.add_argument("--hf-token", help="Hugging Face API token")563    parser.add_argument(564        "--split", default="train", help="Dataset split to use (default: train)"565    )566    parser.add_argument(567        "--max-samples",568        type=int,569        help="Maximum number of samples to process (for testing)",570    )571    parser.add_argument(572        "--private", action="store_true", help="Make output dataset private"573    )574    parser.add_argument(575        "--config",576        help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",577    )578    parser.add_argument(579        "--create-pr",580        action="store_true",581        help="Create a pull request instead of pushing directly (for parallel benchmarking)",582    )583    parser.add_argument(584        "--shuffle",585        action="store_true",586        help="Shuffle the dataset before processing (useful for random sampling)",587    )588    parser.add_argument(589        "--seed",590        type=int,591        default=42,592        help="Random seed for shuffling (default: 42)",593    )594    parser.add_argument(595        "--output-column",596        default="markdown",597        help="Column name for the OCR output text (default: markdown)",598    )599    parser.add_argument(600        "--overwrite",601        action="store_true",602        help="Replace the output column if it already exists in the input dataset "603        "(default: error out to avoid clobbering an existing column).",604    )605    parser.add_argument(606        "--verbose",607        action="store_true",608        help="Log resolved package versions after processing (useful for pinning deps)",609    )610 611    args = parser.parse_args()612 613    main(614        input_dataset=args.input_dataset,615        output_dataset=args.output_dataset,616        image_column=args.image_column,617        batch_size=args.batch_size,618        model=args.model,619        max_model_len=args.max_model_len,620        max_tokens=args.max_tokens,621        gpu_memory_utilization=args.gpu_memory_utilization,622        hf_token=args.hf_token,623        split=args.split,624        max_samples=args.max_samples,625        private=args.private,626        shuffle=args.shuffle,627        seed=args.seed,628        output_column=args.output_column,629        overwrite=args.overwrite,630        verbose=args.verbose,631        config=args.config,632        create_pr=args.create_pr,633    )634