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>
51 lines
2.2 KiB
Python
51 lines
2.2 KiB
Python
"""Per-method quantization specs and vLLM serve-argument construction.
|
|
|
|
kind:
|
|
baseline - serve the unquantized model as-is (the accuracy reference)
|
|
online - vLLM quantizes at load (fp8 dynamic); serve the base model
|
|
compressor - llm-compressor produced a checkpoint; quant config travels with it
|
|
bnb - transformers+bitsandbytes produced a checkpoint; serve with bnb flags
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from typing import Dict, List
|
|
|
|
METHOD_SPECS: Dict[str, Dict] = {
|
|
"bf16": {"kind": "baseline", "t4_deployable": True},
|
|
"fp8": {"kind": "online", "vllm_quant": "fp8", "t4_deployable": False},
|
|
# smoothquant=False: SmoothQuant's default mappings fail to resolve for the
|
|
# qwen3_5 layer layout (each mapping matches all 24 input_layernorms), so
|
|
# int8 is plain GPTQ W8A8.
|
|
"int8": {"kind": "compressor", "modifier": "gptq", "scheme": "W8A8", "smoothquant": False, "t4_deployable": True},
|
|
"awq": {"kind": "compressor", "modifier": "awq", "scheme": "W4A16", "smoothquant": False, "t4_deployable": True},
|
|
"gptq": {"kind": "compressor", "modifier": "gptq", "scheme": "W4A16", "smoothquant": False, "t4_deployable": True},
|
|
"bnb8": {"kind": "bnb", "bits": 8, "t4_deployable": True},
|
|
"bnb4": {"kind": "bnb", "bits": 4, "t4_deployable": True},
|
|
}
|
|
|
|
|
|
def method_names() -> List[str]:
|
|
return list(METHOD_SPECS.keys())
|
|
|
|
|
|
def vllm_serve_args(method: str, model_path: str, base_model: str, port: int) -> List[str]:
|
|
spec = METHOD_SPECS[method]
|
|
kind = spec["kind"]
|
|
serve_model = base_model if kind in ("baseline", "online") else model_path
|
|
args = [
|
|
"--host", "127.0.0.1",
|
|
"--port", str(port),
|
|
"--model", serve_model,
|
|
"--served-model-name", "datalab-to/surya-ocr-2",
|
|
"--max-model-len", "18000",
|
|
"--max-num-seqs", "16",
|
|
"--gpu-memory-utilization", "0.85",
|
|
"--enable-prefix-caching",
|
|
"--mm-processor-kwargs", '{"min_pixels":3136,"max_pixels":6291456}',
|
|
]
|
|
if kind == "online":
|
|
args += ["--quantization", spec["vllm_quant"]]
|
|
elif kind == "bnb":
|
|
args += ["--quantization", "bitsandbytes", "--load-format", "bitsandbytes"]
|
|
return args
|