uv-scripts/ocr

1# /// script2# requires-python = ">=3.11"3# dependencies = [4#     "datasets>=4.0.0",5#     "huggingface-hub",6#     "pillow",7#     "requests",8# ]9# ///10 11"""12Convert document images to markdown using LightOnOCR-2 via an in-job vLLM server.13 14Same model, message shape, and sampling as lighton-ocr2.py, but serves the15model behind `vllm serve` inside the job (the model card's own documented16path) and posts images concurrently — continuous batching stays fed instead17of draining at each offline batch barrier, and a bad image fails one request18instead of a whole batch. Measured ~1.8x the offline recipe's inference19throughput on a 100-page historical-scan smoke test (l4x1, concurrency 32).20 21This script is the *driver* half: it expects the server on localhost22(started by the job command below), loads the input dataset, posts images23concurrently, and pushes the result dataset. The driver has no torch/vllm24deps, so `uv run` starts in seconds while the server warms up in parallel.25 26Run on HF Jobs (standard uv-run shape — the script starts `vllm serve` itself27as a subprocess when no server is already reachable; the only thing to get28right is the --image flag, which provides the `vllm` binary):29 30  hf jobs uv run --detach --flavor l4x1 -s HF_TOKEN --timeout 4h \\31      --image vllm/vllm-openai:v0.22.1 \\32      https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr2-server.py \\33      <input-dataset> <output-dataset>34 35To use an already-running or remote endpoint instead, pass --server URL — the36script only spawns a server when the (default localhost) URL is unreachable.37The serve flags live in SERVE_ARGS below (the model card's own recommended38command: `--limit-mm-per-prompt`, `--mm-processor-cache-gb 0`,39`--no-enable-prefix-caching` — OCR never reuses images, so the caches only40cost memory).41 42Model: lightonai/LightOnOCR-2-1B (1B, Apache-2.0)43- Message is the image ONLY (no text prompt) — LightOnOCR-2's trained format.44- Images resized client-side so the longest dimension is 1540px (training45  resolution at 200 DPI), same as the offline recipe.46- Sampling per the card: temperature 0.2, top_p 0.9, max_tokens 4096.47"""48 49import argparse50import atexit51import base6452import concurrent.futures53import io54import json55import logging56import os57import shutil58import subprocess59import sys60import threading61import time62from datetime import datetime63from typing import Any, Dict, Union64from urllib.parse import urlparse65 66import requests67from datasets import load_dataset68from huggingface_hub import DatasetCard, login69from PIL import Image70 71logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")72logger = logging.getLogger(__name__)73 74MODEL = "lightonai/LightOnOCR-2-1B"75 76DEFAULT_TARGET_SIZE = 1540  # longest dimension; LightOnOCR-2 training resolution77 78# The serve command this script spawns when no server is reachable — the model79# card's own recommended flags; single source of truth for the serving config.80SERVE_ARGS = [81    "vllm", "serve", MODEL,82    "--limit-mm-per-prompt", '{"image": 1}',83    "--mm-processor-cache-gb", "0",84    "--no-enable-prefix-caching",85    "--max-model-len", "8192",86    "--gpu-memory-utilization", "0.8",87    "--port", "8000",88]89 90RUN_COMMAND = (91    "hf jobs uv run --detach --flavor l4x1 -s HF_TOKEN --timeout 4h \\\n"92    "    --image vllm/vllm-openai:v0.22.1 \\\n"93    "    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr2-server.py \\\n"94    "    <input-dataset> <output-dataset>"95)96 97 98def ensure_output_columns_free(dataset, columns, overwrite=False):99    """Fail fast if an output column would collide with an existing input column.100 101    Adding a column that already exists silently overwrites it (e.g. a ground-truth102    `text`/`markdown` column) or crashes on push with a duplicate-column error only103    *after* inference has run. Catch it up front. With overwrite=True, drop the clashing104    column(s) here instead (logged) so the later add_column is clean.105    """106    clash = [c for c in columns if c in dataset.column_names]107    if not clash:108        return dataset109    if overwrite:110        logger.warning(f"--overwrite: replacing existing column(s) {clash}")111        return dataset.remove_columns(clash)112    logger.error(113        f"Output column(s) {clash} already exist in the input dataset "114        f"(columns: {dataset.column_names})."115    )116    logger.error("Choose a different --output-column, or pass --overwrite to replace them.")117    sys.exit(1)118 119 120def to_pil_image(image: Union[Image.Image, Dict[str, Any], str]) -> Image.Image:121    """Convert a dataset image cell (PIL image, bytes dict, or path) to RGB PIL."""122    if isinstance(image, Image.Image):123        pil_img = image124    elif isinstance(image, dict) and "bytes" in image:125        pil_img = Image.open(io.BytesIO(image["bytes"]))126    elif isinstance(image, str):127        pil_img = Image.open(image)128    else:129        raise ValueError(f"Unsupported image type: {type(image)}")130    return pil_img.convert("RGB")131 132 133def encode_image(image, target_size: int) -> str:134    """RGB-convert, resize longest dimension to target_size, return base64 PNG."""135    img = to_pil_image(image)136    if target_size:137        w, h = img.size138        if max(w, h) != target_size:139            scale = target_size / max(w, h)140            img = img.resize(141                (max(1, int(w * scale)), max(1, int(h * scale))),142                Image.Resampling.LANCZOS,143            )144    buf = io.BytesIO()145    img.save(buf, format="PNG")146    return base64.b64encode(buf.getvalue()).decode()147 148 149def server_alive(server: str) -> bool:150    try:151        return requests.get(f"{server}/health", timeout=5).status_code == 200152    except requests.RequestException:153        return False154 155 156def wait_for_server(server: str, timeout_s: int, proc: "subprocess.Popen | None" = None):157    logger.info(f"Waiting for server at {server}...")158    deadline = time.time() + timeout_s159    while time.time() < deadline:160        if server_alive(server):161            logger.info("Server is ready")162            return163        if proc is not None and proc.poll() is not None:164            logger.error(f"Spawned vllm serve exited with code {proc.returncode} before becoming ready")165            sys.exit(1)166        time.sleep(10)167    logger.error(f"Server did not become ready within {timeout_s}s")168    sys.exit(1)169 170 171def ensure_server(server: str, timeout_s: int = 1800):172    """Use a reachable server; otherwise spawn `vllm serve` ourselves; else fail fast.173 174    Spawning is only attempted for a localhost URL — a remote --server that is175    down is the user's to fix, not ours to shadow with a local model.176    """177    if server_alive(server):178        logger.info(f"Using already-running server at {server}")179        return180 181    host = urlparse(server).hostname or ""182    if host not in ("127.0.0.1", "localhost", "0.0.0.0"):183        logger.info(f"Remote server {server} not up yet — waiting for it")184        wait_for_server(server, timeout_s)185        return186 187    if shutil.which("vllm") is None:188        logger.error("No server is running and the `vllm` binary is not on PATH.")189        logger.error("Run this script on a vLLM image so it can start the server itself:\n")190        logger.error(RUN_COMMAND)191        logger.error("\n(or start `vllm serve` yourself / pass --server URL of a running endpoint)")192        sys.exit(1)193 194    logger.info(f"Starting server: {' '.join(SERVE_ARGS)}")195    proc = subprocess.Popen(SERVE_ARGS)  # logs interleave with ours on stdout/stderr196    atexit.register(proc.terminate)  # don't leave a GPU server behind on local runs197    wait_for_server(server, timeout_s, proc=proc)198 199 200def ocr_one(201    server: str,202    image,203    target_size: int,204    max_tokens: int,205    temperature: float,206    top_p: float,207    timeout_s: int,208    retries: int = 2,209) -> str:210    """OCR a single image via the chat completions API. Returns raw model text."""211    b64 = encode_image(image, target_size)212    payload = {213        "model": MODEL,214        "messages": [215            {216                "role": "user",217                "content": [218                    # Image ONLY — LightOnOCR-2 uses no text prompt.219                    {220                        "type": "image_url",221                        "image_url": {"url": f"data:image/png;base64,{b64}"},222                    },223                ],224            }225        ],226        "temperature": temperature,227        "top_p": top_p,228        "max_tokens": max_tokens,229    }230    last_err = None231    for attempt in range(retries + 1):232        try:233            resp = requests.post(234                f"{server}/v1/chat/completions", json=payload, timeout=timeout_s235            )236            resp.raise_for_status()237            return resp.json()["choices"][0]["message"]["content"]238        except Exception as e:239            last_err = e240            if attempt < retries:241                time.sleep(10 * (attempt + 1))242    raise RuntimeError(f"request failed after {retries + 1} attempts: {last_err}")243 244 245def create_dataset_card(246    source_dataset: str,247    model: str,248    num_samples: int,249    num_errors: int,250    processing_time: str,251    images_per_sec: float,252    concurrency: int,253    max_tokens: int,254    temperature: float,255    target_size: int,256    image_column: str = "image",257    split: str = "train",258) -> str:259    """Create a dataset card documenting the OCR process."""260    model_name = model.split("/")[-1]261 262    # Canonical provenance stamp (see AGENTS.md): Jobs claim gated on JOB_ID, set by HF Jobs in-container.263    on_jobs = os.environ.get("JOB_ID") is not None264    hw = os.environ.get("ACCELERATOR") or ""265    origin = (266        "Produced on [Hugging Face Jobs](https://huggingface.co/docs/huggingface_hub/guides/jobs)"267        + (f" (`{hw}`)" if hw else "")268    ) if on_jobs else "Generated"269    jobs_tag = "\n- hf-jobs" if on_jobs else ""270 271    return f"""---272tags:273- ocr274- document-processing275- lighton-ocr276- markdown277- uv-script278- generated{jobs_tag}279---280 281# Document OCR using {model_name} (server mode)282 283This dataset contains OCR results from images in [{source_dataset}](https://huggingface.co/datasets/{source_dataset}) using LightOnOCR-2 (1B), served behind an in-job vLLM server with concurrent requests (continuous batching).284 285## Processing Details286 287- **Source Dataset**: [{source_dataset}](https://huggingface.co/datasets/{source_dataset})288- **Model**: [{model}](https://huggingface.co/{model})289- **Number of Samples**: {num_samples:,}290- **Failed Requests**: {num_errors:,} (marked `[OCR ERROR]`)291- **Processing Time**: {processing_time}292- **Throughput**: {images_per_sec:.2f} images/sec293- **Processing Date**: {datetime.now().strftime("%Y-%m-%d %H:%M UTC")}294 295### Configuration296 297- **Mode**: vLLM server (`vllm serve`) + concurrent driver, {concurrency} concurrent requests298- **Image Column**: `{image_column}`299- **Dataset Split**: `{split}`300- **Target Image Size**: {target_size}px (longest dimension)301- **Max Output Tokens**: {max_tokens:,}302- **Temperature**: {temperature}303 304## Dataset Structure305 306The dataset contains all original columns plus:307- `markdown`: The extracted text in markdown format308- `inference_info`: JSON list tracking all OCR models applied to this dataset309 310## Reproduction311 312{origin} with the [`lighton-ocr2-server.py`](https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr2-server.py) recipe from [uv-scripts](https://huggingface.co/uv-scripts) — see the script docstring for the single `hf jobs run` command that starts the server and driver together. The offline-vLLM sibling recipe is [`lighton-ocr2.py`](https://huggingface.co/datasets/uv-scripts/ocr/raw/main/lighton-ocr2.py).313"""314 315 316def main(317    input_dataset: str,318    output_dataset: str,319    image_column: str = "image",320    server: str = "http://127.0.0.1:8000",321    concurrency: int = 32,322    max_tokens: int = 4096,323    temperature: float = 0.2,324    top_p: float = 0.9,325    target_size: int = DEFAULT_TARGET_SIZE,326    request_timeout: int = 1800,327    hf_token: str = None,328    split: str = "train",329    max_samples: int = None,330    private: bool = False,331    shuffle: bool = False,332    seed: int = 42,333    output_column: str = "markdown",334    overwrite: bool = False,335    config: str = None,336    create_pr: bool = False,337):338    """Process images from HF dataset through a LightOnOCR-2 vLLM server."""339 340    start_time = datetime.now()341 342    HF_TOKEN = hf_token or os.environ.get("HF_TOKEN")343    if HF_TOKEN:344        login(token=HF_TOKEN)345 346    logger.info(f"Using model: {MODEL} via server {server}")347 348    logger.info(f"Loading dataset: {input_dataset}")349    dataset = load_dataset(input_dataset, split=split)350 351    if image_column not in dataset.column_names:352        raise ValueError(353            f"Column '{image_column}' not found. Available: {dataset.column_names}"354        )355 356    dataset = ensure_output_columns_free(dataset, [output_column], overwrite=overwrite)357 358    if shuffle:359        logger.info(f"Shuffling dataset with seed {seed}")360        dataset = dataset.shuffle(seed=seed)361 362    if max_samples:363        dataset = dataset.select(range(min(max_samples, len(dataset))))364        logger.info(f"Limited to {len(dataset)} samples")365 366    # Reuse a reachable server, else spawn `vllm serve` (needs the vllm binary,367    # i.e. a vllm/vllm-openai image), else fail fast with the correct command.368    ensure_server(server)369 370    n = len(dataset)371    logger.info(f"Processing {n} images, concurrency {concurrency}")372    all_outputs = [None] * n373    errors = 0374    done = 0375    inference_start = time.time()376    lock = threading.Lock()377 378    def worker(i: int) -> None:379        nonlocal errors, done380        try:381            text = ocr_one(382                server,383                dataset[i][image_column],384                target_size,385                max_tokens,386                temperature,387                top_p,388                request_timeout,389            )390            all_outputs[i] = text.strip()391        except Exception as e:392            logger.error(f"Image {i} failed: {e}")393            all_outputs[i] = "[OCR ERROR]"394            with lock:395                errors += 1396        with lock:397            done += 1398            if done % 25 == 0 or done == n:399                rate = done / max(time.time() - inference_start, 1e-9)400                logger.info(f"{done}/{n} done ({rate:.2f} img/s, {errors} errors)")401 402    with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool:403        list(pool.map(worker, range(n)))404 405    inference_secs = time.time() - inference_start406    processing_duration = datetime.now() - start_time407    processing_time_str = f"{processing_duration.total_seconds() / 60:.1f} min"408    images_per_sec = n / inference_secs if inference_secs else 0.0409 410    logger.info(f"Adding '{output_column}' column to dataset")411    dataset = dataset.add_column(output_column, all_outputs)412 413    # Inference info tracking414    inference_entry = {415        "model_id": MODEL,416        "model_name": "LightOnOCR-2-1B",417        "column_name": output_column,418        "timestamp": datetime.now().isoformat(),419        "temperature": temperature,420        "top_p": top_p,421        "max_tokens": max_tokens,422        "target_size": target_size,423        "mode": "vllm-server",424        "concurrency": concurrency,425    }426 427    if "inference_info" in dataset.column_names:428        logger.info("Updating existing inference_info column")429 430        def update_inference_info(example):431            try:432                existing_info = (433                    json.loads(example["inference_info"])434                    if example["inference_info"]435                    else []436                )437            except (json.JSONDecodeError, TypeError):438                existing_info = []439            existing_info.append(inference_entry)440            return {"inference_info": json.dumps(existing_info)}441 442        dataset = dataset.map(update_inference_info)443    else:444        logger.info("Creating new inference_info column")445        inference_list = [json.dumps([inference_entry])] * len(dataset)446        dataset = dataset.add_column("inference_info", inference_list)447 448    # Push to hub with retry and XET fallback449    logger.info(f"Pushing to {output_dataset}")450    max_retries = 3451    for attempt in range(1, max_retries + 1):452        try:453            if attempt > 1:454                logger.warning("Disabling XET (fallback to HTTP upload)")455                os.environ["HF_HUB_DISABLE_XET"] = "1"456            dataset.push_to_hub(457                output_dataset,458                private=private,459                token=HF_TOKEN,460                max_shard_size="500MB",461                **({"config_name": config} if config else {}),462                create_pr=create_pr,463                commit_message=f"Add {MODEL} OCR results ({len(dataset)} samples, server mode)"464                + (f" [{config}]" if config else ""),465            )466            break467        except Exception as e:468            logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")469            if attempt < max_retries:470                delay = 30 * (2 ** (attempt - 1))471                logger.info(f"Retrying in {delay}s...")472                time.sleep(delay)473            else:474                logger.error("All upload attempts failed. OCR results are lost.")475                sys.exit(1)476 477    logger.info("Creating dataset card")478    card_content = create_dataset_card(479        source_dataset=input_dataset,480        model=MODEL,481        num_samples=len(dataset),482        num_errors=errors,483        processing_time=processing_time_str,484        images_per_sec=images_per_sec,485        concurrency=concurrency,486        max_tokens=max_tokens,487        temperature=temperature,488        target_size=target_size,489        image_column=image_column,490        split=split,491    )492 493    card = DatasetCard(card_content)494    card.push_to_hub(output_dataset, token=HF_TOKEN)495 496    logger.info("Done! LightOnOCR-2 server-mode processing complete.")497    logger.info(498        f"Dataset available at: https://huggingface.co/datasets/{output_dataset}"499    )500    logger.info(f"Processing time: {processing_time_str}")501    logger.info(502        f"Throughput: {images_per_sec:.2f} images/sec "503        f"(inference only, excl. dataset load/push; {errors} errors)"504    )505 506 507if __name__ == "__main__":508    if len(sys.argv) == 1:509        print("=" * 70)510        print("LightOnOCR-2 Document Processing (vLLM server mode)")511        print("=" * 70)512        print("\nSame model + outputs as lighton-ocr2.py, but drives an in-job")513        print("`vllm serve` with concurrent requests — no batch barriers,")514        print("per-image (not per-batch) failure isolation.")515        print("\nThe server must already be running (the job command starts")516        print("both — see the module docstring for the full `hf jobs run`).")517        print("\nExamples:")518        print("\n1. Basic OCR (server on localhost:8000):")519        print("   uv run lighton-ocr2-server.py input-dataset output-dataset")520        print("\n2. Test with a small sample:")521        print("   uv run lighton-ocr2-server.py large-dataset test --max-samples 10 --shuffle")522        print("\nFor full help: uv run lighton-ocr2-server.py --help")523        sys.exit(0)524 525    parser = argparse.ArgumentParser(526        description="Document OCR using LightOnOCR-2 via an in-job vLLM server",527        formatter_class=argparse.RawDescriptionHelpFormatter,528        epilog="""529Examples:530  uv run lighton-ocr2-server.py my-docs analyzed-docs531  uv run lighton-ocr2-server.py large-dataset test --max-samples 50 --shuffle532See the module docstring for the full `hf jobs run` command (server + driver in one job).533        """,534    )535 536    parser.add_argument("input_dataset", help="Input dataset ID from Hugging Face Hub")537    parser.add_argument("output_dataset", help="Output dataset ID for Hugging Face Hub")538    parser.add_argument(539        "--image-column",540        default="image",541        help="Column containing images (default: image)",542    )543    parser.add_argument(544        "--server",545        default="http://127.0.0.1:8000",546        help="vLLM server base URL (default: in-job localhost:8000)",547    )548    parser.add_argument(549        "--concurrency",550        type=int,551        default=32,552        help="Concurrent OCR requests (default: 32; vLLM queues excess internally, "553        "so this mainly needs to be high enough to keep continuous batching fed)",554    )555    parser.add_argument(556        "--max-tokens",557        type=int,558        default=4096,559        help="Maximum tokens to generate (default: 4096, the model card value)",560    )561    parser.add_argument(562        "--temperature",563        type=float,564        default=0.2,565        help="Sampling temperature (default: 0.2, the model card value)",566    )567    parser.add_argument(568        "--top-p",569        type=float,570        default=0.9,571        help="Top-p sampling (default: 0.9, the model card value)",572    )573    parser.add_argument(574        "--target-size",575        type=int,576        default=DEFAULT_TARGET_SIZE,577        help=f"Resize images so the longest dimension is this many pixels before upload "578        f"(default: {DEFAULT_TARGET_SIZE}, the model's training resolution); 0 disables",579    )580    parser.add_argument(581        "--request-timeout",582        type=int,583        default=1800,584        help="Per-request timeout in seconds (default: 1800)",585    )586    parser.add_argument("--hf-token", help="Hugging Face API token")587    parser.add_argument(588        "--split", default="train", help="Dataset split to use (default: train)"589    )590    parser.add_argument(591        "--max-samples",592        type=int,593        help="Maximum number of samples to process (for testing)",594    )595    parser.add_argument(596        "--private", action="store_true", help="Make output dataset private"597    )598    parser.add_argument(599        "--config",600        help="Config/subset name when pushing to Hub (for benchmarking multiple models in one repo)",601    )602    parser.add_argument(603        "--create-pr",604        action="store_true",605        help="Create a pull request instead of pushing directly (for parallel benchmarking)",606    )607    parser.add_argument(608        "--shuffle", action="store_true", help="Shuffle dataset before processing"609    )610    parser.add_argument(611        "--seed",612        type=int,613        default=42,614        help="Random seed for shuffling (default: 42)",615    )616    parser.add_argument(617        "--output-column",618        default="markdown",619        help="Column name for output text (default: markdown)",620    )621    parser.add_argument(622        "--overwrite",623        action="store_true",624        help="Replace the output column if it already exists in the input dataset "625        "(default: error out to avoid clobbering an existing column).",626    )627 628    args = parser.parse_args()629 630    main(631        input_dataset=args.input_dataset,632        output_dataset=args.output_dataset,633        image_column=args.image_column,634        server=args.server,635        concurrency=args.concurrency,636        max_tokens=args.max_tokens,637        temperature=args.temperature,638        top_p=args.top_p,639        target_size=args.target_size,640        request_timeout=args.request_timeout,641        hf_token=args.hf_token,642        split=args.split,643        max_samples=args.max_samples,644        private=args.private,645        shuffle=args.shuffle,646        seed=args.seed,647        output_column=args.output_column,648        overwrite=args.overwrite,649        config=args.config,650        create_pr=args.create_pr,651    )652