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>
This commit is contained in:
Fu Dai
@@ -0,0 +1,207 @@
|
||||
"""llama.cpp backend: spawns the upstream `llama-server` binary natively.
|
||||
|
||||
Install:
|
||||
- macOS: brew install llama.cpp (Metal build, MPS)
|
||||
- Linux: brew install llama.cpp OR github.com/ggml-org/llama.cpp/releases
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
from openai import OpenAI
|
||||
|
||||
from surya.inference.backends.base import Backend, ServerHandle
|
||||
from surya.inference.backends.openai_client import chat_completions_batch
|
||||
from surya.inference.backends.spawn import (
|
||||
SpawnHandle,
|
||||
SpawnError,
|
||||
attach_or_spawn,
|
||||
)
|
||||
from surya.inference.schema import BatchInputItem, BatchOutputItem
|
||||
from surya.logging import get_logger
|
||||
from surya.settings import settings
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def _resolve_llama_server_binary() -> str:
|
||||
binary = settings.LLAMA_CPP_BINARY
|
||||
if binary and os.path.isfile(binary):
|
||||
return binary
|
||||
found = shutil.which(binary or "llama-server")
|
||||
if found:
|
||||
return found
|
||||
raise SpawnError(
|
||||
"llama-server binary not found. Install with:\n"
|
||||
" macOS: brew install llama.cpp\n"
|
||||
" Linux: brew install llama.cpp OR download from\n"
|
||||
" https://github.com/ggml-org/llama.cpp/releases\n"
|
||||
"Or set LLAMA_CPP_BINARY in your env to the binary path."
|
||||
)
|
||||
|
||||
|
||||
def _download_gguf_files() -> tuple[str, str]:
|
||||
"""Download model + mmproj GGUFs from HF Hub. Returns local paths."""
|
||||
repo = settings.SURYA_GGUF_REPO
|
||||
model_file = settings.SURYA_GGUF_MODEL_FILE
|
||||
mmproj_file = settings.SURYA_GGUF_MMPROJ_FILE
|
||||
logger.info(f"Downloading {model_file} and {mmproj_file} from {repo}")
|
||||
model_path = hf_hub_download(repo_id=repo, filename=model_file)
|
||||
mmproj_path = hf_hub_download(repo_id=repo, filename=mmproj_file)
|
||||
return model_path, mmproj_path
|
||||
|
||||
|
||||
def _health_url(port: int) -> str:
|
||||
return f"http://{settings.SURYA_INFERENCE_HOST}:{port}"
|
||||
|
||||
|
||||
def _openai_url(port: int) -> str:
|
||||
return f"http://{settings.SURYA_INFERENCE_HOST}:{port}/v1"
|
||||
|
||||
|
||||
class LlamaCppBackend(Backend):
|
||||
name = "llamacpp"
|
||||
|
||||
def __init__(self):
|
||||
self.handle: Optional[ServerHandle] = None
|
||||
self._client: Optional[OpenAI] = None
|
||||
|
||||
def start(self) -> ServerHandle:
|
||||
if self.handle is not None:
|
||||
return self.handle
|
||||
|
||||
# If user pinned an external server, attach without spawning.
|
||||
# No binary or GGUF download needed in that case.
|
||||
if settings.SURYA_INFERENCE_URL:
|
||||
spawned = attach_or_spawn(
|
||||
backend=self.name,
|
||||
expected_model_name=settings.SURYA_MODEL_CHECKPOINT,
|
||||
spawn_fn=lambda port: SpawnHandle(
|
||||
pid=None, cleanup_id="", cleanup_kind="process"
|
||||
), # never called
|
||||
health_url_for=_health_url,
|
||||
openai_url_for=_openai_url,
|
||||
startup_timeout=settings.SURYA_INFERENCE_STARTUP_TIMEOUT,
|
||||
)
|
||||
self.handle = ServerHandle(
|
||||
base_url=spawned.base_url,
|
||||
model_name=spawned.model_name,
|
||||
spawned_by_us=spawned.spawned_by_us,
|
||||
)
|
||||
self._client = OpenAI(api_key="EMPTY", base_url=self.handle.base_url)
|
||||
return self.handle
|
||||
|
||||
binary = _resolve_llama_server_binary()
|
||||
|
||||
# Pre-download GGUFs so the spawn doesn't race the download
|
||||
if (
|
||||
settings.SURYA_GGUF_LOCAL_MODEL_PATH
|
||||
and settings.SURYA_GGUF_LOCAL_MMPROJ_PATH
|
||||
):
|
||||
model_path = settings.SURYA_GGUF_LOCAL_MODEL_PATH
|
||||
mmproj_path = settings.SURYA_GGUF_LOCAL_MMPROJ_PATH
|
||||
else:
|
||||
model_path, mmproj_path = _download_gguf_files()
|
||||
|
||||
# Total KV-cache budget. llama-server divides --ctx-size across
|
||||
# --parallel slots, so a too-small total silently truncates outputs
|
||||
# once each slot's share fills. Scale with parallel by default;
|
||||
# SURYA_INFERENCE_CTX_SIZE overrides to a fixed value if set.
|
||||
parallel = settings.SURYA_INFERENCE_PARALLEL
|
||||
per_slot = settings.SURYA_INFERENCE_CTX_PER_SLOT
|
||||
ctx_size = settings.SURYA_INFERENCE_CTX_SIZE
|
||||
if ctx_size is None:
|
||||
ctx_size = max(16384, parallel * per_slot)
|
||||
effective_per_slot = ctx_size // max(parallel, 1)
|
||||
logger.info(
|
||||
f"llama-server ctx-size={ctx_size} "
|
||||
f"(~{effective_per_slot}/slot × {parallel} parallel slots)"
|
||||
)
|
||||
if effective_per_slot < per_slot:
|
||||
logger.warning(
|
||||
f"per-slot ctx ({effective_per_slot}) is below recommended "
|
||||
f"{per_slot}; outputs may truncate. Raise "
|
||||
f"SURYA_INFERENCE_CTX_SIZE or SURYA_INFERENCE_CTX_PER_SLOT, "
|
||||
f"or lower SURYA_INFERENCE_PARALLEL."
|
||||
)
|
||||
|
||||
def spawn_fn(port: int) -> SpawnHandle:
|
||||
cmd = [
|
||||
binary,
|
||||
"-m",
|
||||
model_path,
|
||||
"--mmproj",
|
||||
mmproj_path,
|
||||
"-ngl",
|
||||
str(settings.LLAMA_CPP_NGL),
|
||||
"--host",
|
||||
settings.SURYA_INFERENCE_HOST,
|
||||
"--port",
|
||||
str(port),
|
||||
"--parallel",
|
||||
str(parallel),
|
||||
"--ctx-size",
|
||||
str(ctx_size),
|
||||
"--no-mmproj-offload" if settings.LLAMA_CPP_NO_MMPROJ_OFFLOAD else "",
|
||||
"--alias",
|
||||
settings.SURYA_MODEL_CHECKPOINT,
|
||||
"--jinja",
|
||||
]
|
||||
cmd = [c for c in cmd if c]
|
||||
for extra in (settings.LLAMA_CPP_EXTRA_ARGS or "").split():
|
||||
cmd.append(extra)
|
||||
logger.info(f"Spawning: {' '.join(cmd)}")
|
||||
log_path = Path("~/.cache/datalab/surya/llamacpp_server.log").expanduser()
|
||||
log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
log_fp = open(log_path, "ab")
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=log_fp,
|
||||
stderr=subprocess.STDOUT,
|
||||
start_new_session=True,
|
||||
)
|
||||
return SpawnHandle(
|
||||
pid=proc.pid, cleanup_id=str(proc.pid), cleanup_kind="process"
|
||||
)
|
||||
|
||||
spawned = attach_or_spawn(
|
||||
backend=self.name,
|
||||
expected_model_name=settings.SURYA_MODEL_CHECKPOINT,
|
||||
spawn_fn=spawn_fn,
|
||||
health_url_for=_health_url,
|
||||
openai_url_for=_openai_url,
|
||||
startup_timeout=settings.SURYA_INFERENCE_STARTUP_TIMEOUT,
|
||||
)
|
||||
self.handle = ServerHandle(
|
||||
base_url=spawned.base_url,
|
||||
model_name=spawned.model_name,
|
||||
spawned_by_us=spawned.spawned_by_us,
|
||||
)
|
||||
self._client = OpenAI(
|
||||
api_key="EMPTY",
|
||||
base_url=self.handle.base_url,
|
||||
)
|
||||
return self.handle
|
||||
|
||||
def stop(self) -> None:
|
||||
# atexit handler in spawn.py owns cleanup; nothing to do here.
|
||||
self.handle = None
|
||||
self._client = None
|
||||
|
||||
def generate(self, batch: List[BatchInputItem]) -> List[BatchOutputItem]:
|
||||
if self.handle is None or self._client is None:
|
||||
self.start()
|
||||
return chat_completions_batch(
|
||||
batch,
|
||||
client=self._client,
|
||||
model_name=self.handle.model_name,
|
||||
timeout=settings.SURYA_INFERENCE_TIMEOUT_SECONDS,
|
||||
max_workers=settings.SURYA_INFERENCE_PARALLEL,
|
||||
request_logprobs_default=settings.SURYA_INFERENCE_LOGPROBS,
|
||||
)
|
||||
Reference in New Issue
Block a user