Suya OCR API — vLLM-backed, OpenAI-compatible OCR service

FastAPI service wrapping the Surya-OCR-2 model (datalab-to) served through vLLM:
legacy /v1/api/ai/* endpoints, an OpenAI-compatible /v1/chat/completions endpoint,
a coalescing request batcher, a local OCR CLI, Docker packaging, multilingual
example outputs, and quantization/concurrency benchmarks.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

scripts/quant/build_model.py

@@ -0,0 +1,144 @@
+"""Produce vLLM-loadable quantized checkpoints.
+
+compressor methods (int8/awq/gptq) use llm-compressor `oneshot`.
+bnb methods (bnb8/bnb4) use transformers + BitsAndBytesConfig + save_pretrained.
+baseline/online (bf16/fp8) need no build.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import List
+
+from scripts.quant.recipes import METHOD_SPECS
+
+
+def needs_build(method: str) -> bool:
+    return METHOD_SPECS[method]["kind"] in ("compressor", "bnb")
+
+
+def _is_built(out_dir: Path) -> bool:
+    return (out_dir / "config.json").exists()
+
+
+CALIB_MAX_SEQ_LEN = 8192
+# Never quantize the LM head or the vision tower: vision modules are shape-fragile
+# (see the FP8 Marlin failure) and contribute little to decode cost.
+QUANT_IGNORE = ["re:.*lm_head", "re:.*visual.*", "re:.*vision.*"]
+
+
+def _calibration_dataset(calib_images: List[Path], processor):
+    """HF Dataset of pre-tokenized multimodal samples (batch dim kept), matching
+    llm-compressor's multimodal-vision examples; re-tensorized by _data_collator.
+    Activation calibration only needs representative forward passes, not the
+    exact training prompt."""
+    from datasets import Dataset
+    from PIL import Image
+
+    samples = []
+    for path in calib_images:
+        image = Image.open(path).convert("RGB")
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "image"},
+                {"type": "text", "text": "OCR this document."},
+            ],
+        }]
+        prompt = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+        inputs = processor(
+            text=[prompt], images=[image],
+            padding=False, truncation=True, max_length=CALIB_MAX_SEQ_LEN,
+        )
+        samples.append({k: (v.tolist() if hasattr(v, "tolist") else v) for k, v in inputs.items()})
+    return Dataset.from_list(samples)
+
+
+def _data_collator(batch):
+    import torch
+
+    assert len(batch) == 1
+    return {key: torch.tensor(value) for key, value in batch[0].items()}
+
+
+def _build_compressor(method: str, base_model: str, out_dir: Path, calib_images: List[Path]) -> None:
+    from llmcompressor import oneshot
+    from llmcompressor.modifiers.quantization import GPTQModifier
+    from transformers import AutoProcessor
+
+    try:
+        from transformers import AutoModelForImageTextToText as _AutoModel
+    except ImportError:  # older transformers
+        from transformers import AutoModelForCausalLM as _AutoModel
+
+    spec = METHOD_SPECS[method]
+    model = _AutoModel.from_pretrained(base_model, dtype="auto", device_map="auto")
+    processor = AutoProcessor.from_pretrained(base_model)
+    dataset = _calibration_dataset(calib_images, processor)
+
+    if spec["modifier"] == "awq":
+        # llm-compressor main: AWQ is a transform paired with a QuantizationModifier.
+        from llmcompressor.modifiers.quantization import QuantizationModifier
+        try:
+            from llmcompressor.modifiers.transform.awq import AWQModifier
+        except ImportError:  # older layouts keep AWQModifier under modifiers.awq
+            from llmcompressor.modifiers.awq import AWQModifier
+        recipe = [
+            AWQModifier(duo_scaling=False),
+            QuantizationModifier(scheme=spec["scheme"], ignore=QUANT_IGNORE),
+        ]
+    else:
+        recipe = []
+        if spec.get("smoothquant"):
+            from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
+            recipe.append(SmoothQuantModifier(smoothing_strength=0.8))
+        recipe.append(GPTQModifier(targets="Linear", scheme=spec["scheme"], ignore=QUANT_IGNORE))
+
+    oneshot(
+        model=model,
+        dataset=dataset,
+        recipe=recipe,
+        max_seq_length=CALIB_MAX_SEQ_LEN,
+        num_calibration_samples=len(dataset),
+        data_collator=_data_collator,
+    )
+    out_dir.mkdir(parents=True, exist_ok=True)
+    model.save_pretrained(out_dir, save_compressed=True)
+    processor.save_pretrained(out_dir)
+
+
+def _build_bnb(method: str, base_model: str, out_dir: Path) -> None:
+    import torch
+    from transformers import AutoProcessor, BitsAndBytesConfig
+
+    try:
+        from transformers import AutoModelForImageTextToText as _AutoModel
+    except ImportError:
+        from transformers import AutoModelForCausalLM as _AutoModel
+
+    bits = METHOD_SPECS[method]["bits"]
+    if bits == 8:
+        config = BitsAndBytesConfig(load_in_8bit=True)
+    else:
+        config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+        )
+    model = _AutoModel.from_pretrained(base_model, quantization_config=config, device_map="auto")
+    model.save_pretrained(out_dir)
+    AutoProcessor.from_pretrained(base_model).save_pretrained(out_dir)
+
+
+def build_model(method: str, base_model: str, out_dir: Path, calib_images: List[Path]) -> Path:
+    if not needs_build(method):
+        return Path(base_model)
+    out_dir = Path(out_dir)
+    if _is_built(out_dir):
+        return out_dir
+    kind = METHOD_SPECS[method]["kind"]
+    if kind == "compressor":
+        _build_compressor(method, base_model, out_dir, calib_images)
+    else:
+        _build_bnb(method, base_model, out_dir)
+    return out_dir