Fu DaiandClaude Opus 4.8
1a585693be
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>
166 lines
7.0 KiB
Python
166 lines
7.0 KiB
Python
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()
|