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>
2026-06-17 10:20:02 +04:00
commit 1a585693be
147 changed files with 13827 additions and 0 deletions
@@ -0,0 +1,165 @@
+import os
+from typing import Callable, Dict, Optional
+
+import torch
+from dotenv import find_dotenv
+from pydantic import computed_field
+from pydantic_settings import BaseSettings
+from pathlib import Path
+from platformdirs import user_cache_dir
+
+
+class Settings(BaseSettings):
+    # General
+    TORCH_DEVICE: Optional[str] = None
+    IMAGE_DPI: int = 96  # used for layout + text detection (coarse structure)
+    IMAGE_DPI_HIGHRES: int = 192  # used for recognition + table rec (fine glyphs)
+    IN_STREAMLIT: bool = False
+    DISABLE_TQDM: bool = False
+    S3_BASE_URL: str = "https://models.datalab.to"
+    PARALLEL_DOWNLOAD_WORKERS: int = 10
+    MODEL_CACHE_DIR: str = str(Path(user_cache_dir("datalab")) / "models")
+    LOGLEVEL: str = "INFO"
+
+    # Paths
+    RESULT_DIR: str = "results"
+    BASE_DIR: str = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    FONT_DIR: str = os.path.join(BASE_DIR, "static", "fonts")
+
+    @computed_field
+    def TORCH_DEVICE_MODEL(self) -> str:
+        if self.TORCH_DEVICE is not None:
+            return self.TORCH_DEVICE
+        if torch.cuda.is_available():
+            return "cuda"
+        if torch.backends.mps.is_available():
+            return "mps"
+        return "cpu"
+
+    # ---- Surya2 inference (VLM-backed: vllm | llamacpp) ---------------------
+    SURYA_MODEL_CHECKPOINT: str = "datalab-to/surya-ocr-2"
+    SURYA_GGUF_REPO: str = "datalab-to/surya-ocr-2-gguf"
+    SURYA_GGUF_MODEL_FILE: str = "surya-2.gguf"
+    SURYA_GGUF_MMPROJ_FILE: str = "surya-2-mmproj.gguf"
+    # If set, used directly instead of HF download (handy for local-conversion testing)
+    SURYA_GGUF_LOCAL_MODEL_PATH: Optional[str] = None
+    SURYA_GGUF_LOCAL_MMPROJ_PATH: Optional[str] = None
+
+    # Backend selection
+    SURYA_INFERENCE_BACKEND: Optional[str] = None  # "vllm" | "llamacpp" | None (auto)
+    SURYA_INFERENCE_URL: Optional[str] = None  # external server, skip spawn
+    SURYA_INFERENCE_AUTOSTART: bool = True
+    # Leave an auto-spawned server running after the process exits so later
+    # commands attach to it instead of re-spawning (avoids repeated startup /
+    # model-load cost). Stop it manually when done — see `surya/inference`.
+    SURYA_INFERENCE_KEEP_ALIVE: bool = False
+    SURYA_INFERENCE_HOST: str = "127.0.0.1"
+    SURYA_INFERENCE_PORT: Optional[int] = None  # None = pick a free port
+    SURYA_INFERENCE_PARALLEL: int = 8
+    # Max concurrent in-flight chat-completion requests to the inference server
+    # per batch. Cap to roughly VLLM_MAX_NUM_SEQS so block fan-out keeps the
+    # GPU's sequence slots full without flooding the queue. Tuned empirically.
+    SURYA_INFERENCE_MAX_INFLIGHT: int = 16
+    # Per-parallel-slot KV-cache budget for the llama.cpp backend. Worst-case
+    # one OCR request: ~2k for image prefill + SURYA_MAX_TOKENS_FULL_PAGE
+    # (8192) generation + ~2k prompt/chat-template overhead ≈ 12k. Below this
+    # llama-server silently truncates outputs once a slot fills.
+    SURYA_INFERENCE_CTX_PER_SLOT: int = 12288
+    # Optional override for the *total* ctx passed to llama-server. When None
+    # (default), total = max(16384, PARALLEL * CTX_PER_SLOT). Set this only
+    # if you've hand-tuned for a specific machine.
+    SURYA_INFERENCE_CTX_SIZE: Optional[int] = None
+    SURYA_INFERENCE_TIMEOUT_SECONDS: float = 600.0
+    SURYA_INFERENCE_STARTUP_TIMEOUT: float = 600.0
+    SURYA_INFERENCE_LOGPROBS: bool = True
+    SURYA_INFERENCE_MAX_RETRIES: int = 1
+    # Force layout/table_rec output through a JSON schema via guided decoding.
+    # Eliminates malformed-JSON failures at small decode-throughput cost.
+    SURYA_GUIDED_LAYOUT: bool = True
+    # Disabled: with no minItems in TABLE_REC_JSON_SCHEMA, the constrained
+    # decoder closes the array after one element at temperature=0. The model
+    # produces well-formed JSON without the schema.
+    SURYA_GUIDED_TABLE_REC: bool = False
+
+    # Token budgets
+    SURYA_MAX_TOKENS_LAYOUT: int = 3072
+    SURYA_MAX_TOKENS_TABLE_REC: int = 3072
+    SURYA_MAX_TOKENS_BLOCK_CEILING: int = 8192
+    SURYA_MAX_TOKENS_FULL_PAGE: int = 6144
+    SURYA_MAX_BLOCKS_PER_PAGE: int = 80
+
+    BBOX_SCALE: int = 1000
+
+    # vllm
+    VLLM_DOCKER_IMAGE: str = "vllm/vllm-openai:v0.20.1"
+    VLLM_API_KEY: str = "EMPTY"
+    VLLM_GPUS: str = "0"
+    VLLM_GPU_TYPE: str = "4090"
+    # bfloat16 needs an Ampere+ GPU (compute capability >= 8.0). On older cards
+    # (e.g. T4 / Turing) vllm refuses to start with bf16 — set float16 there.
+    VLLM_DTYPE: str = "bfloat16"
+    VLLM_MAX_MODEL_LEN: int = 18000
+    VLLM_GPU_MEMORY_UTILIZATION: float = 0.85
+    # MTP speculative decoding only feeds the nested-Docker spawn path (vllm.py),
+    # which is disabled in production; the real launch (start_single_container.sh)
+    # never passes --speculative-config. Benchmarked +27-36% SLOWER at OCR's 16-wide
+    # block concurrency (GPU already compute-bound) — see
+    # docs/quantization_benchmark_results.md §5b. Do not wire MTP into the launch.
+    VLLM_ENABLE_MTP: bool = True
+    VLLM_MTP_TOKENS: int = 2
+    VLLM_EXTRA_ARGS: Optional[str] = None
+    DOCKER_HF_CACHE_PATH: str = "~/.cache/huggingface"
+
+    # llama.cpp
+    LLAMA_CPP_BINARY: str = "llama-server"
+    LLAMA_CPP_NGL: int = 99  # all layers on GPU (Metal on macOS, CUDA on Linux GPU); harmless no-op on pure-CPU builds
+    LLAMA_CPP_NO_MMPROJ_OFFLOAD: bool = False
+    LLAMA_CPP_EXTRA_ARGS: Optional[str] = None
+
+    # ---- Detection (kept) ---------------------------------------------------
+    DETECTOR_BATCH_SIZE: Optional[int] = None
+    DETECTOR_MODEL_CHECKPOINT: str = "s3://text_detection/2025_05_07"
+    DETECTOR_IMAGE_CHUNK_HEIGHT: int = 1400
+    DETECTOR_TEXT_THRESHOLD: float = 0.6
+    DETECTOR_BLANK_THRESHOLD: float = 0.35
+    DETECTOR_POSTPROCESSING_CPU_WORKERS: int = min(8, os.cpu_count())
+    DETECTOR_MIN_PARALLEL_THRESH: int = 3
+    DETECTOR_BOX_Y_EXPAND_MARGIN: float = 0.05
+
+    # ---- OCR Error (kept) ---------------------------------------------------
+    OCR_ERROR_MODEL_CHECKPOINT: str = "s3://ocr_error_detection/2025_02_18"
+    OCR_ERROR_BATCH_SIZE: Optional[int] = None
+
+    # ---- Debug / draw fonts (label rendering on annotated images) ----------
+    RECOGNITION_RENDER_FONTS: Dict[str, str] = {
+        "all": os.path.join(FONT_DIR, "GoNotoCurrent-Regular.ttf"),
+        "zh": os.path.join(FONT_DIR, "GoNotoCJKCore.ttf"),
+        "ja": os.path.join(FONT_DIR, "GoNotoCJKCore.ttf"),
+        "ko": os.path.join(FONT_DIR, "GoNotoCJKCore.ttf"),
+    }
+    RECOGNITION_FONT_DL_BASE: str = (
+        "https://github.com/satbyy/go-noto-universal/releases/download/v7.0"
+    )
+
+    @computed_field
+    def MODEL_DTYPE(self) -> torch.dtype:
+        if self.TORCH_DEVICE_MODEL == "cpu":
+            return torch.float32
+        return torch.float16
+
+    @computed_field
+    def MODEL_DTYPE_BFLOAT(self) -> torch.dtype:
+        if self.TORCH_DEVICE_MODEL == "cpu":
+            return torch.float32
+        return torch.bfloat16
+
+    @computed_field
+    def INFERENCE_MODE(self) -> Callable:
+        return torch.inference_mode
+
+    class Config:
+        env_file = find_dotenv("local.env")
+        extra = "ignore"
+
+
+settings = Settings()