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,25 @@
|
||||
# Keep the build context small — none of this is needed at runtime.
|
||||
.git
|
||||
.venv-quant
|
||||
.venv
|
||||
venv
|
||||
*.venv
|
||||
.pytest_cache
|
||||
__pycache__
|
||||
**/__pycache__
|
||||
*.pyc
|
||||
results
|
||||
baseline_outputs
|
||||
eval_set
|
||||
docs
|
||||
tests
|
||||
logs/*.log
|
||||
concurrency_*.json
|
||||
concurrency_*.csv
|
||||
concurrency_*.png
|
||||
concurrency_sweep.py
|
||||
concurrency_test.py
|
||||
*.md
|
||||
temp_image_*.png
|
||||
temp_image_*.jpg
|
||||
temp_image_*.jpeg
|
||||
+27
@@ -0,0 +1,27 @@
|
||||
# Python bytecode
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
|
||||
# Logs
|
||||
logs/
|
||||
*.log
|
||||
|
||||
# Virtualenvs, caches, agent tooling
|
||||
.venv*/
|
||||
.pytest_cache/
|
||||
.agents/
|
||||
.codex/
|
||||
|
||||
# Benchmark / eval / quantization artifacts (kept on disk, never in git)
|
||||
results/
|
||||
baseline_outputs/
|
||||
eval_set/
|
||||
|
||||
# Local docs / working notes
|
||||
docs/
|
||||
|
||||
# Scratch OCR inputs/outputs
|
||||
temp_image_*
|
||||
output.json
|
||||
compare_out/
|
||||
*.pdf
|
||||
@@ -0,0 +1,125 @@
|
||||
# Repository Instructions
|
||||
|
||||
## Project Context
|
||||
|
||||
This repo serves a FastAPI OCR API from `api.py`. The API keeps the legacy
|
||||
request/response shape:
|
||||
|
||||
```json
|
||||
{"data": ..., "message": "success", "code": 200}
|
||||
```
|
||||
|
||||
Input validation is handled with Pydantic in `Info`. Keep external payload
|
||||
fields stable unless the user explicitly asks for a breaking API change.
|
||||
|
||||
Primary endpoints:
|
||||
|
||||
- `POST /v1/api/ai/suya_ocr/`: legacy OCR route.
|
||||
- `POST /v1/api/ai/suya_ocr_vllm/`: vLLM-backed OCR route.
|
||||
- `GET /v1/api/ai/suya_ocr_vllm/health`: reports vLLM backend attachment.
|
||||
- `POST /v1/api/ai/suya_text_det/`, `/suya_layout_det/`, `/suya_table_rec/`.
|
||||
|
||||
## Local Run
|
||||
|
||||
Use the conda Python on this host; Homebrew `python3` may not have the API
|
||||
dependencies installed. For local direct `api.py` runs, start a vLLM OpenAI
|
||||
server separately and point the API at it.
|
||||
|
||||
```bash
|
||||
PYTHONDONTWRITEBYTECODE=1 \
|
||||
PYTHONPATH=/path/to/suya-ocr-api \
|
||||
SURYA_INFERENCE_BACKEND=vllm \
|
||||
SURYA_INFERENCE_URL=http://127.0.0.1:8000/v1 \
|
||||
SURYA_INFERENCE_AUTOSTART=false \
|
||||
SURYA_INFERENCE_PARALLEL=8 \
|
||||
SURYA_MAX_TOKENS_FULL_PAGE=6144 \
|
||||
/opt/conda/bin/python api.py
|
||||
```
|
||||
|
||||
The API listens on `http://127.0.0.1:5002`.
|
||||
|
||||
For Kubernetes/container deployment, use `Dockerfile`. It is based on
|
||||
`vllm/vllm-openai:v0.20.1` and runs vLLM in the background, then `api.py` in
|
||||
the foreground through `scripts/start_single_container.sh`.
|
||||
|
||||
## vLLM Backend Notes
|
||||
|
||||
The repo vendors the Surya source under `surya/`; runtime must not depend on
|
||||
`/path/to/surya`. `vllm_tools.py` must set Surya/vLLM environment
|
||||
defaults before importing `surya.settings` or Surya predictor classes.
|
||||
Important defaults:
|
||||
|
||||
- `SURYA_INFERENCE_BACKEND=vllm`
|
||||
- `SURYA_INFERENCE_URL=http://127.0.0.1:8000/v1`
|
||||
- `SURYA_INFERENCE_AUTOSTART=false`
|
||||
- `SURYA_INFERENCE_LOGPROBS=false`
|
||||
- `SURYA_INFERENCE_MAX_RETRIES=1`
|
||||
- `SURYA_MAX_TOKENS_FULL_PAGE=6144`
|
||||
- `SURYA_MAX_BLOCKS_PER_PAGE=80`
|
||||
- `SUYA_OCR_MODE=block`
|
||||
|
||||
The Dockerfile is tuned conservatively for a T4 machine:
|
||||
|
||||
- `VLLM_GPU_TYPE=t4`
|
||||
- `VLLM_DTYPE=float16`
|
||||
- `SURYA_INFERENCE_PARALLEL=8`
|
||||
- `VLLM_MAX_NUM_SEQS=16`
|
||||
- `VLLM_MAX_BATCHED_TOKENS=4096`
|
||||
- `VLLM_GPU_MEMORY_UTILIZATION=0.85`
|
||||
|
||||
If running on A100 hardware, do not silently change the Dockerfile T4 defaults;
|
||||
ask or document the deployment target first.
|
||||
|
||||
## Docker/NVIDIA
|
||||
|
||||
The production container must not run nested Docker. Kubernetes should provide
|
||||
GPU access to the single container. The old Surya Docker auto-spawn path is
|
||||
disabled unless `SUYA_ALLOW_NESTED_DOCKER=true` is explicitly set.
|
||||
|
||||
For host-level debugging only, Docker GPU runtime can be checked with:
|
||||
|
||||
```bash
|
||||
docker info --format '{{json .Runtimes}}'
|
||||
docker run --rm --runtime nvidia --gpus all --entrypoint nvidia-smi vllm/vllm-openai:v0.20.1 -L
|
||||
```
|
||||
|
||||
If the named runtime is missing, use:
|
||||
|
||||
```bash
|
||||
sudo nvidia-ctk runtime configure --runtime=docker --config=/etc/docker/daemon.json
|
||||
sudo systemctl restart docker
|
||||
```
|
||||
|
||||
`scripts/docker_nvidia_runtime_compat.sh` is legacy support for local debugging,
|
||||
not the Kubernetes deployment path.
|
||||
|
||||
## Concurrency Benchmark
|
||||
|
||||
Run the comparison benchmark with the same image and settings when updating
|
||||
results:
|
||||
|
||||
```bash
|
||||
/opt/conda/bin/python concurrency_test.py \
|
||||
--image /path/to/surya/static/images/excerpt.png \
|
||||
--requests 16 \
|
||||
--concurrency 16 \
|
||||
--timeout 900 \
|
||||
--output concurrency_test_results.json
|
||||
```
|
||||
|
||||
After a successful run, update:
|
||||
|
||||
- `concurrency_test_results.json`
|
||||
- `IMPLEMENTATION_PLAN_AND_TEST_RESULTS.md`
|
||||
|
||||
Report latency columns as seconds. Mean latency is the average elapsed time per
|
||||
request, not wall-clock time for the whole benchmark.
|
||||
|
||||
## Editing Guidance
|
||||
|
||||
- Use `rg`/`rg --files` for search.
|
||||
- Use `apply_patch` for manual edits.
|
||||
- Do not revert user changes or generated benchmark artifacts unless requested.
|
||||
- Keep logging structured and include request IDs for API paths.
|
||||
- Avoid import-time OCR work; model or backend attachment should happen only as
|
||||
needed by predictor setup and requests.
|
||||
+66
@@ -0,0 +1,66 @@
|
||||
FROM vllm/vllm-openai:v0.20.1
|
||||
|
||||
ARG DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
RUN apt-get update && \
|
||||
apt-get install --no-install-recommends -y \
|
||||
ffmpeg \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libcairo2 \
|
||||
libgirepository1.0-dev \
|
||||
libdbus-1-3 \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
ENV PYTHONPATH=/opt/suya-ocr
|
||||
|
||||
ENV SURYA_MODEL_CHECKPOINT=datalab-to/surya-ocr-2
|
||||
ENV SURYA_INFERENCE_BACKEND=vllm
|
||||
ENV SURYA_INFERENCE_URL=http://127.0.0.1:8000/v1
|
||||
ENV SURYA_INFERENCE_AUTOSTART=false
|
||||
ENV SURYA_INFERENCE_KEEP_ALIVE=false
|
||||
ENV SURYA_INFERENCE_PARALLEL=8
|
||||
ENV SURYA_INFERENCE_LOGPROBS=false
|
||||
ENV SURYA_INFERENCE_MAX_RETRIES=1
|
||||
# Cap concurrent in-flight chat-completion requests to vLLM per coalesced batch.
|
||||
# Must track VLLM_MAX_NUM_SEQS below: the request batcher otherwise ran only
|
||||
# SURYA_INFERENCE_PARALLEL(=8) wide, leaving half of vLLM's 16 sequence slots
|
||||
# idle. Setting this to 16 closed a 33% mean / 32% p95 regression (see
|
||||
# docs/diagnosis_baseline.md). Do not exceed VLLM_MAX_NUM_SEQS (overcommit only
|
||||
# queues).
|
||||
ENV SURYA_INFERENCE_MAX_INFLIGHT=16
|
||||
|
||||
ENV SUYA_OCR_MODE=block
|
||||
ENV SUYA_MAX_BATCH_SIZE=8
|
||||
ENV SUYA_BATCH_WAIT_MS=25
|
||||
ENV SUYA_MAX_QUEUE_SIZE=128
|
||||
ENV SUYA_VLLM_IMAGE_FORMAT=JPEG
|
||||
ENV SUYA_VLLM_JPEG_QUALITY=92
|
||||
ENV SURYA_MAX_BLOCKS_PER_PAGE=80
|
||||
ENV SURYA_MAX_TOKENS_FULL_PAGE=6144
|
||||
|
||||
# GPU/dtype defaults target T4 (deployment hardware). Do not change these
|
||||
# without confirming the deployment target.
|
||||
ENV VLLM_GPU_TYPE=t4
|
||||
ENV VLLM_DTYPE=float16
|
||||
ENV VLLM_GPU_MEMORY_UTILIZATION=0.85
|
||||
ENV VLLM_MAX_MODEL_LEN=18000
|
||||
ENV VLLM_MAX_NUM_SEQS=16
|
||||
ENV VLLM_MAX_BATCHED_TOKENS=4096
|
||||
# Latency-tuned vLLM serve flags (CUDA graph, prefix caching, chunked prefill)
|
||||
# are applied in scripts/start_single_container.sh — see that file and
|
||||
# docs/quantization_benchmark_results.md §5 for the measured rationale.
|
||||
|
||||
WORKDIR /opt/suya-ocr
|
||||
|
||||
COPY requirements.txt pyproject.toml ./
|
||||
RUN pip install --no-cache-dir uv==0.8.15 && \
|
||||
python3 -m uv pip install --system -r requirements.txt
|
||||
|
||||
COPY . .
|
||||
|
||||
EXPOSE 5002
|
||||
|
||||
ENTRYPOINT []
|
||||
CMD ["/opt/suya-ocr/scripts/start_single_container.sh"]
|
||||
@@ -0,0 +1,182 @@
|
||||
# vLLM Backend Plan and Test Results
|
||||
|
||||
Date: 2026-06-09
|
||||
|
||||
## Plan
|
||||
|
||||
1. Keep the legacy OCR endpoint at `/v1/api/ai/suya_ocr/` for baseline comparison.
|
||||
2. Add `/v1/api/ai/suya_ocr_vllm/`, backed by Surya2 `SuryaInferenceManager(method="vllm")`.
|
||||
3. Use the Surya vLLM backend settings from the referenced Surya checkout:
|
||||
- `SURYA_INFERENCE_BACKEND=vllm`
|
||||
- `SURYA_INFERENCE_KEEP_ALIVE=true`
|
||||
- T4 default: `VLLM_GPU_TYPE=t4`, `VLLM_DTYPE=float16`, `SURYA_INFERENCE_PARALLEL=16`
|
||||
- optional tuning through `VLLM_GPUS`, `VLLM_GPU_MEMORY_UTILIZATION`, and `SURYA_INFERENCE_URL`
|
||||
4. Add a concurrency benchmark that sends the same payload to the legacy and vLLM endpoints and writes latency/throughput JSON.
|
||||
5. Verify syntax locally, then run the benchmark on a host with the Surya2 package, Docker, and NVIDIA GPU access.
|
||||
|
||||
## Implemented Files
|
||||
|
||||
- `vllm_tools.py`: vLLM-backed Surya2 predictor setup and OCR response formatting.
|
||||
- `api.py`: new `/v1/api/ai/suya_ocr_vllm/` endpoint and health endpoint.
|
||||
- `concurrency_test.py`: repeatable old-vs-new endpoint concurrency benchmark.
|
||||
- `concurrency_test_results.json`: placeholder result file that the benchmark overwrites with measured results.
|
||||
- `requirements.txt`: dependency alignment for Surya2/vLLM OpenAI client support.
|
||||
- `Dockerfile`: T4 vLLM defaults. The referenced backend maps T4 16 GB VRAM to `max_num_seqs=16` and `max_num_batched_tokens=4096`; `float16` is required because T4 does not support `bfloat16`.
|
||||
|
||||
## Single-Container Kubernetes Deployment
|
||||
|
||||
The deployment now uses `vllm/vllm-openai:v0.20.1` as the base image. The container starts the vLLM OpenAI-compatible server in the background and runs `api.py` in the foreground through `scripts/start_single_container.sh`.
|
||||
|
||||
The API no longer depends on nested Docker. It attaches to:
|
||||
|
||||
```text
|
||||
SURYA_INFERENCE_URL=http://127.0.0.1:8000/v1
|
||||
```
|
||||
|
||||
The Surya source needed by the API is vendored into this repo under `surya/`, so the container does not depend on `/path/to/surya` at runtime.
|
||||
|
||||
Serving defaults for T4:
|
||||
|
||||
```text
|
||||
SUYA_OCR_MODE=block
|
||||
SUYA_MAX_BATCH_SIZE=8
|
||||
SUYA_BATCH_WAIT_MS=25
|
||||
SURYA_INFERENCE_PARALLEL=8
|
||||
SURYA_INFERENCE_MAX_RETRIES=1
|
||||
SURYA_MAX_BLOCKS_PER_PAGE=80
|
||||
VLLM_DTYPE=float16
|
||||
VLLM_MAX_NUM_SEQS=16
|
||||
VLLM_MAX_BATCHED_TOKENS=4096
|
||||
```
|
||||
|
||||
Speed-related changes:
|
||||
|
||||
- `/v1/api/ai/suya_ocr_vllm/` uses `vllm_batcher.py` to coalesce concurrent HTTP requests into batch Surya calls.
|
||||
- `vllm_tools.py` batches layout and recognition across images.
|
||||
- Default OCR mode is block mode, avoiding long full-page generations and hidden fallback bursts.
|
||||
- Vendored Surya caps block OCR per page with `SURYA_MAX_BLOCKS_PER_PAGE`.
|
||||
- Retry amplification is limited with `SURYA_INFERENCE_MAX_RETRIES=1`.
|
||||
- The endpoint skips drawing annotated images because the API response only returns JSON/text.
|
||||
- Base64 input is decoded once per request and the same PIL image object is reused for low/high-res OCR paths unless a caller explicitly asks for a copy.
|
||||
- vLLM image payload encoding defaults to high-quality JPEG (`SUYA_VLLM_IMAGE_FORMAT=JPEG`, `SUYA_VLLM_JPEG_QUALITY=92`) to reduce CPU serialization and local vLLM transfer size versus PNG. Set `SUYA_VLLM_IMAGE_FORMAT=PNG` if lossless transport is required.
|
||||
|
||||
## Benchmark Command
|
||||
|
||||
```bash
|
||||
python concurrency_test.py \
|
||||
--image /path/to/surya/ttt.png \
|
||||
--requests 20 \
|
||||
--concurrency 8 \
|
||||
--output concurrency_test_results.json
|
||||
```
|
||||
|
||||
## Current Test Results
|
||||
|
||||
Import verification passed with `PYTHONPATH=/path/to/surya`.
|
||||
|
||||
Concurrency benchmark run:
|
||||
|
||||
```bash
|
||||
/opt/conda/bin/python concurrency_test.py \
|
||||
--image temp_image_c3c56948-282e-453b-8fac-6c482243d1e5.jpg \
|
||||
--requests 16 \
|
||||
--concurrency 16 \
|
||||
--timeout 900 \
|
||||
--output concurrency_test_results.json
|
||||
```
|
||||
|
||||
Run started at `2026-06-11T14:35:17+0400`. Docker NVIDIA runtime and GPU checks passed. The local Python environment does not include the `vllm` package, so this validation attached the updated API to the already-warm local `vllm/vllm-openai:v0.20.1` server on port `43401` with `SURYA_INFERENCE_URL=http://127.0.0.1:43401/v1`. The API path still used the single-container attach mode (`SURYA_INFERENCE_AUTOSTART=false`) and did not start nested Docker.
|
||||
|
||||
| Endpoint | Requests | Concurrency | Success | Failed | Wall seconds | Throughput rps | Mean latency | P50 latency | P95 latency | Max latency |
|
||||
| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
|
||||
| `/v1/api/ai/suya_ocr/` | 16 | 16 | 16 | 0 | 43.62 | 0.3668 | 40.81 | 41.84 | 43.39 | 43.62 |
|
||||
| `/v1/api/ai/suya_ocr_vllm/` | 16 | 16 | 16 | 0 | 83.59 | 0.1914 | 59.08 | 60.41 | 83.52 | 83.54 |
|
||||
|
||||
Latency values are per request in seconds. Mean latency is the average request elapsed time, not the total benchmark wall time.
|
||||
|
||||
## Approach A Outcome (2026-06-12)
|
||||
|
||||
Diagnosis + work-reduction tuning (see `docs/diagnosis_baseline.md` and
|
||||
`docs/tier2_results.md`). Measured on the A100 dev host against the warm
|
||||
`surya-vllm-43401` server — same host/server as the table above, so comparable.
|
||||
**Production is T4; the fix below is hardware-independent, the absolute numbers
|
||||
are not.**
|
||||
|
||||
**Root cause:** the request batcher coalesced requests then dispatched them only
|
||||
`SURYA_INFERENCE_PARALLEL=8` wide, so only 8 of vLLM's 16 sequence slots were
|
||||
used (poller `peak_running=8`). The GPU sat half-idle, which is why the vLLM
|
||||
endpoint was *slower* than the flooding legacy path.
|
||||
|
||||
**Fix (Tier 1):** scale the worker pool to the in-flight block count, capped by a
|
||||
new `SURYA_INFERENCE_MAX_INFLIGHT` (default 16, = `VLLM_MAX_NUM_SEQS`). This fills
|
||||
all 16 slots (`peak_running=16`).
|
||||
|
||||
| vLLM endpoint, 16/16 | mean | p50 | p95 | peak_running |
|
||||
| --- | ---: | ---: | ---: | ---: |
|
||||
| MI=8 (before) | 50.22 | 49.48 | 72.69 | 8 |
|
||||
| MI=16 (fix) | 33.75 | 36.11 | 49.15 | 16 |
|
||||
|
||||
- Mean −33%, p95 −32%; single request 6.16 s → ~5.3 s. CER divergence 0.0009
|
||||
(lossless; gate ≤ 0.005). The vLLM endpoint now beats legacy on mean and median.
|
||||
- MI=24 overcommits (worse). Batch-size / batch-wait are within noise.
|
||||
|
||||
**Tier 2 (token ceiling, recognition DPI, block cap):** no latency benefit on
|
||||
this sparse page — decode dominates and the page is light — so all Tier-2 knobs
|
||||
are left at defaults. They remain relevant for dense / pathological pages.
|
||||
|
||||
**Locked change:** `SURYA_INFERENCE_MAX_INFLIGHT=16` added to `Dockerfile` and
|
||||
`scripts/start_single_container.sh`. Re-confirm absolute latencies on a T4 before
|
||||
treating them as production SLAs; the regression fix itself applies unchanged.
|
||||
|
||||
## FP8 Experiment
|
||||
|
||||
I attempted to start the latest vLLM image with multimodal encoder FP8 enabled:
|
||||
|
||||
```bash
|
||||
--dtype bfloat16 \
|
||||
--mm-encoder-attn-dtype fp8
|
||||
```
|
||||
|
||||
The first attempt on `cuda:0` failed before serving traffic because GPU 0 did not have enough free memory for the default `--gpu-memory-utilization 0.85` target.
|
||||
|
||||
The second attempt on `cuda:1` booted farther, but vLLM aborted during model initialization with:
|
||||
|
||||
```text
|
||||
ValueError: mm_encoder_attn_dtype='fp8' requires the FlashInfer cuDNN backend with cuDNN >= 9.17.1 on a GPU with native FP8 support.
|
||||
```
|
||||
|
||||
That means:
|
||||
|
||||
- no FP8 latency number was produced
|
||||
- the current A100 hardware is not a valid target for this exact FP8 path
|
||||
- the failure is architectural, not a tuning issue
|
||||
|
||||
To measure a real FP8 speed change, the next engineer should use a GPU with native FP8 support and a vLLM build that satisfies the FlashInfer cuDNN requirement for multimodal encoder FP8.
|
||||
|
||||
## NVIDIA Docker Fix
|
||||
|
||||
Docker currently has no registered `nvidia` runtime. The expected daemon config is:
|
||||
|
||||
```json
|
||||
{
|
||||
"data-root": "/data/docker_cache",
|
||||
"runtimes": {
|
||||
"nvidia": {
|
||||
"args": [],
|
||||
"path": "nvidia-container-runtime"
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
This user cannot write `/etc/docker/daemon.json` or restart Docker because `sudo` requires a password. The host does support Docker GPU passthrough through modern `--gpus all`, so the API now patches Surya's vLLM spawn path to use `scripts/docker_nvidia_runtime_compat.sh`. That wrapper strips the legacy `--runtime nvidia` pair and preserves `--gpus device=...`.
|
||||
|
||||
An admin can still register the named runtime permanently with:
|
||||
|
||||
```bash
|
||||
sudo nvidia-ctk runtime configure --runtime=docker --config=/etc/docker/daemon.json
|
||||
sudo systemctl restart docker
|
||||
docker run --rm --runtime nvidia --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi
|
||||
```
|
||||
|
||||
The repo includes `scripts/configure_nvidia_docker.sh` with those commands.
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2026 Fu Dai
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,308 @@
|
||||
# Suya OCR API
|
||||
|
||||
A FastAPI service that wraps the [Surya-OCR-2](https://huggingface.co/datalab-to/surya-ocr-2)
|
||||
document-OCR model and serves it through **vLLM** for production-grade throughput and latency.
|
||||
A single 650M vision-language model (`qwen3_5` architecture) performs both **layout detection**
|
||||
and **text recognition** — there is no separate torch detection model in the hot path.
|
||||
|
||||
```
|
||||
HTTP client ──► FastAPI app (surya/endpoint/, served via api.py shim :5002)
|
||||
├─ legacy.py /v1/api/ai/* envelope endpoints
|
||||
├─ openai.py /v1/chat/completions (OpenAI-compatible)
|
||||
├─ request batcher (vllm_batcher.py) ← coalesces concurrent requests
|
||||
│ ▼
|
||||
└─ vLLM OpenAI server (:8000) ← Surya-OCR-2 VLM
|
||||
```
|
||||
|
||||
> **Inbound OpenAI:** clients call `/v1/chat/completions`. **Outbound OpenAI:** the service itself calls the vLLM server via `surya/inference/backends/openai_client.py`.
|
||||
|
||||
---
|
||||
|
||||
## What was improved
|
||||
|
||||
This service started as a torch-spawn pipeline and was re-engineered for serving, with the
|
||||
measured results summarized below:
|
||||
|
||||
| Area | Change | Result |
|
||||
| --- | --- | --- |
|
||||
| **Serving engine** | Recognition + layout now run through a vLLM **OpenAI-compatible** server instead of in-process torch | Continuous batching, CUDA graphs, prefix caching |
|
||||
| **vLLM tuning** | CUDA graph on, prefix caching on, chunked prefill on (pinned in the launch script) | CUDA graph is the biggest lever — **~9.5× vs eager**; prefix cache ~15% |
|
||||
| **Concurrency fix** | Matched the batcher's HTTP fan-out (`SURYA_INFERENCE_MAX_INFLIGHT`) to vLLM's 16 sequence slots | Closed a **33% mean / 32% p95** regression |
|
||||
| **Request batching** | Added a coalescing batcher so concurrent page requests share vLLM batches | Fills all 16 slots under load |
|
||||
| **Quantization study** | Benchmarked fp8/int8/awq/gptq/bnb4/bnb8 vs BF16 on A100 | **No method beats BF16 on latency**; int8 (GPTQ W8A8) near-lossless at ~40% smaller |
|
||||
| **MTP speculative decode** | Tested `--speculative-config mtp` | **Rejected: +27–36% slower** at 16-wide block concurrency |
|
||||
|
||||
**Bottom line:** on A100 the latency levers are *serving config + batching*, not weight
|
||||
quantization or speculative decoding. Keep BF16 with CUDA graphs and prefix caching on.
|
||||
|
||||
---
|
||||
|
||||
## Examples — multilingual OCR
|
||||
|
||||
Real outputs from **this service** (block mode, single A100) on public-domain newspaper/magazine
|
||||
scans: the original page, the same page with detected text-block boxes (coloured by block type /
|
||||
reading order; orange = non-text region skipped), and an excerpt of the recognised text. Full
|
||||
text for each sample lives in [`assets/`](assets/).
|
||||
|
||||
| Language | Original | Detected boxes | OCR text (excerpt) |
|
||||
| --- | --- | --- | --- |
|
||||
| **Chinese**<br><sub>叻報, 1890</sub> | <img src="assets/chinese_original.jpg" width="180"> | <img src="assets/chinese_boxes.png" width="180"> | 大清光緒十六年<br>本館新聞除禮拜外日出一張<br>諸君賜閱本報者無論本埠外埠…<br>工務局告示 |
|
||||
| **English**<br><sub>The Nation, 1846</sub> | <img src="assets/english_original.png" width="180"> | <img src="assets/english_boxes.png" width="180"> | VOL. IV. No. 181.<br>DUBLIN, SATURDAY, MARCH 28, 1846.<br>PRICE 6 D.<br>DAVIS TESTIMONIAL. |
|
||||
| **Arabic**<br><sub>Al-Ahram, 1981</sub> | <img src="assets/arabic_original.jpg" width="180"> | <img src="assets/arabic_boxes.png" width="180"> | رئيس مجلس الادارة — عبد الله عبد البخاري<br>المنطقات المنتصرة<br>السنة ١٠٠٧ — العدد ٣٩٦٣٨ |
|
||||
| **Russian**<br><sub>Виттова Пляска, 1905</sub> | <img src="assets/russian_original.jpg" width="180"> | <img src="assets/russian_boxes.png" width="180"> | Виттова Пляска<br>ОДНОДНЕВНАЯ ГАЗЕТА. ПАЛИ—ТИКО—ФИ—НАНСОВАЯ<br>Цѣна 5 копѣекъ. |
|
||||
| **French**<br><sub>Le Miroir des Sports, 1937</sub> | <img src="assets/french_original.jpg" width="180"> | <img src="assets/french_boxes.png" width="180"> | LE MIROIR DES SPORTS<br>Le plus fort tirage des hebdomadaires sportifs<br>Mardi 10 Août 1937 |
|
||||
|
||||
> Excerpts are the model's **raw** output on hard, dense historical scans — ornate mastheads
|
||||
> (e.g. the gothic “THE NATION”) are the failure cases. Reproduce any row with, e.g.,
|
||||
> `python -m surya.scripts.ocr_text assets/russian_original.jpg`, or via the API/CLI examples below.
|
||||
|
||||
---
|
||||
|
||||
## API
|
||||
|
||||
FastAPI app served on port **5002**. Interactive docs: `http://<host>:5002/v1/api/ai/swagger`.
|
||||
|
||||
All OCR endpoints accept the same JSON body and return the same envelope.
|
||||
|
||||
**Request body:**
|
||||
```json
|
||||
{
|
||||
"file": "<base64-encoded image bytes>",
|
||||
"type": "png", // one of: png | jpg | jpeg | gif
|
||||
"skip_text_detection": false, // true ⇒ full-page OCR, skip layout
|
||||
"skip_table_detection": false, // table endpoint only
|
||||
"recognize_math": false,
|
||||
"ocr_with_boxes": true
|
||||
}
|
||||
```
|
||||
|
||||
**Response envelope:**
|
||||
```json
|
||||
{ "data": { "ocr_text_json": {...}, "text_lines": "line1\nline2\n...", "elapsed_seconds": 4.64 },
|
||||
"message": "success",
|
||||
"code": 200 }
|
||||
```
|
||||
|
||||
| Method | Path | Purpose | Backend |
|
||||
| --- | --- | --- | --- |
|
||||
| POST | `/v1/api/ai/suya_ocr_vllm/` | **Primary OCR** — layout + recognition, request-batched | vLLM |
|
||||
| GET | `/v1/api/ai/suya_ocr_vllm/health` | Backend info (model, base URL, mode, parallelism) | — |
|
||||
| POST | `/v1/api/ai/suya_ocr/` | OCR (routes to the same vLLM path) | vLLM |
|
||||
| POST | `/v1/api/ai/suya_layout_det/` | Layout detection only | vLLM |
|
||||
| POST | `/v1/api/ai/suya_table_rec/` | Table structure recognition | vLLM |
|
||||
| POST | `/image2text` | OCR a multipart file upload (no base64) | vLLM |
|
||||
|
||||
> The legacy `/v1/api/ai/suya_text_det/` endpoint targets the standalone torch detection model,
|
||||
> which is **not loaded** in the vLLM deployment (the VLM does detection internally). It will
|
||||
> return an error envelope — use `/suya_ocr_vllm/` instead.
|
||||
|
||||
### OpenAI-compatible endpoint (`POST /v1/chat/completions`)
|
||||
|
||||
Point any OpenAI client at `http://<host>:5002/v1`. Attach the page as an `image_url`
|
||||
content part; the joined OCR text comes back in `choices[0].message.content`, and the
|
||||
structured per-line JSON + timing come back in a non-standard `surya` field.
|
||||
|
||||
OCR options are passed through `extra_body` (the OpenAI SDK delivers these at the top
|
||||
level of the request, which is where the server reads them):
|
||||
|
||||
| Field | Default | Purpose |
|
||||
| --- | --- | --- |
|
||||
| `mode` | `block` | `block` (layout→per-block), `full_page` (one call), or `table` |
|
||||
| `skip_text_detection` | `false` | Full-page OCR, skip layout (same as `mode:full_page`) |
|
||||
| `recognize_math` | `false` | Math-aware recognition |
|
||||
| `skip_table_detection` | `false` | `mode:table` only — treat whole image as one table |
|
||||
| `ocr_with_boxes` | `true` | When `false`, omit `surya.ocr_text_json` (text only) |
|
||||
|
||||
Streaming (`stream:true`) is **not supported** and returns HTTP 400 — OCR completes
|
||||
all at once, so there is no token stream.
|
||||
|
||||
### OpenAI-compatible backend
|
||||
|
||||
Recognition and layout are issued as **OpenAI Chat Completions** calls to the vLLM server
|
||||
(`POST {SURYA_INFERENCE_URL}/chat/completions`, default `http://127.0.0.1:8000/v1`), one per
|
||||
image crop, with the page image passed as a base64 `image_url` content part. This means the
|
||||
inference tier is a standard vLLM OpenAI server — you can point `SURYA_INFERENCE_URL` at any
|
||||
OpenAI-compatible vLLM endpoint (local, remote, or a shared cluster).
|
||||
|
||||
---
|
||||
|
||||
## Concurrency support
|
||||
|
||||
Two cooperating layers keep the GPU saturated without overcommitting:
|
||||
|
||||
1. **Request batcher** (`vllm_batcher.py`): a background thread coalesces incoming page
|
||||
requests that share the same options into one batch — up to `SUYA_MAX_BATCH_SIZE` (8),
|
||||
waiting at most `SUYA_BATCH_WAIT_MS` (25 ms). Requests queue up to `SUYA_MAX_QUEUE_SIZE`
|
||||
(128); callers block on a future until their page completes.
|
||||
2. **vLLM sequence slots**: within a page, the detected blocks are fanned out over HTTP up to
|
||||
`SURYA_INFERENCE_MAX_INFLIGHT` (16) concurrent workers, matching vLLM's `--max-num-seqs` (16)
|
||||
so every sequence slot is filled. (Overcommitting past 16 only queues — see the diagnosis.)
|
||||
|
||||
**Measured behaviour** (single A100, `benchmarks/concurrency_sweep_results.csv`) — 1→50 concurrent
|
||||
clients, **0 failures** at every level. Throughput is GPU-bound and plateaus around
|
||||
**0.20 req/s** on one model replica:
|
||||
|
||||
| Concurrency | Success / Failed | Throughput (req/s) | Mean latency (s) | p95 latency (s) |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| 1 | 1 / 0 | 0.13 | 7.66 | 7.66 |
|
||||
| 4 | 4 / 0 | 0.19 | 20.95 | 20.96 |
|
||||
| 8 | 8 / 0 | 0.19 | 34.05 | 41.46 |
|
||||
| 12 | 12 / 0 | 0.20 | 45.55 | 59.96 |
|
||||
| 20 | 20 / 0 | 0.20 | 66.68 | 92.91 |
|
||||
| 30 | 30 / 0 | 0.20 | 90.50 | 147.47 |
|
||||
| 50 | 50 / 0 | 0.21 | 149.29 | 240.87 |
|
||||
|
||||
|
||||
|
||||
Scale throughput by adding GPU replicas behind a load balancer, each with its own vLLM server.
|
||||
|
||||
---
|
||||
|
||||
## How to use
|
||||
|
||||
### Run with Docker (recommended)
|
||||
|
||||
```bash
|
||||
docker build -t suya-ocr-api .
|
||||
docker run --rm --gpus all -p 5002:5002 suya-ocr-api
|
||||
```
|
||||
|
||||
The container entrypoint (`scripts/start_single_container.sh`) starts the vLLM server with the
|
||||
tuned flags, waits for it to become healthy, then launches the API on `:5002`.
|
||||
|
||||
> **Deployment target:** the Dockerfile defaults target **T4** (`VLLM_DTYPE=float16`,
|
||||
> `VLLM_GPU_TYPE=t4`). On bf16-capable GPUs (A100/L40/4090) set `VLLM_DTYPE=bfloat16`.
|
||||
|
||||
### API mode — call the HTTP service (curl)
|
||||
|
||||
```bash
|
||||
# encode an image and OCR it
|
||||
B64=$(base64 -w0 page.png)
|
||||
curl -s http://localhost:5002/v1/api/ai/suya_ocr_vllm/ \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d "{\"file\":\"$B64\",\"type\":\"png\",\"ocr_with_boxes\":true}" \
|
||||
| python3 -c "import sys,json; d=json.load(sys.stdin)['data']; print(d['text_lines']); print('elapsed:', d['elapsed_seconds'])"
|
||||
|
||||
# full-page OCR (skip layout) — good for single-column pages
|
||||
curl -s http://localhost:5002/v1/api/ai/suya_ocr_vllm/ \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d "{\"file\":\"$B64\",\"type\":\"png\",\"skip_text_detection\":true}"
|
||||
|
||||
# health / backend info
|
||||
curl -s http://localhost:5002/v1/api/ai/suya_ocr_vllm/health
|
||||
```
|
||||
|
||||
### Demo: OCR one image via the OpenAI API
|
||||
|
||||
**OpenAI Python SDK** (the headline "it's really OpenAI-compatible" path):
|
||||
|
||||
```python
|
||||
import base64
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:5002/v1", api_key="not-needed")
|
||||
|
||||
with open("page.png", "rb") as f:
|
||||
b64 = base64.b64encode(f.read()).decode()
|
||||
|
||||
resp = client.chat.completions.create(
|
||||
model="surya-ocr",
|
||||
messages=[{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}},
|
||||
{"type": "text", "text": "Extract the text"},
|
||||
],
|
||||
}],
|
||||
extra_body={"mode": "block", "ocr_with_boxes": True},
|
||||
)
|
||||
print(resp.choices[0].message.content) # plain OCR text
|
||||
print(resp.model_dump()["surya"]["ocr_text_json"]) # structured boxes
|
||||
```
|
||||
|
||||
**curl** (data-URL base64 image):
|
||||
|
||||
```bash
|
||||
B64=$(base64 -w0 page.png)
|
||||
curl -s http://localhost:5002/v1/chat/completions \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d "{\"model\":\"surya-ocr\",\"messages\":[{\"role\":\"user\",\"content\":[{\"type\":\"image_url\",\"image_url\":{\"url\":\"data:image/png;base64,$B64\"}},{\"type\":\"text\",\"text\":\"ocr\"}]}]}" \
|
||||
| python3 -c "import sys,json; print(json.load(sys.stdin)['choices'][0]['message']['content'])"
|
||||
```
|
||||
|
||||
### CLI mode — local OCR (no HTTP server)
|
||||
|
||||
Run OCR straight from the command line. The `surya` package spins up its own vLLM-backed
|
||||
inference backend, so the FastAPI server does **not** need to be running. Input can be a single
|
||||
image, a PDF, or a folder of them; results are written as `results.json`.
|
||||
|
||||
```bash
|
||||
# OCR an image or PDF → writes results/surya/<name>/results.json
|
||||
python -m surya.scripts.ocr_text page.png
|
||||
|
||||
# pick an output dir and OCR only some PDF pages (0-indexed; ranges allowed)
|
||||
python -m surya.scripts.ocr_text doc.pdf --output_dir out --page_range 0,2-4
|
||||
|
||||
# also save annotated bbox images, and print per-stage timing
|
||||
python -m surya.scripts.ocr_text page.png --images --debug
|
||||
```
|
||||
|
||||
By default the CLI **autostarts a vLLM server** (`SURYA_INFERENCE_AUTOSTART=true`) and tears it
|
||||
down on exit. To reuse a vLLM server that is already running (for example the one started by the
|
||||
Docker container or `scripts/start_single_container.sh`), point the CLI at it instead of spawning
|
||||
a new one:
|
||||
|
||||
```bash
|
||||
SURYA_INFERENCE_URL=http://127.0.0.1:8000/v1 python -m surya.scripts.ocr_text page.png
|
||||
```
|
||||
|
||||
Pass `--keep_server` to leave a spawned server up so back-to-back commands reuse it. Sibling
|
||||
commands share the same options (run with `--help`):
|
||||
`surya.scripts.detect_layout`, `surya.scripts.detect_text`, `surya.scripts.table_recognition`.
|
||||
|
||||
### Run without Docker
|
||||
|
||||
Requires a CUDA GPU with the matching vLLM. Start the two processes (vLLM first, then the API)
|
||||
exactly as `scripts/start_single_container.sh` does, or run that script directly inside an
|
||||
environment that already has the dependencies from `requirements.txt`.
|
||||
|
||||
---
|
||||
|
||||
## Configuration
|
||||
|
||||
Set via environment variables (defaults shown; see `Dockerfile` and `surya/settings.py`):
|
||||
|
||||
| Variable | Default | Purpose |
|
||||
| --- | --- | --- |
|
||||
| `SURYA_INFERENCE_URL` | `http://127.0.0.1:8000/v1` | vLLM OpenAI endpoint |
|
||||
| `SURYA_INFERENCE_PARALLEL` | `8` | Recognition worker hint |
|
||||
| `SURYA_INFERENCE_MAX_INFLIGHT` | `16` | Max concurrent HTTP calls to vLLM — **keep = `VLLM_MAX_NUM_SEQS`** |
|
||||
| `SUYA_OCR_MODE` | `block` | `block` (layout→per-block) or `full_page` (single call) |
|
||||
| `SUYA_MAX_BATCH_SIZE` | `8` | Request-batcher coalesce size |
|
||||
| `SUYA_BATCH_WAIT_MS` | `25` | Request-batcher wait window |
|
||||
| `SUYA_MAX_QUEUE_SIZE` | `128` | Request-batcher backlog cap |
|
||||
| `VLLM_DTYPE` | `float16` | `bfloat16` on A100/L40/4090; `float16` on T4 |
|
||||
| `VLLM_MAX_MODEL_LEN` | `18000` | vLLM context length |
|
||||
| `VLLM_MAX_NUM_SEQS` | `16` | vLLM sequence slots |
|
||||
| `VLLM_GPU_MEMORY_UTILIZATION` | `0.85` | vLLM GPU memory fraction |
|
||||
|
||||
---
|
||||
|
||||
## Benchmarks
|
||||
|
||||
- `benchmarks/` — concurrency sweep scripts and results (see the table above)
|
||||
- `scripts/quant/` — the quantization benchmarking toolchain (build, serve, capture, score, plot)
|
||||
|
||||
---
|
||||
|
||||
## Acknowledgements
|
||||
|
||||
This project is a serving wrapper around **[Surya OCR](https://github.com/datalab-to/surya)** by
|
||||
**[Datalab](https://www.datalab.to/)**. All layout detection and text recognition is done by their
|
||||
[`datalab-to/surya-ocr-2`](https://huggingface.co/datalab-to/surya-ocr-2) model — this repo only
|
||||
adds the vLLM serving, request batching, and API/CLI around it. Huge thanks to the Surya team for
|
||||
building and open-sourcing such a capable multilingual OCR model.
|
||||
|
||||
- Official site: <https://www.datalab.to/>
|
||||
- Surya GitHub: <https://github.com/datalab-to/surya>
|
||||
- Model: <https://huggingface.co/datalab-to/surya-ocr-2>
|
||||
@@ -0,0 +1,60 @@
|
||||
# Single Call Timing Result
|
||||
|
||||
Test image: `/path/to/suya-ocr-api/temp_image_c3c56948-282e-453b-8fac-6c482243d1e5.jpg`
|
||||
|
||||
Endpoint: `/v1/api/ai/suya_ocr_vllm/`
|
||||
|
||||
Run time: 2026-06-10 16:18 Asia/Dubai
|
||||
|
||||
Environment:
|
||||
|
||||
- `SURYA_INFERENCE_URL=http://127.0.0.1:43401/v1`
|
||||
- `SURYA_INFERENCE_PARALLEL=8`
|
||||
- `SUYA_OCR_MODE=block`
|
||||
- `SUYA_TIMING_ENABLED=true`
|
||||
|
||||
Summary:
|
||||
|
||||
| Stage | Time ms | Notes |
|
||||
|---|---:|---|
|
||||
| Client observed request | 7812.80 | Single HTTP call to API |
|
||||
| API request total | 7800.86 | FastAPI middleware duration |
|
||||
| Base64/PIL decode | 33.38 | 683,466 byte JPEG, 1410x1033 |
|
||||
| API wait for vLLM batcher | 7755.73 | Includes queue, layout, recognition |
|
||||
| OCR batch total | 7730.19 | Surya vLLM OCR path |
|
||||
| Layout predictor total | 2114.89 | One page layout call |
|
||||
| Layout vLLM generate | 2053.21 | Includes backend attach/start check |
|
||||
| vLLM backend attach/start check | 97.84 | Health/model lookup on first request |
|
||||
| Layout chat completion | 1902.34 | One vLLM `/chat/completions` call |
|
||||
| Layout parse/output boxes | 61.48 | JSON parse, bbox conversion, blank filtering |
|
||||
| Recognition block total | 5615.26 | Block OCR after layout |
|
||||
| Recognition build crop batch | 1.09 | Crop 17 OCR blocks, skip 1 block |
|
||||
| Recognition vLLM generate | 5613.12 | 17 block calls, parallelism 8 |
|
||||
| Recognition output assembly | 0.28 | Clean HTML and build page result |
|
||||
| Response assembly | 0.01 | No bbox image drawing |
|
||||
|
||||
Block OCR details:
|
||||
|
||||
| Metric | Value |
|
||||
|---|---:|
|
||||
| Layout blocks returned in response | 18 |
|
||||
| OCR block requests sent to vLLM | 17 |
|
||||
| Skipped blocks | 1 |
|
||||
| Sum of requested block max tokens | 6210 |
|
||||
| Generated block tokens | 4454 |
|
||||
| Sum of block chat completion times | 33259.06 ms |
|
||||
| Mean block chat completion time | 1956.42 ms |
|
||||
| Max block chat completion time | 4305.79 ms |
|
||||
| Wall time for all block completions | 5613.02 ms |
|
||||
|
||||
Conclusion:
|
||||
|
||||
The bottleneck is vLLM generation, not Python parsing or image decode. The single request spends about 72% of wall time in block OCR generation (`5613 ms / 7801 ms`) and about 24% in layout generation (`1902 ms / 7801 ms`). CPU-side crop/build/parse/assembly is negligible for this image.
|
||||
|
||||
Most useful next optimization directions:
|
||||
|
||||
1. Reduce number of OCR block requests per page.
|
||||
2. Reduce block max token budget.
|
||||
3. Merge or batch adjacent text blocks before recognition.
|
||||
4. Avoid running layout when caller can provide/skip detection.
|
||||
5. Tune vLLM scheduling/parallelism after reducing block count.
|
||||
@@ -0,0 +1,6 @@
|
||||
import uvicorn
|
||||
|
||||
from surya.endpoint.app import app # noqa: F401 (re-exported for `uvicorn api:app`)
|
||||
|
||||
if __name__ == "__main__":
|
||||
uvicorn.run("api:app", host="0.0.0.0", port=5002, workers=1)
|
||||
+64
@@ -0,0 +1,64 @@
|
||||
# import base64
|
||||
# import requests
|
||||
# import os
|
||||
|
||||
# # Configuration
|
||||
# API_URL = "http://localhost:5002/v1/api/ai/suya_ocr"
|
||||
|
||||
# # Helper to encode image
|
||||
# def encode_image_to_base64(path):
|
||||
# with open(path, "rb") as image_file:
|
||||
# return base64.b64encode(image_file.read()).decode("utf-8")
|
||||
|
||||
# def test_ocr_endpoint_requests():
|
||||
# image_path = "/path/to/surya/ttt.png"
|
||||
# assert os.path.exists(image_path), f"Image file not found: {image_path}"
|
||||
|
||||
# encoded_image = encode_image_to_base64(image_path)
|
||||
|
||||
# payload = {
|
||||
# "file": encoded_image,
|
||||
# "type": "jpg",
|
||||
# "skip_text_detection": False,
|
||||
# "skip_table_detection": False,
|
||||
# "recognize_math": False,
|
||||
# "ocr_with_boxes": True
|
||||
# }
|
||||
|
||||
# # Send request using requests
|
||||
# response = requests.post(API_URL, json=payload)
|
||||
|
||||
# assert response.status_code == 200
|
||||
# response_json = response.json()
|
||||
# assert response_json["code"] == 200
|
||||
# assert "ocr_text_json" in response_json["data"]
|
||||
# assert "text_lines" in response_json["data"]
|
||||
|
||||
# print("OCR Response Text:")
|
||||
# print(response_json["data"]["text_lines"])
|
||||
|
||||
|
||||
# import time
|
||||
# tik = time.time()
|
||||
# test_ocr_endpoint_requests()
|
||||
# tok = time.time()
|
||||
# print("time consumption: ", tok - tik)
|
||||
|
||||
|
||||
|
||||
import requests
|
||||
|
||||
# API 地址
|
||||
url = "http://127.0.0.1:5002/image2text"
|
||||
|
||||
# 测试用图像路径
|
||||
image_path = "/path/to/surya/ttt.png" # 替换为你自己的图像路径
|
||||
|
||||
# 发起 POST 请求
|
||||
with open(image_path, "rb") as image_file:
|
||||
files = {"file": ("ttt.png", image_file, "image/png")}
|
||||
response = requests.post(url, files=files)
|
||||
|
||||
# 打印响应
|
||||
print("Status Code:", response.status_code)
|
||||
print("Response JSON:", response.json())
|
||||
@@ -0,0 +1,15 @@
|
||||
<img alt="A colorful illustration of a structural steel brick facade with a complex graphic design of a steel brick facade an
|
||||
تسس الأطراف سنة ١٨٧٥ : سانيم ويشمارة نقلا
|
||||
رئيس مجلس الادارة رئيس الشعب وير عبد الله عبد البخاري إبراهيم تلفيح
|
||||
<img alt="A small rectangular label with a black border. It features the text 'Al-Akram' in a large, bold serif font at the bottom center. Below the text, the label contains the dates '13 OCTOBER 1981' and '1 APRI
|
||||
- السنة __1٠٠٧__ المدد ٣٩٦٣٨ -
|
||||
المنطقات المنتصرة
|
||||
عن اليوم المسيحي ع بعض وعرف منذ باحد ابن محمد المرق حول من الطائف السلامة
|
||||
اعدادة تأليف الوزارة بنفس تشكيلها الحد للحافظون باقيون لحين صدور قرار
|
||||
يعلن شسعب مصر اليوم أسلام الصالم لجمسع ارانته الخسرة في مواصلة مسير الوطس منصب أنور السادات حيث يتفق ملايين الواطئين على صناعيق الاستفتاء ليؤكد شسع مصر وأجماعه على تأييد الانس والمبادئ والسياسة الحكيمة التي وضعها الرئيس مغطة في شخص ثانيه ومرشح مجلس الشسع محمد حسنى مبارك .
|
||||
الاجتهان المسيّدات مص : الفَتْيَسارُ مِمْسارُك وتَمِمسارُ شَيْرِ مَمْيَس هُن عَلَيْهِ السَّمالَات عَلَت السِّيرَة جِبلِن السَّادَت وعَى تنتَلِ النَّواء مِن مسيدات. مص الثَّنِي الَّذِي مَن مَلَونا : ان كل من يجب رسم ويجب السلادات يجب ان يؤدى الجابه ويعرض على النَّواء بصرى في الامتداد لأخير السيد حسن ملوك دليسا الجوهرية كانه شريك بصرىّه من البل بناء مص والسلام : واغشته أن الشاركة في الامتداد والجابه فوهى وعلى الجوهر تأديته سواء كل رجل. أو لمرانة . والجواب على الجوهر على الجوهر المتوازية ساداته صلة .
|
||||
اسم الزعيم المسادات
|
||||
على مسندان التطوير وهابعة الدوليسة وجزيرة الكورسسان والق الجلس الخصوم المعارضة القادرة السن على الطاق اسم التنمية التي المتحالف على موجات التعوين لا واضحة مثلا في بالذيان تعتبر / لما وقت من أجل بحث حتى استحديث غير سنويا الجائزة والفم . كذلك عبر الإيتاء التوجيه المتبع بصفحة الإسبانية السن الذي انتم التنمية المستحالة على بوليزية الفرسان الحالة على فتاة الموسوعين والتالية تباين في بالحكم الجيولوجي الحر والحقوق والملكين المستحلة التي الحكم الحكما المستحلة المتبعة .
|
||||
حتى الطريق الرئيسي للانسياتولية -- كما قرر مجلس جامعة الالزينية الملائي التم الزعيم الراحل طبي جنيفةللالزينية وهو وفي ملابرح ترى الإيطالي الفسي والتنيدي للمحافظة الملائي اسم الزعيم الراحل طبي الواجه الجديدة الفريتوليز واجه سيود من الثوريا ولاني فستطيع الان ويتم استورداع ٢ الاند
|
||||
شعب مصر يقول كلمته اليوم ص (5)
|
||||
@isaadawi
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.3 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 358 KiB |
@@ -0,0 +1,27 @@
|
||||
大清光緒十六年
|
||||
本館新聞除禮拜外日出一張 諸君賜閱本報者無論本埠外埠均可函致本館取
|
||||
本埠外埠均可向我希龍照
|
||||
閱如外埠閱報照加寄費可
|
||||
也除日報照常刊送外復不
|
||||
工務局告示
|
||||
本局為便利商民起見，特將本局之工務局，改名為工務局，並將各項工程，由本局負責。凡有需要之工程，請向本局申請。此布。
|
||||
升印書籍 單件前有英國名
|
||||
體活字及各式花邊等延有
|
||||
西人排印工藝精所有中
|
||||
西文字均能刊印並有石版印字大小機器均用精工刊
|
||||
第二章
|
||||
吹笛
|
||||
巴 巴 巴 巴
|
||||
百四十
|
||||
第十四卷
|
||||
九字
|
||||
印工緻純倫前延有文士書
|
||||
寫各體文字現由上海購到各種書籍出售其書目另行刊布所有各書均詳列如左
|
||||
邦他所有名者置焉却宜與本館所發來札原擇其有關世道人心者附登否亦須無
|
||||
礙於人而後始錄其有謬罵成文以及支離費解者概不登報嗣後 諸君來札請祈
|
||||
登報嗣後，諸翁來札尚祈自諒，更須寫實，倘里而後敢登，但來稿登與不登，均不能
|
||||
庚寅正月初八日
|
||||
第五次每週發行日元一張 每張售價五角 每張可換一張日本銀行票
|
||||
叻報
|
||||
本館開設在舊巴虱口門牌第十一號
|
||||
本館啟事 ○啟者本埠自開設至今將益十餘所列新聞時事彙錄 諸君移步 惠顧焉 誠感良深 本館更欲大振報端 業於廈門 汕頭及南洋各埠 經營 採訪 新聞 如有要事 須為 刊登 以便 諸君 藉 閱 凡 近聞有 不肖之徒 在 外 妄 為 流言 有 礙 本館 本館 奉 札 外 處 之 信 稱 本館 寄 諸 君 以 查 報 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報 不肖者 有 強 明 指 出 廣 告 諸 君 之 所 自 報
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.4 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 428 KiB |
@@ -0,0 +1,72 @@
|
||||
Н Н А Т І Є
|
||||
VOL. IV. No. 181.
|
||||
DUBLIN, SATURDAY, MARCH 28, 1846.
|
||||
PRICE 6 D .
|
||||
DAVIS TESTIMONIAL.
|
||||
THE NATIVES OF IRELAND.
|
||||
THE Committee appointed for the purpose of receiving and analyzing the Telegraphic Statistics of LATE THOMAS DAVIS, ESQ., has received to leave open to public competition among American scientists and the public and the property of Specimens of a Status, in Sketch-model, together with a basin of Clay, lists size, to be sent for inspection on or before the 30th of June next. Any further information required may be obtained by applying to the Secretaries,
|
||||
LOWER DOMINICK-STREET, DUBLIN.
|
||||
THOMAS DAVIS
|
||||
HAS the pleasure of announcing that the PORTRAIT OF THE LATE THOMAS DAVIS, Esq.,
|
||||
LATE THOMAS DAVIS, Esq., P. W. BURTON, Esq., R.H.A., In the case of Sir John W.
|
||||
Is now ready for delivery
|
||||
in India Paper .. .. . 2s. 6d.
|
||||
.. .. . . . . . . . 1s. 0d.
|
||||
ished by JAMES DUFFY. Sole Agent
|
||||
in Ringrose, 11, Sherrard-street, Golden
|
||||
THE CABINET EDITION.
|
||||
LODGE'S PORTRAITS. —The First Volume, containing Thirty Portraits, in new earth, handsomely bound in crimson cloth, price 6 c . 6d. See also Lodge's Portraits, price 6 c . 64c. sec. or in Forty-eight Numbers, price 1 s . each. Dubin: J. M. GILLASHAN, 21, D'Olier-street; William Smith, 113, Fleet-street, London; Fraser and Co., Edinburgh.
|
||||
JUST PUBLISHED,
|
||||
THE FORTUNES OF TOROLOGH O'BRIEN:
|
||||
A TALE OF THE WORLD OF KING DEBORAH AND HER LASH AND D'Olier-street;
|
||||
London: William S. Orr and Co., Anne-corner, Pater-
|
||||
monter-row; Edinburgh: Fraser and Co. Sold by all Book-
|
||||
sellers.
|
||||
THE DUBLIN UNIVERSITY MAGAZINE
|
||||
NOTICE. JAMES DUFFY
|
||||
BEGS to announce that he has removed an IMMENSE STOCK OF BOOKS FROM 23, ANGELESA-STREET, To the Extensive Premises Above Occupied by MR. JOHN HARVIE, No. 10, WELLINGTON-QUAY.
|
||||
Now Ready, Part III., Price 1s., VALENTINE MCCLUTCHY, THE IRISH AGENT WITH ILLUSTRATIONS BY FRIZ. Dublin: Published by JAMES DUPPY; London Simpkin, Marshall, and Co.
|
||||
KNIGHT'S WEEKLY VOLUME, Printed in United States and Europe
|
||||
Price 1s. Just Published,
|
||||
VEGETABLE SUBSTANCES USED FOR
|
||||
THE FOOD OF MAN. Vol. I.
|
||||
London: Charles Knight and Co., 22, Ludgate-street.
|
||||
Just Published, 8rd., No. I, price 6d., PAYNE'S ILLUSTRATED LONDON
|
||||
A t the present moment no Illustrated Work exists, giving an adequate idea of the vast improvements of modern London. These, together with the great accommodation afforded for visiting them, imperatively mand such an undertaking as is now brought before it public.
|
||||
The Work will be published in Fortnightly Parts, at St pence, containing Five highly-finished steel Engraving and 16 to 24 pages of historical and descriptive letter-pre and in Monthly Parts, at One Shiling, to be completed about 20 Parts.
|
||||
THE NATIONAL LIBRARY FOR IRELAND,
|
||||
THE NATIONAL LIBRARY FOR IRELAND. BY DISTINGUISHED LITERARY ERISHMEN, Is now Publishing, in Fourpenny Monthly Volumes, complete.
|
||||
NATIONAL ENGRAVINGS, beautifully coloured, price One Penny, arc simultaneously issued with this Series.
|
||||
issued with this Series. Already Published, THE LIFE, TIMES, AND SPEECHES OF DANIEL O'CONNELL ESQ. M.P.
|
||||
DANIEL O'CONNELL, ESQ., M.P. (Fifth Edition.)
|
||||
NOW READY, PRICE FOURPENG THE LIFE AND WRITINGS OF JOHN R. CURRAN
|
||||
THE NEW WOOLLEN WAREHOUSE 17, DAME-STREET.
|
||||
17, DAME-STREET, IS NOW OPEN, EDWARD K. HOLROYD AND CO., PROPRIETORS.
|
||||
THE support of the Public is sought to proprietors of this Establishment with respect to the fact that the real and legitimate stock is secured by every one of the states provided for them the First and Leading Woollen Trade in D position they are resolved to attain, by SELLING WOOLLEN GOODS AT PRICES WHICH WILL BE PROVIDED BY THE SUPPORT.
|
||||
The extent of Stock will be commensurate with portence of the city of Dublin, and will invariably plete in its selection of Noveltics for the Season. The System now in use by many honest Trade marking their Goods, the Lowest Price, in Plain F will be adopted by E. K. H. and Co.
|
||||
At the present day, such and so many are the use of advertisers, that the Public wisely view with such professions of exclusive advantages.
|
||||
The Stock will comprise every class of WO GOODS, including the finest West of England Coded.
|
||||
E. K. H. and Co. would remark that their not to establish a Trade in low Woollens, but and First-class Goods ; from buyers of such goods.
|
||||
"Tailors will find a really EXTENSIVE AMENT OF TRIMMINGS.
|
||||
TERMS—CASH ONLY. E. K. HOLROYD AND COMPANY, INC.
|
||||
E. K. HOLBOYD AND COMPANY, 17 STREET, DUBLIN.
|
||||
MECHANICS INSTITUTE. TO BOOKSELLERS. T HE Board of Directors will receive for furnishing 300. worth of Books to a Catalogue of the Books required (with the &c.) can be had by applying to the Librarian, M
|
||||
THE PHOENIX PARK, DUBLIN.
|
||||
THE IRISH FARMERS' JOURNAL Wednesday last contains a beautifully painted of the Phoenix Park, the only Public Park in concise account is also given of the improvemented by Mr. Decimus Burton to the Common Woods and Forests.
|
||||
LOYAL NATIONAL REPEAL ASSOCIATION
|
||||
THE LOYAL NATIONAL REPEAL ACTION OF IRELAND will meet on MAY DAY, the 30th of MARCH, at ONE o'Clock, at the CONCILIATION HALL,
|
||||
CORN-EXCHANGE ROOMS. Ladies submitted fees will enter by the Third Executive.
|
||||
Ladies admitted free will enter by the Third Front I and proceed to the portion of the Side and End Gallery which have set a spart for them. For Members having been admitted to Found in the property of the Side are preserved under the Second Side Gallery. The access for them to the Seats under the Left-hand Gallery, is by the Door of the Hall nearest to the Exchange; and the access to the Seats under the hand Gallery, is by the Door in White's-lane.
|
||||
By Order, T. M. RAY, Secret
|
||||
BUTLER and ARCHER, 58, North STREET, have received their Annual Extraordinary of AGRICULTURAL SEEDS, consisting of 56 SEED OATS of the different vertexes, YETCHES, RED and WHITE CLOVERS, TRI COW-GRASS, SCOTCH and ITALIAN-RYER SEEDS.
|
||||
B. and A. have taken great care in the selecting Stock, and the quality of the various Articles can it is such as they feel assured will give satisfaction they can with great confidence recommend.
|
||||
NEW FARM SEEDS.
|
||||
JOHN HODRICHT,
|
||||
JOHN THOMAS-STEELEY,
|
||||
RESPECTFULLY makes known the a
|
||||
following New-York imported Seeds, which he is extensively supplied, having been pu
|
||||
competent judges in the best markets. He is to
|
||||
dispose of them at very Redwood New England
|
||||
Cover and Tree Shore, Hosteton, Potato, and Early August
|
||||
New York.
|
||||
Scotch, Hopeau, Patao, and Bary Angus of English and Irish Spring Vetches, True Art Grass. All warranted news and genuine. The public is also associated with and Refined Sagars, Tess, Wince, and White her offers to the Public on liberal terms, and for they will be found to give general satisfaction. 28, THOMAS-STREET.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 952 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
@@ -0,0 +1,5 @@
|
||||
Le Numéro : 75 Centimes.
|
||||
Mardi 10 Août 1937. - N° 963
|
||||
LE MIROIR DES SPORTS
|
||||
Le plus fort tirage des hebdomadaires sportifs
|
||||
LOUISETE FLEURET, 18 ANS , championne de France de natation de demi-fond, à battu dimanche, au stade nautique des Tournelles, les records de France des 800 mètres, 1.000 mètres et 1.500 mètres nage libre, respectivement en 12° 7" 2/10, 15° 16" 4/10 et 23° 13" 8/10. Voici, sortant de l'eau après son triple exploit, la jeune championne qui ne paraît nullement déprimée par son effort et qui conserve son frais sourire habituel.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 2.0 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 277 KiB |
@@ -0,0 +1,48 @@
|
||||
Виттова Плясна
|
||||
ОДНОДНЕВНАЯ ГАЗЕТА. ПАЛИ—ТИКО—ФИ—НАНСОВАЯ.
|
||||
Цѣна 5 копѣекъ.
|
||||
Условія подписки см. на послѣдней страницѣ.
|
||||
СКОРО ПОСТУПЯТЪ ВЪ ПРОДАЖУ НОВЫЯ ПАПИРОСЫ
|
||||
25 штукъ 25 копѣекъ.
|
||||
МОШЕННИКЪ 10 шт. 6 кош.
|
||||
Об Почтением:
|
||||
Альянсъ-Израелитъ.
|
||||
Чудо!
|
||||
„ТАРЛЭ“
|
||||
Чудо!
|
||||
Новость! Новость!
|
||||
СПЕЦІАЛЬНО ДЛЯ САМООБОРОНЫ.
|
||||
Револьверы повой спетемы «Таря». Главное преимущество.
|
||||
ЗА НЕНАДОВНОСТЬУ ПРОДАЕТСЯ ЗДАНЕ. Указоренники заброшен. 1-9, сорость Д. С. С. Борова.
|
||||
ЕБ ЗАЛЬ СОЕДИНЕНЫХЪ ПРИСУТСТВИ при благосклонномъ участии любимъяшихъ русской публикой социалъ демократовъ и прочихъ особъ.
|
||||
СОСТОИТСЯ ВЕЧЕРЪ, по программѣ всѣхъ революціонныхъ партій. Подъѣздъ со своей платформы.
|
||||
СОДВРЖАНІЕ.
|
||||
Объ редкали.—Текатарицы.—Фальтелян.—Равоварская: Правительства.—Стактонорали.—Замянки.—Мустепильная возрось.—Раболай о режиме.—О вроеперагах.—Ос катури.—Хроника.—Учебная одъван.—Театры и журналы.—Въезд кат. врояждан.—Обиходная редентура.—Объявления.
|
||||
С.Петербург, без числа. Вступили из период дель нашего издания священные узбурги наших и читаемых, что издание наше возник соответствует избранному для него активно. Пометили, что витвовали шляхов происходить темери у наст. Итакъ, товарищи внереды!
|
||||
Телеграммы
|
||||
Театрафиаго Агентетог.
|
||||
МОССА - 20 мафа. Назвистие молод дочи, орбитант на Петрура, был мир казаков, и углавл их молочное свото. Как при последних последних денях после последних днев были их получены прочие правители ста. Вдохновл влад и составил требомай просто последний день последний представитель образа, чтобы молочное пате жил в падании строить и влясти во приблизку из Баратов. Железное бластко привело к последним дням последних днев. МОССА 21 мафа. Вернула сообщения казакам не точе. Молод дочи на Петрура приглашил казаков правительство прибрестить почти прогрессивно вырожить. Казак отманяли оснавные и то, что ени не желают распать все наставить. Железное обратило себе
|
||||
ФЕЛЬЕТОНЪ.
|
||||
КОВЫЙ РЕВИЗОРЪ.
|
||||
Под Я присвоены васть, посподь, чтобы сообщить вам присворетное нам бить и на присвоенных местах.
|
||||
Как Рисковский.
|
||||
Тор, да я нас как слабуют, учрестить но собрать себя на восток.
|
||||
Ано, Одь Ворчат на!
|
||||
Лич, Луч, Посподь После мой!
|
||||
Одь Я как усиль будет правдучественно, после того, что я присвоен местам и согласно новому зат интервализивая — при шля, важно по серебряной после и ушля — при шля, после того, что я присвоен местам и согласно новому просто, успехам, затяги увачуюю спосочу. "Русское Спосоч", посочу посподу.
|
||||
Ано, Одь Ну, у „постояна“ спосочу нати.
|
||||
Тор, Нати, Тати, воть пом прямо пом прямо пом прямо пом прямо пом прямо пом (тамамо) пом прямо пом спроять какого-то же желан
|
||||
вит до него из замалых, и так, усмотряясь на простые ветви нестояния.
|
||||
По своей стороне, в том случае Давида и нескольких Дина стоят во время своего братства и противления с ними и нестоянием их отряда.
|
||||
В том случае Давид и нестояние их отряда не получают своего общения и забили прох. В том случае Давид и нестояние их отряда не получают своего общения.
|
||||
В ООСА 21 января. В листоке Давида и нестояния их отряда на простые ветви нестояния из Дав, так как они в листоке выбили нестояние их отряда и нестояние их отряда не получают своего общения и забили прох. В том случае Давид и нестояние их отряда не получают своего общения и забили прох.
|
||||
В ООСА 22 января. В листоке Давида и нестояния их отряда на простые ветви нестояния из Дав, так как они в листоке выбили нестояние их отряда и нестояние их отряда не получают своего общения и забили прох.
|
||||
В ООСА 23 января. В листоке Давида и нестояния их отряда на простые ветви нестояния из Дав, так как они в листоке выбили нестояние их отряда и нестояние их отряда не получают своего общения и забили прох.
|
||||
В ООСА 24 января. В листоке Давида и нестояния их отряда на простые ветви нестояния из Дав, так как они в листоке выбили нестояние их отряда и нестояние их отряда не получают своего общения и забили прох.
|
||||
В ООСА 25 января. В листоке Давида и нестояния их отряда на простые ветви нестояния из Дав, так как они в листоке выбили нестояние их отряда и нестояние их отряда не получают своего общения и забили прох.
|
||||
В ООСА 26 января. В листоке Давида и нестояния их отряда на простые ветви нестояния из Дав, так как они в листоке выбили нестояние их отряда и нестояние их отряда не получают своего общения и забили прох.
|
||||
В ООСА 27 января. В листоке Давида и нестояния их отряда на простые ветви нестояния из Дав, так как они в листоке выбили нестояние их отряда и нестояние их отряда не получают своего общения и забили прох.
|
||||
В ООСА 28 января. В листоке Давида и нестояния их отряда на простые ветви нестояния из Дав, так как они в листоке выбили нестояние их отряда и нестояние их отряда не получают своего общения и забили прох.
|
||||
В ООСА 29 января. В листоке Давида и нестояния их отряда на простые ветви нестояния из Дав, так как они в листоке выбили нестояние их отряда и нестояние их отряда не получают своего общения и забили прох.
|
||||
В ООСА 30 января. В листоке Давида и нестояния их отряда на простые ветви нестояния из Дав, так как они в листоке выбили нестояние их отряда и нестояние их отряда не получают своего общения и забили прох.
|
||||
В ООСА 31
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 218 KiB |
@@ -0,0 +1,267 @@
|
||||
"""Concurrency benchmark + old-vs-new comparison for the OCR service.
|
||||
|
||||
Two subcommands:
|
||||
|
||||
sweep Drive one running service across a list of concurrency levels and
|
||||
record latency/throughput/success per level. Reuses the request
|
||||
harness in concurrency_test.py.
|
||||
|
||||
compare Load two sweep result files (old + new) and emit a markdown table,
|
||||
a CSV, and comparison plots (latency & throughput vs concurrency,
|
||||
plus a per-level speedup bar).
|
||||
|
||||
Typical flow (see scripts/run_concurrency_comparison.sh):
|
||||
|
||||
# against the new vLLM service (its optimized endpoint)
|
||||
python concurrency_compare.py sweep --label new \
|
||||
--url http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/ \
|
||||
--image temp_image_c3c56948-282e-453b-8fac-6c482243d1e5.jpg \
|
||||
--levels 1,2,4,8,16 --out results/compare/new.json
|
||||
|
||||
# ... swap containers, then against the old service (its OCR endpoint)
|
||||
python concurrency_compare.py sweep --label old \
|
||||
--url http://127.0.0.1:5002/v1/api/ai/suya_ocr/ \
|
||||
--image temp_image_c3c56948-282e-453b-8fac-6c482243d1e5.jpg \
|
||||
--levels 1,2,4,8,16 --out results/compare/old.json
|
||||
|
||||
python concurrency_compare.py compare \
|
||||
--old results/compare/old.json --new results/compare/new.json \
|
||||
--out-dir results/compare
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from concurrency_test import _encode_image, run_endpoint
|
||||
|
||||
|
||||
def _parse_levels(value: str) -> List[int]:
|
||||
return [int(x.strip()) for x in value.split(",") if x.strip()]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- sweep -----
|
||||
|
||||
def run_sweep(args: argparse.Namespace) -> None:
|
||||
image_path = Path(args.image)
|
||||
image_type = image_path.suffix.lstrip(".").lower() or "png"
|
||||
payload = {
|
||||
"file": _encode_image(image_path),
|
||||
"type": "jpg" if image_type == "jpeg" else image_type,
|
||||
"skip_text_detection": False,
|
||||
"skip_table_detection": False,
|
||||
"recognize_math": False,
|
||||
"ocr_with_boxes": True,
|
||||
}
|
||||
|
||||
levels = _parse_levels(args.levels)
|
||||
rows: List[Dict[str, Any]] = []
|
||||
for concurrency in levels:
|
||||
requests_count = max(concurrency, concurrency * args.reqs_per_level)
|
||||
print(f"[{args.label}] concurrency={concurrency} requests={requests_count}", flush=True)
|
||||
row = run_endpoint(
|
||||
args.label, args.url, payload,
|
||||
requests_count=requests_count, concurrency=concurrency, timeout=args.timeout,
|
||||
)
|
||||
rows.append(row)
|
||||
lat = row["latency_seconds"]
|
||||
print(f" -> success={row['success']}/{row['requests']} "
|
||||
f"rps={row['throughput_rps']:.3f} mean={lat['mean']:.2f}s p95={lat['p95']:.2f}s",
|
||||
flush=True)
|
||||
|
||||
out = {
|
||||
"label": args.label,
|
||||
"url": args.url,
|
||||
"image": str(image_path),
|
||||
"levels": levels,
|
||||
"reqs_per_level": args.reqs_per_level,
|
||||
"started_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
|
||||
"results": rows,
|
||||
}
|
||||
out_path = Path(args.out)
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
out_path.write_text(json.dumps(out, indent=2), encoding="utf-8")
|
||||
print(f"wrote {out_path}")
|
||||
|
||||
|
||||
# -------------------------------------------------------------- compare -----
|
||||
|
||||
def _index_by_concurrency(sweep: Dict[str, Any]) -> Dict[int, Dict[str, Any]]:
|
||||
return {row["concurrency"]: row for row in sweep["results"]}
|
||||
|
||||
|
||||
def _comparison_rows(old: Dict[str, Any], new: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
old_by_c = _index_by_concurrency(old)
|
||||
new_by_c = _index_by_concurrency(new)
|
||||
concurrencies = sorted(set(old_by_c) | set(new_by_c))
|
||||
rows = []
|
||||
for c in concurrencies:
|
||||
o = old_by_c.get(c)
|
||||
n = new_by_c.get(c)
|
||||
row: Dict[str, Any] = {"concurrency": c}
|
||||
row["old_mean_s"] = round(o["latency_seconds"]["mean"], 2) if o else None
|
||||
row["new_mean_s"] = round(n["latency_seconds"]["mean"], 2) if n else None
|
||||
row["old_p95_s"] = round(o["latency_seconds"]["p95"], 2) if o else None
|
||||
row["new_p95_s"] = round(n["latency_seconds"]["p95"], 2) if n else None
|
||||
row["old_rps"] = round(o["throughput_rps"], 3) if o else None
|
||||
row["new_rps"] = round(n["throughput_rps"], 3) if n else None
|
||||
row["old_fail"] = o["failed"] if o else None
|
||||
row["new_fail"] = n["failed"] if n else None
|
||||
if o and n and o["latency_seconds"]["mean"] and n["latency_seconds"]["mean"]:
|
||||
row["latency_speedup"] = round(o["latency_seconds"]["mean"] / n["latency_seconds"]["mean"], 2)
|
||||
else:
|
||||
row["latency_speedup"] = None
|
||||
if o and n and o["throughput_rps"]:
|
||||
row["throughput_gain"] = round(n["throughput_rps"] / o["throughput_rps"], 2)
|
||||
else:
|
||||
row["throughput_gain"] = None
|
||||
rows.append(row)
|
||||
return rows
|
||||
|
||||
|
||||
def _write_table(rows: List[Dict[str, Any]], old_label: str, new_label: str, path: Path) -> str:
|
||||
cols = [
|
||||
("concurrency", "conc"),
|
||||
("old_mean_s", f"{old_label} mean(s)"),
|
||||
("new_mean_s", f"{new_label} mean(s)"),
|
||||
("latency_speedup", "latency ×"),
|
||||
("old_p95_s", f"{old_label} p95(s)"),
|
||||
("new_p95_s", f"{new_label} p95(s)"),
|
||||
("old_rps", f"{old_label} rps"),
|
||||
("new_rps", f"{new_label} rps"),
|
||||
("throughput_gain", "rps ×"),
|
||||
("old_fail", f"{old_label} fail"),
|
||||
("new_fail", f"{new_label} fail"),
|
||||
]
|
||||
header = "| " + " | ".join(label for _, label in cols) + " |"
|
||||
sep = "| " + " | ".join("---" for _ in cols) + " |"
|
||||
lines = [header, sep]
|
||||
for r in rows:
|
||||
cells = []
|
||||
for key, _ in cols:
|
||||
v = r.get(key)
|
||||
cells.append("" if v is None else str(v))
|
||||
lines.append("| " + " | ".join(cells) + " |")
|
||||
md = "\n".join(lines) + "\n"
|
||||
path.write_text(md, encoding="utf-8")
|
||||
return md
|
||||
|
||||
|
||||
def _write_csv(rows: List[Dict[str, Any]], path: Path) -> None:
|
||||
if not rows:
|
||||
return
|
||||
with path.open("w", newline="", encoding="utf-8") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
|
||||
def _write_plots(rows, old_label, new_label, out_dir: Path) -> List[Path]:
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
x = [r["concurrency"] for r in rows]
|
||||
paths = []
|
||||
|
||||
# 1. latency (mean + p95) vs concurrency
|
||||
fig, ax = plt.subplots(figsize=(9, 6))
|
||||
ax.plot(x, [r["old_mean_s"] for r in rows], marker="o", color="#c0392b", label=f"{old_label} mean")
|
||||
ax.plot(x, [r["old_p95_s"] for r in rows], marker="^", color="#c0392b", linestyle="--", label=f"{old_label} p95")
|
||||
ax.plot(x, [r["new_mean_s"] for r in rows], marker="o", color="#27ae60", label=f"{new_label} mean")
|
||||
ax.plot(x, [r["new_p95_s"] for r in rows], marker="^", color="#27ae60", linestyle="--", label=f"{new_label} p95")
|
||||
ax.set_xlabel("Concurrency (simultaneous requests)")
|
||||
ax.set_ylabel("Latency per request (s)")
|
||||
ax.set_title("OCR latency vs concurrency — old vs new")
|
||||
ax.grid(True, alpha=0.3)
|
||||
ax.legend()
|
||||
fig.tight_layout()
|
||||
p = out_dir / "latency_vs_concurrency.png"
|
||||
fig.savefig(p, dpi=160)
|
||||
plt.close(fig)
|
||||
paths.append(p)
|
||||
|
||||
# 2. throughput vs concurrency
|
||||
fig, ax = plt.subplots(figsize=(9, 6))
|
||||
ax.plot(x, [r["old_rps"] for r in rows], marker="o", color="#c0392b", label=old_label)
|
||||
ax.plot(x, [r["new_rps"] for r in rows], marker="o", color="#27ae60", label=new_label)
|
||||
ax.set_xlabel("Concurrency (simultaneous requests)")
|
||||
ax.set_ylabel("Throughput (requests/s)")
|
||||
ax.set_title("OCR throughput vs concurrency — old vs new")
|
||||
ax.grid(True, alpha=0.3)
|
||||
ax.legend()
|
||||
fig.tight_layout()
|
||||
p = out_dir / "throughput_vs_concurrency.png"
|
||||
fig.savefig(p, dpi=160)
|
||||
plt.close(fig)
|
||||
paths.append(p)
|
||||
|
||||
# 3. per-level latency speedup bar
|
||||
fig, ax = plt.subplots(figsize=(9, 6))
|
||||
speedups = [r["latency_speedup"] or 0 for r in rows]
|
||||
ax.bar([str(c) for c in x], speedups, color="#2980b9")
|
||||
ax.axhline(1.0, color="gray", linestyle="--", linewidth=1)
|
||||
for i, v in enumerate(speedups):
|
||||
ax.text(i, v, f"{v:.2f}×", ha="center", va="bottom")
|
||||
ax.set_xlabel("Concurrency")
|
||||
ax.set_ylabel(f"Latency speedup ({old_label} mean / {new_label} mean)")
|
||||
ax.set_title("Per-level latency speedup (>1 = new is faster)")
|
||||
ax.grid(True, axis="y", alpha=0.3)
|
||||
fig.tight_layout()
|
||||
p = out_dir / "latency_speedup.png"
|
||||
fig.savefig(p, dpi=160)
|
||||
plt.close(fig)
|
||||
paths.append(p)
|
||||
|
||||
return paths
|
||||
|
||||
|
||||
def run_compare(args: argparse.Namespace) -> None:
|
||||
old = json.loads(Path(args.old).read_text(encoding="utf-8"))
|
||||
new = json.loads(Path(args.new).read_text(encoding="utf-8"))
|
||||
old_label = old.get("label", "old")
|
||||
new_label = new.get("label", "new")
|
||||
out_dir = Path(args.out_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
rows = _comparison_rows(old, new)
|
||||
md = _write_table(rows, old_label, new_label, out_dir / "comparison_table.md")
|
||||
_write_csv(rows, out_dir / "comparison_table.csv")
|
||||
plots = _write_plots(rows, old_label, new_label, out_dir)
|
||||
|
||||
print(md)
|
||||
print("plots:")
|
||||
for p in plots:
|
||||
print(f" {p}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
sub = parser.add_subparsers(dest="cmd", required=True)
|
||||
|
||||
sp = sub.add_parser("sweep", help="benchmark one running service across concurrency levels")
|
||||
sp.add_argument("--label", required=True, help="short name for this service, e.g. old / new")
|
||||
sp.add_argument("--url", required=True, help="full endpoint URL to POST to")
|
||||
sp.add_argument("--image", required=True, help="path to a png/jpg test image")
|
||||
sp.add_argument("--levels", default="1,2,4,8,16", help="comma-separated concurrency levels")
|
||||
sp.add_argument("--reqs-per-level", type=int, default=2, help="requests = level * this (>=level)")
|
||||
sp.add_argument("--timeout", type=float, default=900)
|
||||
sp.add_argument("--out", required=True, help="output JSON path")
|
||||
sp.set_defaults(func=run_sweep)
|
||||
|
||||
cp = sub.add_parser("compare", help="compare two sweep result files")
|
||||
cp.add_argument("--old", required=True, help="old-service sweep JSON")
|
||||
cp.add_argument("--new", required=True, help="new-service sweep JSON")
|
||||
cp.add_argument("--out-dir", default="results/compare")
|
||||
cp.set_defaults(func=run_compare)
|
||||
|
||||
args = parser.parse_args()
|
||||
args.func(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,125 @@
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from concurrency_test import DEFAULT_NEW_ENDPOINT, _encode_image, run_endpoint
|
||||
|
||||
|
||||
def parse_levels(value: str) -> list[int]:
|
||||
return [int(item.strip()) for item in value.split(",") if item.strip()]
|
||||
|
||||
|
||||
def write_csv(path: Path, rows: list[dict]) -> None:
|
||||
columns = [
|
||||
"concurrency",
|
||||
"requests",
|
||||
"success",
|
||||
"failed",
|
||||
"wall_seconds",
|
||||
"throughput_rps",
|
||||
"latency_min",
|
||||
"latency_mean",
|
||||
"latency_p50",
|
||||
"latency_p95",
|
||||
"latency_max",
|
||||
]
|
||||
with path.open("w", newline="", encoding="utf-8") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=columns)
|
||||
writer.writeheader()
|
||||
for row in rows:
|
||||
lat = row["latency_seconds"]
|
||||
writer.writerow(
|
||||
{
|
||||
"concurrency": row["concurrency"],
|
||||
"requests": row["requests"],
|
||||
"success": row["success"],
|
||||
"failed": row["failed"],
|
||||
"wall_seconds": row["wall_seconds"],
|
||||
"throughput_rps": row["throughput_rps"],
|
||||
"latency_min": lat["min"],
|
||||
"latency_mean": lat["mean"],
|
||||
"latency_p50": lat["p50"],
|
||||
"latency_p95": lat["p95"],
|
||||
"latency_max": lat["max"],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def write_plot(path: Path, rows: list[dict]) -> None:
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
x = [row["concurrency"] for row in rows]
|
||||
mean = [row["latency_seconds"]["mean"] for row in rows]
|
||||
p50 = [row["latency_seconds"]["p50"] for row in rows]
|
||||
p95 = [row["latency_seconds"]["p95"] for row in rows]
|
||||
|
||||
plt.figure(figsize=(10, 6))
|
||||
plt.plot(x, mean, marker="o", label="mean")
|
||||
plt.plot(x, p50, marker="o", label="p50")
|
||||
plt.plot(x, p95, marker="o", label="p95")
|
||||
plt.xlabel("Concurrency")
|
||||
plt.ylabel("Latency per request (seconds)")
|
||||
plt.title("vLLM OCR latency vs concurrency")
|
||||
plt.grid(True, alpha=0.3)
|
||||
plt.legend()
|
||||
plt.tight_layout()
|
||||
plt.savefig(path, dpi=160)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Sweep vLLM OCR endpoint concurrency levels.")
|
||||
parser.add_argument("--image", required=True)
|
||||
parser.add_argument("--url", default=DEFAULT_NEW_ENDPOINT)
|
||||
parser.add_argument("--levels", default="1,2,4,6,8,10,12,20,30,40,50")
|
||||
parser.add_argument("--timeout", type=float, default=900)
|
||||
parser.add_argument("--output-json", default="concurrency_sweep_results.json")
|
||||
parser.add_argument("--output-csv", default="concurrency_sweep_results.csv")
|
||||
parser.add_argument("--output-plot", default="concurrency_sweep_latency.png")
|
||||
args = parser.parse_args()
|
||||
|
||||
image_path = Path(args.image)
|
||||
image_type = image_path.suffix.lstrip(".").lower() or "png"
|
||||
payload = {
|
||||
"file": _encode_image(image_path),
|
||||
"type": "jpg" if image_type == "jpeg" else image_type,
|
||||
"skip_text_detection": False,
|
||||
"skip_table_detection": False,
|
||||
"recognize_math": False,
|
||||
"ocr_with_boxes": True,
|
||||
}
|
||||
|
||||
levels = parse_levels(args.levels)
|
||||
rows = []
|
||||
for concurrency in levels:
|
||||
print(f"running concurrency={concurrency}", flush=True)
|
||||
rows.append(
|
||||
run_endpoint(
|
||||
"vllm",
|
||||
args.url,
|
||||
payload,
|
||||
requests_count=concurrency,
|
||||
concurrency=concurrency,
|
||||
timeout=args.timeout,
|
||||
)
|
||||
)
|
||||
|
||||
result = {
|
||||
"started_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
|
||||
"image": str(image_path),
|
||||
"url": args.url,
|
||||
"levels": levels,
|
||||
"results": rows,
|
||||
}
|
||||
Path(args.output_json).write_text(json.dumps(result, indent=2), encoding="utf-8")
|
||||
write_csv(Path(args.output_csv), rows)
|
||||
write_plot(Path(args.output_plot), rows)
|
||||
print(json.dumps(result, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 91 KiB |
@@ -0,0 +1,12 @@
|
||||
concurrency,requests,success,failed,wall_seconds,throughput_rps,latency_min,latency_mean,latency_p50,latency_p95,latency_max
|
||||
1,1,1,0,7.6577067924663424,0.13058739738949013,7.656610286794603,7.656610286794603,7.656610286794603,7.656610286794603,7.656610286794603
|
||||
2,2,2,0,11.936453193426132,0.16755395992349534,11.934747929684818,11.935045637656003,11.934747929684818,11.935343345627189,11.935343345627189
|
||||
4,4,4,0,20.956660461612046,0.19087010582278188,20.942796256393194,20.951773355016485,20.954541454091668,20.95546007528901,20.95546007528901
|
||||
6,6,6,0,29.489064288325608,0.20346525550406608,29.45817438978702,29.47708616297071,29.48345213010907,29.487088727764785,29.487088727764785
|
||||
8,8,8,0,41.46563811413944,0.192930830534405,11.84945331979543,34.050638656364754,41.457184289582074,41.46130123361945,41.46130123361945
|
||||
10,10,10,0,50.19571524951607,0.19922019141058875,12.104127056896687,42.553733562212436,50.15600565075874,50.19275312870741,50.19275312870741
|
||||
12,12,12,0,59.99224958010018,0.20002583807059768,25.362822842784226,45.54706535985073,59.93020099774003,59.95823861565441,59.98695467971265
|
||||
20,20,20,0,99.85960746835917,0.20028117981874766,16.85126264579594,66.68330737100914,55.060471390374005,92.9106959477067,99.76074412371963
|
||||
30,30,30,0,147.6483807535842,0.20318543181362822,20.873554840683937,90.49832232653473,97.37336550559849,147.4727698881179,147.48508568760008
|
||||
40,40,40,0,195.21150322351605,0.20490595758693703,25.34921144787222,114.89935769808945,102.41363409627229,194.96774306707084,194.99985321611166
|
||||
50,50,50,0,241.15029445569962,0.20733957680978582,88.28880738746375,149.28542947791516,126.26868482958525,240.86813501361758,240.87839913833886
|
||||
|
@@ -0,0 +1,218 @@
|
||||
{
|
||||
"started_at": "2026-06-10T15:58:07+0400",
|
||||
"image": "/path/to/suya-ocr-api/temp_image_c3c56948-282e-453b-8fac-6c482243d1e5.jpg",
|
||||
"url": "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/",
|
||||
"levels": [
|
||||
1,
|
||||
2,
|
||||
4,
|
||||
6,
|
||||
8,
|
||||
10,
|
||||
12,
|
||||
20,
|
||||
30,
|
||||
40,
|
||||
50
|
||||
],
|
||||
"results": [
|
||||
{
|
||||
"name": "vllm",
|
||||
"url": "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/",
|
||||
"requests": 1,
|
||||
"concurrency": 1,
|
||||
"success": 1,
|
||||
"failed": 0,
|
||||
"wall_seconds": 7.6577067924663424,
|
||||
"throughput_rps": 0.13058739738949013,
|
||||
"latency_seconds": {
|
||||
"min": 7.656610286794603,
|
||||
"mean": 7.656610286794603,
|
||||
"p50": 7.656610286794603,
|
||||
"p95": 7.656610286794603,
|
||||
"max": 7.656610286794603
|
||||
},
|
||||
"errors": []
|
||||
},
|
||||
{
|
||||
"name": "vllm",
|
||||
"url": "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/",
|
||||
"requests": 2,
|
||||
"concurrency": 2,
|
||||
"success": 2,
|
||||
"failed": 0,
|
||||
"wall_seconds": 11.936453193426132,
|
||||
"throughput_rps": 0.16755395992349534,
|
||||
"latency_seconds": {
|
||||
"min": 11.934747929684818,
|
||||
"mean": 11.935045637656003,
|
||||
"p50": 11.934747929684818,
|
||||
"p95": 11.935343345627189,
|
||||
"max": 11.935343345627189
|
||||
},
|
||||
"errors": []
|
||||
},
|
||||
{
|
||||
"name": "vllm",
|
||||
"url": "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/",
|
||||
"requests": 4,
|
||||
"concurrency": 4,
|
||||
"success": 4,
|
||||
"failed": 0,
|
||||
"wall_seconds": 20.956660461612046,
|
||||
"throughput_rps": 0.19087010582278188,
|
||||
"latency_seconds": {
|
||||
"min": 20.942796256393194,
|
||||
"mean": 20.951773355016485,
|
||||
"p50": 20.954541454091668,
|
||||
"p95": 20.95546007528901,
|
||||
"max": 20.95546007528901
|
||||
},
|
||||
"errors": []
|
||||
},
|
||||
{
|
||||
"name": "vllm",
|
||||
"url": "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/",
|
||||
"requests": 6,
|
||||
"concurrency": 6,
|
||||
"success": 6,
|
||||
"failed": 0,
|
||||
"wall_seconds": 29.489064288325608,
|
||||
"throughput_rps": 0.20346525550406608,
|
||||
"latency_seconds": {
|
||||
"min": 29.45817438978702,
|
||||
"mean": 29.47708616297071,
|
||||
"p50": 29.48345213010907,
|
||||
"p95": 29.487088727764785,
|
||||
"max": 29.487088727764785
|
||||
},
|
||||
"errors": []
|
||||
},
|
||||
{
|
||||
"name": "vllm",
|
||||
"url": "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/",
|
||||
"requests": 8,
|
||||
"concurrency": 8,
|
||||
"success": 8,
|
||||
"failed": 0,
|
||||
"wall_seconds": 41.46563811413944,
|
||||
"throughput_rps": 0.192930830534405,
|
||||
"latency_seconds": {
|
||||
"min": 11.84945331979543,
|
||||
"mean": 34.050638656364754,
|
||||
"p50": 41.457184289582074,
|
||||
"p95": 41.46130123361945,
|
||||
"max": 41.46130123361945
|
||||
},
|
||||
"errors": []
|
||||
},
|
||||
{
|
||||
"name": "vllm",
|
||||
"url": "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/",
|
||||
"requests": 10,
|
||||
"concurrency": 10,
|
||||
"success": 10,
|
||||
"failed": 0,
|
||||
"wall_seconds": 50.19571524951607,
|
||||
"throughput_rps": 0.19922019141058875,
|
||||
"latency_seconds": {
|
||||
"min": 12.104127056896687,
|
||||
"mean": 42.553733562212436,
|
||||
"p50": 50.15600565075874,
|
||||
"p95": 50.19275312870741,
|
||||
"max": 50.19275312870741
|
||||
},
|
||||
"errors": []
|
||||
},
|
||||
{
|
||||
"name": "vllm",
|
||||
"url": "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/",
|
||||
"requests": 12,
|
||||
"concurrency": 12,
|
||||
"success": 12,
|
||||
"failed": 0,
|
||||
"wall_seconds": 59.99224958010018,
|
||||
"throughput_rps": 0.20002583807059768,
|
||||
"latency_seconds": {
|
||||
"min": 25.362822842784226,
|
||||
"mean": 45.54706535985073,
|
||||
"p50": 59.93020099774003,
|
||||
"p95": 59.95823861565441,
|
||||
"max": 59.98695467971265
|
||||
},
|
||||
"errors": []
|
||||
},
|
||||
{
|
||||
"name": "vllm",
|
||||
"url": "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/",
|
||||
"requests": 20,
|
||||
"concurrency": 20,
|
||||
"success": 20,
|
||||
"failed": 0,
|
||||
"wall_seconds": 99.85960746835917,
|
||||
"throughput_rps": 0.20028117981874766,
|
||||
"latency_seconds": {
|
||||
"min": 16.85126264579594,
|
||||
"mean": 66.68330737100914,
|
||||
"p50": 55.060471390374005,
|
||||
"p95": 92.9106959477067,
|
||||
"max": 99.76074412371963
|
||||
},
|
||||
"errors": []
|
||||
},
|
||||
{
|
||||
"name": "vllm",
|
||||
"url": "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/",
|
||||
"requests": 30,
|
||||
"concurrency": 30,
|
||||
"success": 30,
|
||||
"failed": 0,
|
||||
"wall_seconds": 147.6483807535842,
|
||||
"throughput_rps": 0.20318543181362822,
|
||||
"latency_seconds": {
|
||||
"min": 20.873554840683937,
|
||||
"mean": 90.49832232653473,
|
||||
"p50": 97.37336550559849,
|
||||
"p95": 147.4727698881179,
|
||||
"max": 147.48508568760008
|
||||
},
|
||||
"errors": []
|
||||
},
|
||||
{
|
||||
"name": "vllm",
|
||||
"url": "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/",
|
||||
"requests": 40,
|
||||
"concurrency": 40,
|
||||
"success": 40,
|
||||
"failed": 0,
|
||||
"wall_seconds": 195.21150322351605,
|
||||
"throughput_rps": 0.20490595758693703,
|
||||
"latency_seconds": {
|
||||
"min": 25.34921144787222,
|
||||
"mean": 114.89935769808945,
|
||||
"p50": 102.41363409627229,
|
||||
"p95": 194.96774306707084,
|
||||
"max": 194.99985321611166
|
||||
},
|
||||
"errors": []
|
||||
},
|
||||
{
|
||||
"name": "vllm",
|
||||
"url": "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/",
|
||||
"requests": 50,
|
||||
"concurrency": 50,
|
||||
"success": 50,
|
||||
"failed": 0,
|
||||
"wall_seconds": 241.15029445569962,
|
||||
"throughput_rps": 0.20733957680978582,
|
||||
"latency_seconds": {
|
||||
"min": 88.28880738746375,
|
||||
"mean": 149.28542947791516,
|
||||
"p50": 126.26868482958525,
|
||||
"p95": 240.86813501361758,
|
||||
"max": 240.87839913833886
|
||||
},
|
||||
"errors": []
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,185 @@
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import statistics
|
||||
import subprocess
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
DEFAULT_OLD_ENDPOINT = "http://127.0.0.1:5002/v1/api/ai/suya_ocr/"
|
||||
DEFAULT_NEW_ENDPOINT = "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/"
|
||||
|
||||
|
||||
def _encode_image(path: Path) -> str:
|
||||
return base64.b64encode(path.read_bytes()).decode("utf-8")
|
||||
|
||||
|
||||
def _preflight() -> Dict[str, Any]:
|
||||
checks: Dict[str, Any] = {}
|
||||
commands = {
|
||||
"docker_runtimes": ["docker", "info", "--format", "{{json .Runtimes}}"],
|
||||
"nvidia_smi": ["nvidia-smi", "-L"],
|
||||
"nvidia_container_runtime": ["which", "nvidia-container-runtime"],
|
||||
"nvidia_ctk": ["which", "nvidia-ctk"],
|
||||
"docker_gpus": [
|
||||
"docker",
|
||||
"run",
|
||||
"--rm",
|
||||
"--gpus",
|
||||
"all",
|
||||
"--entrypoint",
|
||||
"nvidia-smi",
|
||||
"vllm/vllm-openai:v0.20.1",
|
||||
"-L",
|
||||
],
|
||||
}
|
||||
for name, cmd in commands.items():
|
||||
try:
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
|
||||
checks[name] = {
|
||||
"ok": result.returncode == 0,
|
||||
"stdout": result.stdout.strip(),
|
||||
"stderr": result.stderr.strip(),
|
||||
}
|
||||
except Exception as exc:
|
||||
checks[name] = {"ok": False, "stdout": "", "stderr": str(exc)}
|
||||
runtimes = checks.get("docker_runtimes", {}).get("stdout", "")
|
||||
checks["docker_has_nvidia_runtime"] = '"nvidia"' in runtimes
|
||||
if not checks["docker_has_nvidia_runtime"] and checks.get("docker_gpus", {}).get("ok"):
|
||||
checks["recommended_fix"] = (
|
||||
"Docker GPU passthrough works with --gpus. Use the local "
|
||||
"scripts/docker_nvidia_runtime_compat.sh wrapper to strip the legacy "
|
||||
"--runtime nvidia flag from Surya's vLLM spawn command."
|
||||
)
|
||||
else:
|
||||
checks["recommended_fix"] = (
|
||||
"Run: sudo nvidia-ctk runtime configure --runtime=docker "
|
||||
"--config=/etc/docker/daemon.json && sudo systemctl restart docker"
|
||||
)
|
||||
return checks
|
||||
|
||||
|
||||
def _percentile(values: List[float], percentile: float) -> float:
|
||||
if not values:
|
||||
return 0.0
|
||||
ordered = sorted(values)
|
||||
index = min(len(ordered) - 1, int(round((percentile / 100) * (len(ordered) - 1))))
|
||||
return ordered[index]
|
||||
|
||||
|
||||
def _post_once(url: str, payload: Dict[str, Any], timeout: float) -> Dict[str, Any]:
|
||||
start = time.perf_counter()
|
||||
try:
|
||||
response = requests.post(url, json=payload, timeout=timeout)
|
||||
elapsed = time.perf_counter() - start
|
||||
body = response.json() if response.content else {}
|
||||
return {
|
||||
"ok": response.status_code == 200 and body.get("code") == 200,
|
||||
"status_code": response.status_code,
|
||||
"api_code": body.get("code"),
|
||||
"elapsed_seconds": elapsed,
|
||||
"error": body.get("message") if body.get("code") != 200 else None,
|
||||
}
|
||||
except Exception as exc:
|
||||
return {
|
||||
"ok": False,
|
||||
"status_code": None,
|
||||
"api_code": None,
|
||||
"elapsed_seconds": time.perf_counter() - start,
|
||||
"error": str(exc),
|
||||
}
|
||||
|
||||
|
||||
def run_endpoint(
|
||||
name: str,
|
||||
url: str,
|
||||
payload: Dict[str, Any],
|
||||
requests_count: int,
|
||||
concurrency: int,
|
||||
timeout: float,
|
||||
) -> Dict[str, Any]:
|
||||
start = time.perf_counter()
|
||||
samples: List[Dict[str, Any]] = []
|
||||
with ThreadPoolExecutor(max_workers=concurrency) as pool:
|
||||
futures = [
|
||||
pool.submit(_post_once, url, payload, timeout)
|
||||
for _ in range(requests_count)
|
||||
]
|
||||
for future in as_completed(futures):
|
||||
samples.append(future.result())
|
||||
wall_seconds = time.perf_counter() - start
|
||||
latencies = [sample["elapsed_seconds"] for sample in samples]
|
||||
success_count = sum(1 for sample in samples if sample["ok"])
|
||||
return {
|
||||
"name": name,
|
||||
"url": url,
|
||||
"requests": requests_count,
|
||||
"concurrency": concurrency,
|
||||
"success": success_count,
|
||||
"failed": requests_count - success_count,
|
||||
"wall_seconds": wall_seconds,
|
||||
"throughput_rps": requests_count / wall_seconds if wall_seconds else 0.0,
|
||||
"latency_seconds": {
|
||||
"min": min(latencies) if latencies else 0.0,
|
||||
"mean": statistics.mean(latencies) if latencies else 0.0,
|
||||
"p50": _percentile(latencies, 50),
|
||||
"p95": _percentile(latencies, 95),
|
||||
"max": max(latencies) if latencies else 0.0,
|
||||
},
|
||||
"errors": [sample for sample in samples if not sample["ok"]][:10],
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Compare legacy and vLLM OCR endpoint concurrency.")
|
||||
parser.add_argument("--image", required=True, help="Path to an input png/jpg image.")
|
||||
parser.add_argument("--old-url", default=DEFAULT_OLD_ENDPOINT)
|
||||
parser.add_argument("--new-url", default=DEFAULT_NEW_ENDPOINT)
|
||||
parser.add_argument("--requests", type=int, default=20)
|
||||
parser.add_argument("--concurrency", type=int, default=8)
|
||||
parser.add_argument("--timeout", type=float, default=900)
|
||||
parser.add_argument("--output", default="concurrency_test_results.json")
|
||||
parser.add_argument("--skip-old", action="store_true")
|
||||
parser.add_argument("--skip-new", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
image_path = Path(args.image)
|
||||
image_type = image_path.suffix.lstrip(".").lower() or "png"
|
||||
payload = {
|
||||
"file": _encode_image(image_path),
|
||||
"type": "jpg" if image_type == "jpeg" else image_type,
|
||||
"skip_text_detection": False,
|
||||
"skip_table_detection": False,
|
||||
"recognize_math": False,
|
||||
"ocr_with_boxes": True,
|
||||
}
|
||||
|
||||
results = {
|
||||
"started_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
|
||||
"image": str(image_path),
|
||||
"requests": args.requests,
|
||||
"concurrency": args.concurrency,
|
||||
"preflight": _preflight(),
|
||||
"endpoints": [],
|
||||
}
|
||||
if not args.skip_old:
|
||||
results["endpoints"].append(
|
||||
run_endpoint("legacy", args.old_url, payload, args.requests, args.concurrency, args.timeout)
|
||||
)
|
||||
if not args.skip_new:
|
||||
results["endpoints"].append(
|
||||
run_endpoint("vllm", args.new_url, payload, args.requests, args.concurrency, args.timeout)
|
||||
)
|
||||
|
||||
output_path = Path(args.output)
|
||||
output_path.write_text(json.dumps(results, indent=2), encoding="utf-8")
|
||||
print(json.dumps(results, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
+29
@@ -0,0 +1,29 @@
|
||||
version: 1
|
||||
disable_existing_loggers: False
|
||||
formatters:
|
||||
simple:
|
||||
format: "%(asctime)s - %(module)s - %(funcName)s - line:%(lineno)d - %(levelname)s - %(message)s"
|
||||
handlers:
|
||||
console:
|
||||
class: logging.StreamHandler
|
||||
level: INFO
|
||||
formatter: simple
|
||||
info_file_handler:
|
||||
class: logging.handlers.TimedRotatingFileHandler
|
||||
filename: logs/info.log
|
||||
level: INFO
|
||||
formatter: simple
|
||||
encoding: utf8
|
||||
when: d
|
||||
interval: 1
|
||||
backupCount: 30
|
||||
error_file_handler:
|
||||
class: logging.handlers.RotatingFileHandler
|
||||
level: ERROR
|
||||
formatter: simple
|
||||
filename: logs/errors.log
|
||||
backupCount: 20
|
||||
encoding: utf8
|
||||
root:
|
||||
level: INFO
|
||||
handlers: [console, info_file_handler, error_file_handler]
|
||||
@@ -0,0 +1,25 @@
|
||||
from typing import Dict
|
||||
|
||||
import torch
|
||||
|
||||
from surya.common.predictor import BasePredictor
|
||||
from surya.detection import DetectionPredictor
|
||||
from surya.layout import LayoutPredictor
|
||||
from surya.logging import configure_logging
|
||||
from surya.ocr_error import OCRErrorPredictor
|
||||
from surya.recognition import RecognitionPredictor
|
||||
from surya.table_rec import TableRecPredictor
|
||||
|
||||
configure_logging()
|
||||
|
||||
|
||||
def load_predictors(
|
||||
device: str | torch.device | None = None, dtype: torch.dtype | str | None = None
|
||||
) -> Dict[str, BasePredictor]:
|
||||
return {
|
||||
"layout": LayoutPredictor(device=device, dtype=dtype),
|
||||
"ocr_error": OCRErrorPredictor(device=device, dtype=dtype),
|
||||
"recognition": RecognitionPredictor(device=device, dtype=dtype),
|
||||
"detection": DetectionPredictor(device=device, dtype=dtype),
|
||||
"table_rec": TableRecPredictor(device=device, dtype=dtype),
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
[project]
|
||||
name = "suya-ocr-api"
|
||||
version = "0.1.0"
|
||||
description = "FastAPI OCR service using Surya OCR with vLLM-backed inference."
|
||||
requires-python = ">=3.10,<4"
|
||||
dependencies = [
|
||||
"fastapi==0.115.7",
|
||||
"uvicorn==0.37.0",
|
||||
"python-multipart==0.0.20",
|
||||
"requests==2.32.4",
|
||||
"PyYAML==6.0.2",
|
||||
"pillow==10.4.0",
|
||||
"pydantic==2.11.7",
|
||||
"pydantic-settings==2.9.1",
|
||||
"pypdfium2==4.30.0",
|
||||
"pandas==2.3.0",
|
||||
"openai>=1.55.0,<2",
|
||||
"httpx>=0.27.0,<0.28",
|
||||
"beautifulsoup4>=4.12.0,<5",
|
||||
"platformdirs==4.3.8",
|
||||
"filelock==3.18.0",
|
||||
"python-dotenv==1.1.0",
|
||||
]
|
||||
|
||||
[tool.uv]
|
||||
package = false
|
||||
|
||||
[project.optional-dependencies]
|
||||
trtllm = [
|
||||
"flash-linear-attention==0.5.0",
|
||||
]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
pythonpath = ["."]
|
||||
testpaths = ["tests"]
|
||||
@@ -0,0 +1,4 @@
|
||||
# Extra deps for the quantization benchmark (not needed by the production image).
|
||||
llmcompressor>=0.3.0
|
||||
bitsandbytes>=0.43.0
|
||||
matplotlib>=3.8.0
|
||||
@@ -0,0 +1,16 @@
|
||||
fastapi==0.115.7
|
||||
uvicorn==0.37.0
|
||||
python-multipart==0.0.20
|
||||
requests==2.32.4
|
||||
PyYAML==6.0.2
|
||||
pillow==10.4.0
|
||||
pydantic==2.11.7
|
||||
pydantic-settings==2.9.1
|
||||
python-dotenv==1.1.0
|
||||
pypdfium2==4.30.0
|
||||
pandas==2.3.0
|
||||
platformdirs==4.3.8
|
||||
filelock==3.18.0
|
||||
beautifulsoup4>=4.12.0,<5
|
||||
openai>=1.55.0,<2
|
||||
httpx>=0.27.0,<0.28
|
||||
@@ -0,0 +1,54 @@
|
||||
"""Post one or more images to an OCR endpoint and save text_lines per image.
|
||||
|
||||
Used to freeze the accuracy baseline and to capture candidate outputs for CER.
|
||||
|
||||
Usage:
|
||||
python scripts/capture_ocr_text.py \
|
||||
--url http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/ \
|
||||
--out-dir baseline_outputs \
|
||||
temp_image_c3c56948-282e-453b-8fac-6c482243d1e5.jpg
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def _payload(image_path: Path) -> dict:
|
||||
suffix = image_path.suffix.lstrip(".").lower() or "png"
|
||||
return {
|
||||
"file": base64.b64encode(image_path.read_bytes()).decode("utf-8"),
|
||||
"type": "jpg" if suffix == "jpeg" else suffix,
|
||||
"skip_text_detection": False,
|
||||
"skip_table_detection": False,
|
||||
"recognize_math": False,
|
||||
"ocr_with_boxes": True,
|
||||
}
|
||||
|
||||
|
||||
def capture(url: str, image_path: Path, out_dir: Path, timeout: float) -> str:
|
||||
resp = requests.post(url, json=_payload(image_path), timeout=timeout)
|
||||
body = resp.json()
|
||||
text = (body.get("data") or {}).get("text_lines", "") if isinstance(body.get("data"), dict) else ""
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
(out_dir / (image_path.stem + ".txt")).write_text(text, encoding="utf-8")
|
||||
return text
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Capture OCR text_lines per image.")
|
||||
parser.add_argument("images", nargs="+", type=Path)
|
||||
parser.add_argument("--url", required=True)
|
||||
parser.add_argument("--out-dir", type=Path, required=True)
|
||||
parser.add_argument("--timeout", type=float, default=900)
|
||||
args = parser.parse_args()
|
||||
for image in args.images:
|
||||
capture(args.url, image, args.out_dir, args.timeout)
|
||||
print(f"captured {image.name} -> {args.out_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,59 @@
|
||||
"""Character error rate (CER) between a reference baseline and a candidate.
|
||||
|
||||
CER = levenshtein(reference, hypothesis) / len(reference).
|
||||
Used to gate work-reduction changes: keep a change only if CER <= 0.005.
|
||||
|
||||
Usage:
|
||||
python scripts/cer_divergence.py baseline_outputs/ candidate_outputs/
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _levenshtein(a: str, b: str) -> int:
|
||||
if a == b:
|
||||
return 0
|
||||
if not a:
|
||||
return len(b)
|
||||
if not b:
|
||||
return len(a)
|
||||
prev = list(range(len(b) + 1))
|
||||
for i, ca in enumerate(a, 1):
|
||||
cur = [i]
|
||||
for j, cb in enumerate(b, 1):
|
||||
cost = 0 if ca == cb else 1
|
||||
cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + cost))
|
||||
prev = cur
|
||||
return prev[-1]
|
||||
|
||||
|
||||
def cer(reference: str, hypothesis: str) -> float:
|
||||
if not reference:
|
||||
return 0.0 if not hypothesis else 1.0
|
||||
return _levenshtein(reference, hypothesis) / len(reference)
|
||||
|
||||
|
||||
def cer_over_dirs(baseline_dir: Path, candidate_dir: Path) -> dict:
|
||||
per_file = {}
|
||||
for ref_path in sorted(baseline_dir.glob("*.txt")):
|
||||
cand_path = candidate_dir / ref_path.name
|
||||
ref = ref_path.read_text(encoding="utf-8")
|
||||
hyp = cand_path.read_text(encoding="utf-8") if cand_path.exists() else ""
|
||||
per_file[ref_path.name] = cer(ref, hyp)
|
||||
mean = sum(per_file.values()) / len(per_file) if per_file else 0.0
|
||||
return {"mean_cer": mean, "max_cer": max(per_file.values(), default=0.0), "per_file": per_file}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="CER divergence between two output dirs.")
|
||||
parser.add_argument("baseline_dir", type=Path)
|
||||
parser.add_argument("candidate_dir", type=Path)
|
||||
args = parser.parse_args()
|
||||
print(json.dumps(cer_over_dirs(args.baseline_dir, args.candidate_dir), indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+18
@@ -0,0 +1,18 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
if ! command -v nvidia-ctk >/dev/null 2>&1; then
|
||||
echo "nvidia-ctk is not installed. Install NVIDIA Container Toolkit first." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
sudo nvidia-ctk runtime configure --runtime=docker --config=/etc/docker/daemon.json
|
||||
|
||||
if command -v systemctl >/dev/null 2>&1; then
|
||||
sudo systemctl restart docker
|
||||
else
|
||||
sudo service docker restart
|
||||
fi
|
||||
|
||||
docker info --format '{{json .Runtimes}}'
|
||||
docker run --rm --runtime nvidia --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi
|
||||
Executable
+23
@@ -0,0 +1,23 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
args=()
|
||||
skip_next=0
|
||||
|
||||
for arg in "$@"; do
|
||||
if [ "$skip_next" -eq 1 ]; then
|
||||
skip_next=0
|
||||
if [ "$arg" = "nvidia" ]; then
|
||||
continue
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ "$arg" = "--runtime" ]; then
|
||||
skip_next=1
|
||||
continue
|
||||
fi
|
||||
|
||||
args+=("$arg")
|
||||
done
|
||||
|
||||
exec /usr/bin/docker "${args[@]}"
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Parse surya_timing_summary lines from logs/info.log and aggregate spans.
|
||||
|
||||
Usage:
|
||||
python scripts/parse_timing.py logs/info.log
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
import json
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
_BATCH_RE = re.compile(r"batch_size=(\d+)")
|
||||
_REQID_RE = re.compile(r"request_id=(\S+)")
|
||||
|
||||
|
||||
def parse_line(line: str) -> Optional[Dict[str, Any]]:
|
||||
if "surya_timing_summary" not in line or "events=" not in line:
|
||||
return None
|
||||
events_str = line.split("events=", 1)[1].strip()
|
||||
try:
|
||||
events = ast.literal_eval(events_str)
|
||||
except (ValueError, SyntaxError):
|
||||
return None
|
||||
batch_match = _BATCH_RE.search(line)
|
||||
reqid_match = _REQID_RE.search(line)
|
||||
return {
|
||||
"request_id": reqid_match.group(1) if reqid_match else "-",
|
||||
"batch_size": int(batch_match.group(1)) if batch_match else 0,
|
||||
"events": events,
|
||||
}
|
||||
|
||||
|
||||
def aggregate(records: List[Dict[str, Any]]) -> Dict[str, Dict[str, float]]:
|
||||
stats: Dict[str, Dict[str, float]] = defaultdict(
|
||||
lambda: {"count": 0, "total_ms": 0.0, "mean_ms": 0.0, "total_tokens": 0}
|
||||
)
|
||||
for rec in records:
|
||||
for event in rec["events"]:
|
||||
s = stats[event["name"]]
|
||||
s["count"] += 1
|
||||
s["total_ms"] += event.get("duration_ms", 0.0)
|
||||
tokens = (event.get("metadata") or {}).get("token_count")
|
||||
if tokens:
|
||||
s["total_tokens"] += tokens
|
||||
for s in stats.values():
|
||||
s["mean_ms"] = round(s["total_ms"] / s["count"], 2) if s["count"] else 0.0
|
||||
s["total_ms"] = round(s["total_ms"], 2)
|
||||
return dict(stats)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Aggregate surya timing spans from a log file.")
|
||||
parser.add_argument("logfile", type=Path)
|
||||
args = parser.parse_args()
|
||||
records = [
|
||||
rec
|
||||
for rec in (parse_line(line) for line in args.logfile.read_text(encoding="utf-8").splitlines())
|
||||
if rec is not None
|
||||
]
|
||||
print(json.dumps({"records": len(records), "spans": aggregate(records)}, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,54 @@
|
||||
"""Sample a vLLM server's Prometheus /metrics and report peak in-flight sequences.
|
||||
|
||||
vLLM exposes `vllm:num_requests_running` (currently executing on GPU) and
|
||||
`vllm:num_requests_waiting` (queued). Peak running vs VLLM_MAX_NUM_SEQS tells us
|
||||
whether we are saturating or overflowing the GPU's sequence slots.
|
||||
|
||||
Run this in the background during a benchmark, then read its printed summary.
|
||||
|
||||
Usage:
|
||||
python scripts/poll_vllm_metrics.py \
|
||||
--url http://127.0.0.1:8000/metrics --interval 0.25 --duration 120
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
_RUNNING_RE = re.compile(r"^vllm:num_requests_running\S*\s+([0-9.]+)", re.MULTILINE)
|
||||
_WAITING_RE = re.compile(r"^vllm:num_requests_waiting\S*\s+([0-9.]+)", re.MULTILINE)
|
||||
|
||||
|
||||
def _scrape(url: str) -> tuple[float, float]:
|
||||
text = requests.get(url, timeout=5).text
|
||||
running = max((float(m) for m in _RUNNING_RE.findall(text)), default=0.0)
|
||||
waiting = max((float(m) for m in _WAITING_RE.findall(text)), default=0.0)
|
||||
return running, waiting
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Poll vLLM /metrics for peak in-flight sequences.")
|
||||
parser.add_argument("--url", default="http://127.0.0.1:8000/metrics")
|
||||
parser.add_argument("--interval", type=float, default=0.25)
|
||||
parser.add_argument("--duration", type=float, default=120)
|
||||
args = parser.parse_args()
|
||||
|
||||
peak_running = 0.0
|
||||
peak_waiting = 0.0
|
||||
deadline = time.perf_counter() + args.duration
|
||||
while time.perf_counter() < deadline:
|
||||
try:
|
||||
running, waiting = _scrape(args.url)
|
||||
peak_running = max(peak_running, running)
|
||||
peak_waiting = max(peak_waiting, waiting)
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(args.interval)
|
||||
print(f"peak_running={peak_running} peak_waiting={peak_waiting}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1 @@
|
||||
"""Quantization benchmark harness for the Surya-OCR-2 recognition VLM."""
|
||||
@@ -0,0 +1,59 @@
|
||||
"""Summary-row contract and rendering to CSV + markdown."""
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import io
|
||||
from typing import Any, Dict, List
|
||||
|
||||
SUMMARY_FIELDS = [
|
||||
"method",
|
||||
"status",
|
||||
"t4_deployable",
|
||||
"model_size_mb",
|
||||
"mean_latency_s",
|
||||
"p50_latency_s",
|
||||
"p95_latency_s",
|
||||
"throughput_rps",
|
||||
"decode_tok_s",
|
||||
"mean_cer",
|
||||
"max_cer",
|
||||
"mean_bbox_iou",
|
||||
"mean_missed_lines",
|
||||
"mean_extra_lines",
|
||||
"error",
|
||||
]
|
||||
|
||||
|
||||
def build_row(**kwargs: Any) -> Dict[str, Any]:
|
||||
unknown = set(kwargs) - set(SUMMARY_FIELDS)
|
||||
if unknown:
|
||||
raise KeyError(f"unknown summary field(s): {sorted(unknown)}")
|
||||
row = {field: None for field in SUMMARY_FIELDS}
|
||||
row.update(kwargs)
|
||||
return row
|
||||
|
||||
|
||||
def rows_to_csv(rows: List[Dict[str, Any]]) -> str:
|
||||
buf = io.StringIO()
|
||||
writer = csv.DictWriter(buf, fieldnames=SUMMARY_FIELDS)
|
||||
writer.writeheader()
|
||||
for row in rows:
|
||||
writer.writerow(row)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def _fmt(value: Any) -> str:
|
||||
if value is None:
|
||||
return ""
|
||||
if isinstance(value, float):
|
||||
return f"{value:.4f}"
|
||||
return str(value)
|
||||
|
||||
|
||||
def rows_to_markdown(rows: List[Dict[str, Any]]) -> str:
|
||||
header = "| " + " | ".join(SUMMARY_FIELDS) + " |"
|
||||
sep = "| " + " | ".join("---" for _ in SUMMARY_FIELDS) + " |"
|
||||
lines = [header, sep]
|
||||
for row in rows:
|
||||
lines.append("| " + " | ".join(_fmt(row[f]) for f in SUMMARY_FIELDS) + " |")
|
||||
return "\n".join(lines)
|
||||
@@ -0,0 +1,66 @@
|
||||
"""Greedy bbox IoU matching between reference and candidate page boxes."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Sequence
|
||||
|
||||
Box = Sequence[float]
|
||||
|
||||
|
||||
def iou(a: Box, b: Box) -> float:
|
||||
ax0, ay0, ax1, ay1 = a
|
||||
bx0, by0, bx1, by1 = b
|
||||
ix0, iy0 = max(ax0, bx0), max(ay0, by0)
|
||||
ix1, iy1 = min(ax1, bx1), min(ay1, by1)
|
||||
iw, ih = max(0.0, ix1 - ix0), max(0.0, iy1 - iy0)
|
||||
inter = iw * ih
|
||||
if inter <= 0:
|
||||
return 0.0
|
||||
area_a = max(0.0, ax1 - ax0) * max(0.0, ay1 - ay0)
|
||||
area_b = max(0.0, bx1 - bx0) * max(0.0, by1 - by0)
|
||||
union = area_a + area_b - inter
|
||||
return inter / union if union > 0 else 0.0
|
||||
|
||||
|
||||
def match_boxes(ref: List[Box], cand: List[Box], iou_threshold: float = 0.5) -> Dict[str, float]:
|
||||
used = [False] * len(cand)
|
||||
matched_ious: List[float] = []
|
||||
for r in ref:
|
||||
best_j, best_iou = -1, 0.0
|
||||
for j, c in enumerate(cand):
|
||||
if used[j]:
|
||||
continue
|
||||
cur = iou(r, c)
|
||||
if cur > best_iou:
|
||||
best_iou, best_j = cur, j
|
||||
if best_j >= 0 and best_iou >= iou_threshold:
|
||||
used[best_j] = True
|
||||
matched_ious.append(best_iou)
|
||||
matched = len(matched_ious)
|
||||
return {
|
||||
"matched": matched,
|
||||
"missed": len(ref) - matched,
|
||||
"extra": len(cand) - matched,
|
||||
"mean_matched_iou": sum(matched_ious) / matched if matched else 0.0,
|
||||
}
|
||||
|
||||
|
||||
def bbox_iou_over_dirs(ref_dir: Path, cand_dir: Path, iou_threshold: float = 0.5) -> Dict:
|
||||
per_file: Dict[str, Dict[str, float]] = {}
|
||||
for ref_path in sorted(ref_dir.glob("*.json")):
|
||||
ref_boxes = json.loads(ref_path.read_text(encoding="utf-8")).get("boxes", [])
|
||||
cand_path = cand_dir / ref_path.name
|
||||
cand_boxes = (
|
||||
json.loads(cand_path.read_text(encoding="utf-8")).get("boxes", [])
|
||||
if cand_path.exists()
|
||||
else []
|
||||
)
|
||||
per_file[ref_path.name] = match_boxes(ref_boxes, cand_boxes, iou_threshold)
|
||||
n = len(per_file) or 1
|
||||
return {
|
||||
"mean_bbox_iou": sum(v["mean_matched_iou"] for v in per_file.values()) / n,
|
||||
"mean_missed_lines": sum(v["missed"] for v in per_file.values()) / n,
|
||||
"mean_extra_lines": sum(v["extra"] for v in per_file.values()) / n,
|
||||
"per_file": per_file,
|
||||
}
|
||||
@@ -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
|
||||
@@ -0,0 +1,77 @@
|
||||
"""Capture per-page OCR text, bboxes, and latency from the vLLM OCR endpoint."""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def extract_capture(body: Dict[str, Any]) -> Dict[str, Any]:
|
||||
data = body.get("data") or {}
|
||||
if not isinstance(data, dict):
|
||||
data = {}
|
||||
blocks = ((data.get("ocr_text_json") or {}).get("blocks")) or []
|
||||
boxes = [list(b["bbox"]) for b in blocks if isinstance(b, dict) and b.get("bbox")]
|
||||
return {
|
||||
"text": data.get("text_lines", "") or "",
|
||||
"boxes": boxes,
|
||||
"elapsed_seconds": data.get("elapsed_seconds"),
|
||||
}
|
||||
|
||||
|
||||
def write_capture(cap: Dict[str, Any], out_dir: Path, stem: str) -> None:
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
(out_dir / f"{stem}.txt").write_text(cap["text"], encoding="utf-8")
|
||||
(out_dir / f"{stem}.json").write_text(
|
||||
json.dumps({"boxes": cap["boxes"], "elapsed_seconds": cap["elapsed_seconds"]}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
def _payload(image_path: Path) -> Dict[str, Any]:
|
||||
suffix = image_path.suffix.lstrip(".").lower() or "png"
|
||||
return {
|
||||
"file": base64.b64encode(image_path.read_bytes()).decode("utf-8"),
|
||||
"type": "jpg" if suffix == "jpeg" else suffix,
|
||||
"skip_text_detection": False,
|
||||
"skip_table_detection": False,
|
||||
"recognize_math": False,
|
||||
"ocr_with_boxes": True,
|
||||
}
|
||||
|
||||
|
||||
def capture_page(url: str, image_path: Path, out_dir: Path, timeout: float) -> Dict[str, Any]:
|
||||
resp = requests.post(url, json=_payload(image_path), timeout=timeout)
|
||||
cap = extract_capture(resp.json())
|
||||
if cap["elapsed_seconds"] is None:
|
||||
# Latency is a headline metric; a successful response with no timing means
|
||||
# the response contract changed. Surface it loudly instead of silently
|
||||
# dropping the page from the latency stats.
|
||||
print(
|
||||
f"WARNING: no elapsed_seconds in OCR response for {image_path.name}; "
|
||||
"latency for this page will be missing",
|
||||
file=sys.stderr,
|
||||
)
|
||||
write_capture(cap, out_dir, image_path.stem)
|
||||
return cap
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Capture OCR text+boxes+latency per image.")
|
||||
parser.add_argument("images", nargs="+", type=Path)
|
||||
parser.add_argument("--url", required=True)
|
||||
parser.add_argument("--out-dir", type=Path, required=True)
|
||||
parser.add_argument("--timeout", type=float, default=900)
|
||||
args = parser.parse_args()
|
||||
for image in args.images:
|
||||
cap = capture_page(args.url, image, args.out_dir, args.timeout)
|
||||
print(f"captured {image.name}: {len(cap['boxes'])} boxes, {cap['elapsed_seconds']}s")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Assemble the quantization benchmark summary from captured per-page outputs.
|
||||
|
||||
Unlike run_all's inline scoring (which globs the reference and penalizes any page
|
||||
a method did not capture), this scores **candidate-driven**: for each method it
|
||||
scores only the pages that method actually produced, against the matching
|
||||
reference page. That lets fast methods run the full eval set while slow/eager
|
||||
methods (e.g. bitsandbytes) run a representative subset, without the subset being
|
||||
unfairly penalized for missing pages.
|
||||
|
||||
Per-method status/reason and t4-deployability are passed in via a small spec so
|
||||
methods that could not be built or served are recorded as explicit failed rows.
|
||||
|
||||
Run inside the bench container: python3 -m scripts.quant.finalize
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from scripts.cer_divergence import cer
|
||||
from scripts.quant.aggregate import build_row, rows_to_csv, rows_to_markdown
|
||||
from scripts.quant.bbox_iou import match_boxes
|
||||
|
||||
WORK = Path("results/quant")
|
||||
REF = WORK / "reference"
|
||||
RESULTS = WORK / "results"
|
||||
|
||||
# Per-method outcome metadata for this run (A100, vLLM v0.22.1 container).
|
||||
# status "ok" methods are scored from their captured pages — if a method has no
|
||||
# captures on disk it falls back to a failed row with `fallback_error`. "failed"
|
||||
# methods record why outright. bf16 is the reference (CER 0 / IoU 1 by construction).
|
||||
#
|
||||
# int8/awq/gptq were built with llm-compressor main (0.12.x): the PyPI 0.11 line
|
||||
# pins transformers<5 and cannot load qwen3_5, but main supports transformers 5.x.
|
||||
# int8 is plain GPTQ W8A8 (SmoothQuant default mappings unresolvable for qwen3_5).
|
||||
SPECS: List[Dict[str, Any]] = [
|
||||
{"method": "bf16", "status": "ok", "t4_deployable": True, "is_reference": True},
|
||||
{"method": "fp8", "status": "failed", "t4_deployable": False,
|
||||
"error": "vLLM FP8 Marlin kernel: size_n=32 not divisible by tile_n_size=64 "
|
||||
"(model layer shape incompatible with fp8 dynamic)"},
|
||||
{"method": "int8", "status": "ok", "t4_deployable": True, "is_reference": False,
|
||||
"fallback_error": "checkpoint built (llm-compressor main, GPTQ W8A8) but vLLM "
|
||||
"failed to serve it — see /tmp/serve_int8.log"},
|
||||
{"method": "awq", "status": "ok", "t4_deployable": True, "is_reference": False,
|
||||
"fallback_error": "checkpoint built (llm-compressor main, AWQ W4A16) but vLLM "
|
||||
"failed to serve it — see /tmp/serve_awq.log"},
|
||||
{"method": "gptq", "status": "ok", "t4_deployable": True, "is_reference": False,
|
||||
"fallback_error": "checkpoint built (llm-compressor main, GPTQ W4A16) but vLLM "
|
||||
"failed to serve it — see /tmp/serve_gptq.log"},
|
||||
{"method": "bnb8", "status": "ok", "t4_deployable": True, "is_reference": False},
|
||||
{"method": "bnb4", "status": "ok", "t4_deployable": True, "is_reference": False},
|
||||
]
|
||||
|
||||
|
||||
def _load_page(d: Path, stem: str):
|
||||
j = json.loads((d / f"{stem}.json").read_text(encoding="utf-8"))
|
||||
txt = (d / f"{stem}.txt").read_text(encoding="utf-8") if (d / f"{stem}.txt").exists() else ""
|
||||
return j, txt
|
||||
|
||||
|
||||
def _score(method: str, is_reference: bool) -> Dict[str, Any]:
|
||||
cand_dir = REF if is_reference else RESULTS / method
|
||||
stems = sorted(p.stem for p in cand_dir.glob("*.json"))
|
||||
latencies, cers, ious, missed, extra = [], [], [], [], []
|
||||
for stem in stems:
|
||||
cj, ctxt = _load_page(cand_dir, stem)
|
||||
if cj.get("elapsed_seconds") is not None:
|
||||
latencies.append(cj["elapsed_seconds"])
|
||||
rj, rtxt = _load_page(REF, stem)
|
||||
cers.append(cer(rtxt, ctxt))
|
||||
m = match_boxes(rj.get("boxes", []), cj.get("boxes", []), iou_threshold=0.5)
|
||||
ious.append(m["mean_matched_iou"])
|
||||
missed.append(m["missed"])
|
||||
extra.append(m["extra"])
|
||||
latencies.sort()
|
||||
n = len(latencies)
|
||||
model_dir = WORK / "models" / method
|
||||
size = None
|
||||
if model_dir.exists():
|
||||
size = round(sum(f.stat().st_size for f in model_dir.rglob("*") if f.is_file()) / (1024 * 1024), 1)
|
||||
avg = lambda xs: round(sum(xs) / len(xs), 4) if xs else None
|
||||
return {
|
||||
"n_pages": len(stems),
|
||||
"mean_latency_s": round(sum(latencies) / n, 2) if n else None,
|
||||
"p50_latency_s": round(latencies[n // 2], 2) if n else None,
|
||||
"p95_latency_s": round(latencies[min(n - 1, int(n * 0.95))], 2) if n else None,
|
||||
"mean_cer": avg(cers),
|
||||
"max_cer": round(max(cers), 4) if cers else None,
|
||||
"mean_bbox_iou": avg(ious),
|
||||
"mean_missed_lines": avg(missed),
|
||||
"mean_extra_lines": avg(extra),
|
||||
"model_size_mb": size,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
rows = []
|
||||
for spec in SPECS:
|
||||
if spec["status"] != "ok":
|
||||
rows.append(build_row(method=spec["method"], status="failed",
|
||||
t4_deployable=spec["t4_deployable"], error=spec["error"]))
|
||||
continue
|
||||
cand_dir = REF if spec.get("is_reference") else RESULTS / spec["method"]
|
||||
if not any(cand_dir.glob("*.json")):
|
||||
rows.append(build_row(method=spec["method"], status="failed",
|
||||
t4_deployable=spec["t4_deployable"],
|
||||
error=spec.get("fallback_error", "no captures on disk")))
|
||||
continue
|
||||
s = _score(spec["method"], spec.get("is_reference", False))
|
||||
rows.append(build_row(
|
||||
method=spec["method"], status="ok", t4_deployable=spec["t4_deployable"],
|
||||
model_size_mb=s["model_size_mb"], mean_latency_s=s["mean_latency_s"],
|
||||
p50_latency_s=s["p50_latency_s"], p95_latency_s=s["p95_latency_s"],
|
||||
mean_cer=s["mean_cer"], max_cer=s["max_cer"], mean_bbox_iou=s["mean_bbox_iou"],
|
||||
mean_missed_lines=s["mean_missed_lines"], mean_extra_lines=s["mean_extra_lines"],
|
||||
))
|
||||
print(f"{spec['method']}: {s['n_pages']} pages, "
|
||||
f"mean_latency={s['mean_latency_s']}s, mean_cer={s['mean_cer']}, iou={s['mean_bbox_iou']}")
|
||||
(WORK / "summary.csv").write_text(rows_to_csv(rows), encoding="utf-8")
|
||||
(WORK / "summary.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
|
||||
(WORK / "summary.md").write_text(rows_to_markdown(rows), encoding="utf-8")
|
||||
print(f"wrote summary for {len(rows)} methods to {WORK}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,18 @@
|
||||
"""Load the fixed eval-set manifest (one image filename per line)."""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
|
||||
def load_manifest(manifest_path: Path, image_root: Path) -> List[Path]:
|
||||
paths: List[Path] = []
|
||||
for raw in manifest_path.read_text(encoding="utf-8").splitlines():
|
||||
name = raw.strip()
|
||||
if not name or name.startswith("#"):
|
||||
continue
|
||||
image = image_root / name
|
||||
if not image.exists():
|
||||
raise FileNotFoundError(f"manifest image not found: {name} (looked in {image_root})")
|
||||
paths.append(image)
|
||||
return paths
|
||||
@@ -0,0 +1,78 @@
|
||||
"""Speed-vs-accuracy Pareto scatter and per-method bar charts."""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg") # headless
|
||||
import matplotlib.pyplot as plt # noqa: E402
|
||||
|
||||
|
||||
def pareto_points(rows: List[Dict[str, Any]], x_key: str, y_key: str) -> List[Dict[str, Any]]:
|
||||
points = []
|
||||
for row in rows:
|
||||
if row.get("status") != "ok":
|
||||
continue
|
||||
if row.get(x_key) is None or row.get(y_key) is None:
|
||||
continue
|
||||
points.append({
|
||||
"method": row["method"],
|
||||
"x": row[x_key],
|
||||
"y": row[y_key],
|
||||
"t4_deployable": row.get("t4_deployable"),
|
||||
})
|
||||
return points
|
||||
|
||||
|
||||
def _scatter(points, x_label, y_label, title, out_path: Path) -> None:
|
||||
fig, ax = plt.subplots(figsize=(7, 5))
|
||||
for p in points:
|
||||
marker = "o" if p["t4_deployable"] else "x"
|
||||
ax.scatter(p["x"], p["y"], marker=marker, s=80)
|
||||
ax.annotate(p["method"], (p["x"], p["y"]), textcoords="offset points", xytext=(5, 5))
|
||||
ax.set_xlabel(x_label)
|
||||
ax.set_ylabel(y_label)
|
||||
ax.set_title(title + " (o = T4-deployable, x = A100-only)")
|
||||
fig.tight_layout()
|
||||
fig.savefig(out_path, dpi=120)
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def _bar(rows, metric_key, y_label, out_path: Path) -> None:
|
||||
ok = [r for r in rows if r.get("status") == "ok" and r.get(metric_key) is not None]
|
||||
fig, ax = plt.subplots(figsize=(7, 5))
|
||||
ax.bar([r["method"] for r in ok], [r[metric_key] for r in ok])
|
||||
ax.set_ylabel(y_label)
|
||||
ax.set_title(metric_key)
|
||||
fig.tight_layout()
|
||||
fig.savefig(out_path, dpi=120)
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def render_all(rows: List[Dict[str, Any]], out_dir: Path) -> List[Path]:
|
||||
out_dir = Path(out_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
written: List[Path] = []
|
||||
|
||||
p1 = out_dir / "pareto_latency_cer.png"
|
||||
_scatter(pareto_points(rows, "mean_latency_s", "mean_cer"),
|
||||
"mean latency (s) [lower=faster]", "CER vs BF16 [lower=better]",
|
||||
"Speed vs recognition accuracy", p1)
|
||||
written.append(p1)
|
||||
|
||||
p2 = out_dir / "pareto_latency_iou.png"
|
||||
_scatter(pareto_points(rows, "mean_latency_s", "mean_bbox_iou"),
|
||||
"mean latency (s) [lower=faster]", "bbox IoU vs BF16 [higher=better]",
|
||||
"Speed vs detection accuracy", p2)
|
||||
written.append(p2)
|
||||
|
||||
for metric, label in [("mean_latency_s", "mean latency (s)"),
|
||||
("mean_cer", "CER vs BF16"),
|
||||
("mean_bbox_iou", "bbox IoU vs BF16")]:
|
||||
bp = out_dir / f"bar_{metric}.png"
|
||||
_bar(rows, metric, label, bp)
|
||||
written.append(bp)
|
||||
|
||||
return written
|
||||
@@ -0,0 +1,50 @@
|
||||
"""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
|
||||
@@ -0,0 +1,111 @@
|
||||
"""Orchestrate the quantization sweep: baseline -> each method -> aggregate -> plot.
|
||||
|
||||
Each method is isolated: any exception becomes a status="failed" row so one bad
|
||||
method never aborts the sweep.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from scripts.quant.aggregate import build_row, rows_to_csv, rows_to_markdown
|
||||
from scripts.quant.bbox_iou import bbox_iou_over_dirs
|
||||
from scripts.quant.build_model import build_model
|
||||
from scripts.quant.capture import capture_page
|
||||
from scripts.quant.manifest import load_manifest
|
||||
from scripts.quant.recipes import METHOD_SPECS, method_names, vllm_serve_args
|
||||
from scripts.quant.serve import start_server, stop_server, wait_healthy
|
||||
|
||||
# Reuse the existing CER metric.
|
||||
from scripts.cer_divergence import cer_over_dirs
|
||||
|
||||
OCR_URL = "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/"
|
||||
|
||||
|
||||
def _dir_size_mb(path: Path) -> float:
|
||||
total = sum(f.stat().st_size for f in path.rglob("*") if f.is_file())
|
||||
return round(total / (1024 * 1024), 1)
|
||||
|
||||
|
||||
def _measure_method(method: str, base_model: str, work_dir: Path,
|
||||
eval_images: List[Path], reference_dir: Path) -> Dict[str, Any]:
|
||||
"""Build -> serve -> capture -> metrics for one method. Raises on any failure.
|
||||
|
||||
Assumes the API process (api.py) is already running and reads SURYA_INFERENCE_URL
|
||||
from the env to point at the per-method vLLM server (port set by the runbook).
|
||||
Returns a dict of measured summary fields (no method/status/t4 keys).
|
||||
"""
|
||||
model_dir = work_dir / "models" / method
|
||||
out_capture = work_dir / "results" / method
|
||||
port = 8000
|
||||
|
||||
built = build_model(method, base_model, model_dir, eval_images[: min(8, len(eval_images))])
|
||||
serve_args = vllm_serve_args(method, str(built), base_model, port)
|
||||
proc = start_server(serve_args, str(work_dir / f"vllm_{method}.log"))
|
||||
try:
|
||||
if not wait_healthy(port):
|
||||
raise RuntimeError(f"vLLM did not become healthy for {method}")
|
||||
latencies = []
|
||||
for image in eval_images:
|
||||
cap = capture_page(OCR_URL, image, out_capture, timeout=900)
|
||||
if cap["elapsed_seconds"] is not None:
|
||||
latencies.append(cap["elapsed_seconds"])
|
||||
cer = cer_over_dirs(reference_dir, out_capture)
|
||||
iou = bbox_iou_over_dirs(reference_dir, out_capture)
|
||||
finally:
|
||||
stop_server(port, proc)
|
||||
|
||||
latencies.sort()
|
||||
n = len(latencies)
|
||||
size = _dir_size_mb(built) if METHOD_SPECS[method]["kind"] in ("compressor", "bnb") else None
|
||||
return {
|
||||
"mean_latency_s": round(sum(latencies) / n, 3) if n else None,
|
||||
"p50_latency_s": latencies[n // 2] if n else None,
|
||||
"p95_latency_s": latencies[min(n - 1, int(n * 0.95))] if n else None,
|
||||
"mean_cer": cer["mean_cer"],
|
||||
"max_cer": cer["max_cer"],
|
||||
"mean_bbox_iou": iou["mean_bbox_iou"],
|
||||
"mean_missed_lines": iou["mean_missed_lines"],
|
||||
"mean_extra_lines": iou["mean_extra_lines"],
|
||||
"model_size_mb": size,
|
||||
}
|
||||
|
||||
|
||||
def run_method(method: str, base_model: str, work_dir: Path,
|
||||
eval_images: List[Path], reference_dir: Path) -> Dict[str, Any]:
|
||||
t4 = METHOD_SPECS[method]["t4_deployable"]
|
||||
try:
|
||||
metrics = _measure_method(method, base_model, work_dir, eval_images, reference_dir)
|
||||
return build_row(method=method, status="ok", t4_deployable=t4, **metrics)
|
||||
except Exception as exc: # feasibility probe: record and continue
|
||||
return build_row(method=method, status="failed", t4_deployable=t4, error=str(exc))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Run the quantization benchmark sweep.")
|
||||
parser.add_argument("--base-model", default="datalab-to/surya-ocr-2")
|
||||
parser.add_argument("--manifest", type=Path, default=Path("eval_set/manifest.txt"))
|
||||
parser.add_argument("--image-root", type=Path, default=Path("."))
|
||||
parser.add_argument("--work-dir", type=Path, default=Path("results/quant"))
|
||||
parser.add_argument("--methods", nargs="*", default=method_names())
|
||||
args = parser.parse_args()
|
||||
|
||||
eval_images = load_manifest(args.manifest, args.image_root)
|
||||
reference_dir = args.work_dir / "reference"
|
||||
args.work_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
rows = []
|
||||
for method in args.methods:
|
||||
print(f"=== {method} ===", flush=True)
|
||||
rows.append(run_method(method, args.base_model, args.work_dir, eval_images, reference_dir))
|
||||
|
||||
(args.work_dir / "summary.csv").write_text(rows_to_csv(rows), encoding="utf-8")
|
||||
(args.work_dir / "summary.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
|
||||
(args.work_dir / "summary.md").write_text(rows_to_markdown(rows), encoding="utf-8")
|
||||
print(f"wrote summary for {len(rows)} methods to {args.work_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,60 @@
|
||||
#!/usr/bin/env bash
|
||||
# Serve and score compressor-built checkpoints (int8/gptq/awq) over a page subset.
|
||||
# Each method: hardened server teardown -> serve checkpoint -> health/fail watch ->
|
||||
# capture N_PAGES through the OCR API. Run INSIDE the bench container:
|
||||
# bash scripts/quant/score_methods.sh int8 gptq awq
|
||||
# Requires api.py running on :5002 pointed at :8000.
|
||||
set -uo pipefail
|
||||
|
||||
export LD_LIBRARY_PATH="/usr/local/cuda/compat:${LD_LIBRARY_PATH:-}"
|
||||
N_PAGES="${N_PAGES:-3}"
|
||||
OCR_URL="http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/"
|
||||
LOG=/tmp/score_methods.log
|
||||
|
||||
mapfile -t ALL < <(sed '/^#/d;/^$/d' eval_set/manifest.txt)
|
||||
IMAGES=("${ALL[@]:0:$N_PAGES}")
|
||||
|
||||
stop_server() {
|
||||
local pid
|
||||
pid=$(ss -ltnp 2>/dev/null | grep ":8000 " | grep -oP 'pid=\K[0-9]+' | head -1)
|
||||
[ -n "$pid" ] && kill -9 "$pid" 2>/dev/null
|
||||
pkill -9 -f "vllm.entrypoints" 2>/dev/null
|
||||
# EngineCore outlives the API server and holds the GPU; kill compute procs and
|
||||
# poll until memory actually frees (kill returns before GPU release).
|
||||
for p in $(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null); do
|
||||
kill -9 "$p" 2>/dev/null
|
||||
done
|
||||
for _ in $(seq 1 30); do
|
||||
used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | head -1)
|
||||
[ "${used:-99999}" -lt 2000 ] && break
|
||||
sleep 2
|
||||
done
|
||||
}
|
||||
|
||||
for m in "$@"; do
|
||||
echo "=== scoring $m ===" | tee -a "$LOG"
|
||||
stop_server
|
||||
# VLLM_EXTRA_ARGS lets a caller add e.g. --dtype float16 (Exllama W4A16 kernel
|
||||
# on Ampere only supports float16 activations).
|
||||
nohup python3 -m vllm.entrypoints.openai.api_server --host 127.0.0.1 --port 8000 \
|
||||
--model "results/quant/models/$m" --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}' \
|
||||
${VLLM_EXTRA_ARGS:-} \
|
||||
> "/tmp/serve_${m}.log" 2>&1 &
|
||||
ok=0
|
||||
for _ in $(seq 1 84); do
|
||||
if curl -fs http://127.0.0.1:8000/health >/dev/null 2>&1; then ok=1; break; fi
|
||||
if grep -qiE "Engine core initialization failed" "/tmp/serve_${m}.log" 2>/dev/null; then break; fi
|
||||
sleep 5
|
||||
done
|
||||
if [ "$ok" -ne 1 ]; then
|
||||
echo "$m: FAILED to start" | tee -a "$LOG"
|
||||
continue
|
||||
fi
|
||||
python3 -m scripts.quant.capture --url "$OCR_URL" --out-dir "results/quant/results/$m" "${IMAGES[@]}" \
|
||||
>> "$LOG" 2>&1
|
||||
echo "$m: captured $(ls "results/quant/results/$m"/*.json 2>/dev/null | wc -l) pages" | tee -a "$LOG"
|
||||
done
|
||||
stop_server
|
||||
echo "SCORE_DONE" | tee -a "$LOG"
|
||||
@@ -0,0 +1,59 @@
|
||||
"""Launch a vLLM OpenAI server for one quant variant, health-check it, and tear
|
||||
it down by port (never by command-line match — that bit us in Approach A)."""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import time
|
||||
import urllib.request
|
||||
from typing import List, Optional
|
||||
|
||||
# Default seconds to wait for a server to become healthy. Overridable via env so a
|
||||
# method that fails to boot is recorded quickly instead of blocking the sweep.
|
||||
DEFAULT_HEALTH_TIMEOUT = float(os.environ.get("SURYA_QUANT_HEALTH_TIMEOUT", "900"))
|
||||
|
||||
|
||||
def health_url(port: int) -> str:
|
||||
return f"http://127.0.0.1:{port}/health"
|
||||
|
||||
|
||||
def parse_listening_pid(ss_output: str, port: int) -> Optional[int]:
|
||||
for line in ss_output.splitlines():
|
||||
if f":{port} " in line or line.rstrip().endswith(f":{port}"):
|
||||
m = re.search(r"pid=(\d+)", line)
|
||||
if m:
|
||||
return int(m.group(1))
|
||||
return None
|
||||
|
||||
|
||||
def wait_healthy(port: int, timeout: float = DEFAULT_HEALTH_TIMEOUT) -> bool:
|
||||
deadline = time.time() + timeout
|
||||
url = health_url(port)
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
with urllib.request.urlopen(url, timeout=2) as resp:
|
||||
if resp.status == 200:
|
||||
return True
|
||||
except Exception:
|
||||
time.sleep(3)
|
||||
return False
|
||||
|
||||
|
||||
def start_server(serve_args: List[str], log_path: str) -> subprocess.Popen:
|
||||
cmd = ["python", "-m", "vllm.entrypoints.openai.api_server", *serve_args]
|
||||
log = open(log_path, "w")
|
||||
return subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT)
|
||||
|
||||
|
||||
def stop_server(port: int, proc: Optional[subprocess.Popen] = None) -> None:
|
||||
if proc is not None:
|
||||
proc.terminate()
|
||||
try:
|
||||
proc.wait(timeout=30)
|
||||
except Exception:
|
||||
proc.kill()
|
||||
out = subprocess.run(["ss", "-ltnp"], capture_output=True, text=True).stdout
|
||||
pid = parse_listening_pid(out, port)
|
||||
if pid:
|
||||
subprocess.run(["kill", str(pid)])
|
||||
@@ -0,0 +1,76 @@
|
||||
#!/usr/bin/env bash
|
||||
# MTP (multi-token-prediction) speculative-decode sweep for the latency report.
|
||||
# Serves the BF16 model three ways — no speculation (baseline), MTP with 1
|
||||
# speculative token, MTP with 2 — and captures the same page subset through the
|
||||
# OCR API for each. The model ships 1 nextn-predict layer, so MTP=1 is the
|
||||
# expected-valid setting and MTP=2 is exploratory.
|
||||
#
|
||||
# Run INSIDE the bench container: bash scripts/quant/tune_mtp.sh
|
||||
# Requires: api.py running on :5002 pointed at :8000; compat libs on path.
|
||||
set -uo pipefail
|
||||
|
||||
MODEL="${MODEL:-datalab-to/surya-ocr-2}"
|
||||
OCR_URL="http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/"
|
||||
N_PAGES="${N_PAGES:-5}"
|
||||
OUT_ROOT="results/quant/mtp"
|
||||
LOG=/tmp/tune_mtp.log
|
||||
export LD_LIBRARY_PATH="/usr/local/cuda/compat:${LD_LIBRARY_PATH:-}"
|
||||
|
||||
mapfile -t ALL < <(sed '/^#/d;/^$/d' eval_set/manifest.txt)
|
||||
IMAGES=("${ALL[@]:0:$N_PAGES}")
|
||||
|
||||
# config_name | extra vLLM args. The speculative-config JSON is written compact
|
||||
# (no spaces) so it survives word-splitting as a single argv token.
|
||||
CONFIGS=(
|
||||
"baseline|--enable-prefix-caching"
|
||||
"mtp1|--enable-prefix-caching --speculative-config {\"method\":\"mtp\",\"num_speculative_tokens\":1}"
|
||||
"mtp2|--enable-prefix-caching --speculative-config {\"method\":\"mtp\",\"num_speculative_tokens\":2}"
|
||||
)
|
||||
|
||||
base_args() {
|
||||
echo "--host 127.0.0.1 --port 8000 --model $MODEL --served-model-name $MODEL \
|
||||
--max-model-len 18000 --max-num-seqs 16 --gpu-memory-utilization 0.85 --no-enforce-eager \
|
||||
--mm-processor-kwargs {\"min_pixels\":3136,\"max_pixels\":6291456}"
|
||||
}
|
||||
|
||||
stop_server() {
|
||||
local pid
|
||||
pid=$(ss -ltnp 2>/dev/null | grep ":8000 " | grep -oP 'pid=\K[0-9]+' | head -1)
|
||||
[ -n "$pid" ] && kill -9 "$pid" 2>/dev/null
|
||||
pkill -9 -f "vllm.entrypoints" 2>/dev/null
|
||||
# EngineCore outlives the API server and holds the GPU; kill every remaining
|
||||
# CUDA compute process, then poll until the memory is actually freed.
|
||||
for p in $(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null); do
|
||||
kill -9 "$p" 2>/dev/null
|
||||
done
|
||||
for _ in $(seq 1 30); do
|
||||
used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | head -1)
|
||||
[ "${used:-99999}" -lt 2000 ] && break
|
||||
sleep 2
|
||||
done
|
||||
}
|
||||
|
||||
for entry in "${CONFIGS[@]}"; do
|
||||
name="${entry%%|*}"
|
||||
extra="${entry#*|}"
|
||||
echo "=== mtp config: $name ($extra) ===" | tee -a "$LOG"
|
||||
stop_server
|
||||
# shellcheck disable=SC2046
|
||||
nohup python3 -m vllm.entrypoints.openai.api_server $(base_args) $extra \
|
||||
> "/tmp/mtp_${name}.log" 2>&1 &
|
||||
ok=0
|
||||
for _ in $(seq 1 96); do
|
||||
if curl -fs http://127.0.0.1:8000/health >/dev/null 2>&1; then ok=1; break; fi
|
||||
if grep -qiE "Engine core initialization failed|ValueError|RuntimeError|Traceback" "/tmp/mtp_${name}.log" 2>/dev/null; then break; fi
|
||||
sleep 5
|
||||
done
|
||||
if [ "$ok" -ne 1 ]; then
|
||||
echo "$name: FAILED to start (see /tmp/mtp_${name}.log)" | tee -a "$LOG"
|
||||
continue
|
||||
fi
|
||||
python3 -m scripts.quant.capture --url "$OCR_URL" --out-dir "$OUT_ROOT/$name" "${IMAGES[@]}" \
|
||||
>> "$LOG" 2>&1
|
||||
echo "$name: captured $(ls "$OUT_ROOT/$name"/*.json 2>/dev/null | wc -l) pages" | tee -a "$LOG"
|
||||
done
|
||||
stop_server
|
||||
echo "MTP_DONE" | tee -a "$LOG"
|
||||
@@ -0,0 +1,77 @@
|
||||
#!/usr/bin/env bash
|
||||
# vLLM serving-parameter tuning sweep for the OCR latency report.
|
||||
# Serves the BF16 model under several vLLM flag combinations (CUDA graph vs eager,
|
||||
# prefix caching on/off, chunked prefill on/off), captures a fixed page subset
|
||||
# through the OCR API for each, and records per-config latency.
|
||||
#
|
||||
# Run INSIDE the bench container: bash scripts/quant/tune_vllm.sh
|
||||
# Requires: api.py already running on :5002 pointed at :8000; compat libs on path.
|
||||
set -uo pipefail
|
||||
|
||||
MODEL="${MODEL:-datalab-to/surya-ocr-2}"
|
||||
OCR_URL="http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/"
|
||||
N_PAGES="${N_PAGES:-3}"
|
||||
OUT_ROOT="results/quant/tuning"
|
||||
LOG=/tmp/tune_vllm.log
|
||||
export LD_LIBRARY_PATH="/usr/local/cuda/compat:${LD_LIBRARY_PATH:-}"
|
||||
|
||||
mapfile -t ALL < <(sed '/^#/d;/^$/d' eval_set/manifest.txt)
|
||||
IMAGES=("${ALL[@]:0:$N_PAGES}")
|
||||
|
||||
# config_name | extra vLLM args
|
||||
CONFIGS=(
|
||||
"default|--enable-prefix-caching"
|
||||
"eager|--enable-prefix-caching --enforce-eager"
|
||||
"no_prefix_cache|--no-enable-prefix-caching"
|
||||
"no_chunked_prefill|--enable-prefix-caching --no-enable-chunked-prefill"
|
||||
)
|
||||
|
||||
base_args() {
|
||||
echo "--host 127.0.0.1 --port 8000 --model $MODEL --served-model-name $MODEL \
|
||||
--max-model-len 18000 --max-num-seqs 16 --gpu-memory-utilization 0.85 \
|
||||
--mm-processor-kwargs {\"min_pixels\":3136,\"max_pixels\":6291456}"
|
||||
}
|
||||
|
||||
stop_server() {
|
||||
local pid
|
||||
pid=$(ss -ltnp 2>/dev/null | grep ":8000 " | grep -oP 'pid=\K[0-9]+' | head -1)
|
||||
[ -n "$pid" ] && kill -9 "$pid" 2>/dev/null
|
||||
pkill -9 -f "vllm.entrypoints" 2>/dev/null
|
||||
# The EngineCore worker holds the GPU and outlives the API server; kill every
|
||||
# remaining CUDA compute process, then poll until the memory is actually freed
|
||||
# (kill returns immediately, GPU release lags).
|
||||
for p in $(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null); do
|
||||
kill -9 "$p" 2>/dev/null
|
||||
done
|
||||
for _ in $(seq 1 30); do
|
||||
used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | head -1)
|
||||
[ "${used:-99999}" -lt 2000 ] && break
|
||||
sleep 2
|
||||
done
|
||||
}
|
||||
|
||||
for entry in "${CONFIGS[@]}"; do
|
||||
name="${entry%%|*}"
|
||||
extra="${entry#*|}"
|
||||
echo "=== tuning config: $name ($extra) ===" | tee -a "$LOG"
|
||||
stop_server
|
||||
# shellcheck disable=SC2046
|
||||
nohup python3 -m vllm.entrypoints.openai.api_server $(base_args) $extra \
|
||||
> "/tmp/tune_${name}.log" 2>&1 &
|
||||
# wait for health or crash (max 7 min)
|
||||
ok=0
|
||||
for _ in $(seq 1 84); do
|
||||
if curl -fs http://127.0.0.1:8000/health >/dev/null 2>&1; then ok=1; break; fi
|
||||
if grep -qiE "Engine core initialization failed|RuntimeError" "/tmp/tune_${name}.log" 2>/dev/null; then break; fi
|
||||
sleep 5
|
||||
done
|
||||
if [ "$ok" -ne 1 ]; then
|
||||
echo "$name: FAILED to start" | tee -a "$LOG"
|
||||
continue
|
||||
fi
|
||||
python3 -m scripts.quant.capture --url "$OCR_URL" --out-dir "$OUT_ROOT/$name" "${IMAGES[@]}" \
|
||||
>> "$LOG" 2>&1
|
||||
echo "$name: captured $(ls "$OUT_ROOT/$name"/*.json 2>/dev/null | wc -l) pages" | tee -a "$LOG"
|
||||
done
|
||||
stop_server
|
||||
echo "TUNE_DONE" | tee -a "$LOG"
|
||||
Executable
+77
@@ -0,0 +1,77 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
export SURYA_MODEL_CHECKPOINT="${SURYA_MODEL_CHECKPOINT:-datalab-to/surya-ocr-2}"
|
||||
export SURYA_INFERENCE_BACKEND="${SURYA_INFERENCE_BACKEND:-vllm}"
|
||||
export SURYA_INFERENCE_URL="${SURYA_INFERENCE_URL:-http://127.0.0.1:8000/v1}"
|
||||
export SURYA_INFERENCE_AUTOSTART="${SURYA_INFERENCE_AUTOSTART:-false}"
|
||||
export SURYA_INFERENCE_PARALLEL="${SURYA_INFERENCE_PARALLEL:-8}"
|
||||
export SURYA_INFERENCE_LOGPROBS="${SURYA_INFERENCE_LOGPROBS:-false}"
|
||||
export SURYA_INFERENCE_MAX_RETRIES="${SURYA_INFERENCE_MAX_RETRIES:-1}"
|
||||
# Keep in sync with VLLM_MAX_NUM_SEQS so the batcher fills every vLLM sequence
|
||||
# slot (see docs/diagnosis_baseline.md). Overcommitting past it only queues.
|
||||
export SURYA_INFERENCE_MAX_INFLIGHT="${SURYA_INFERENCE_MAX_INFLIGHT:-16}"
|
||||
export SURYA_MAX_TOKENS_FULL_PAGE="${SURYA_MAX_TOKENS_FULL_PAGE:-6144}"
|
||||
export SURYA_MAX_BLOCKS_PER_PAGE="${SURYA_MAX_BLOCKS_PER_PAGE:-80}"
|
||||
export SUYA_OCR_MODE="${SUYA_OCR_MODE:-block}"
|
||||
export SUYA_MAX_BATCH_SIZE="${SUYA_MAX_BATCH_SIZE:-8}"
|
||||
export SUYA_BATCH_WAIT_MS="${SUYA_BATCH_WAIT_MS:-25}"
|
||||
export SUYA_VLLM_IMAGE_FORMAT="${SUYA_VLLM_IMAGE_FORMAT:-JPEG}"
|
||||
export SUYA_VLLM_JPEG_QUALITY="${SUYA_VLLM_JPEG_QUALITY:-92}"
|
||||
export VLLM_DTYPE="${VLLM_DTYPE:-float16}"
|
||||
export VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.85}"
|
||||
export VLLM_MAX_MODEL_LEN="${VLLM_MAX_MODEL_LEN:-18000}"
|
||||
export VLLM_MAX_NUM_SEQS="${VLLM_MAX_NUM_SEQS:-16}"
|
||||
export VLLM_MAX_BATCHED_TOKENS="${VLLM_MAX_BATCHED_TOKENS:-4096}"
|
||||
|
||||
# vLLM serving parameters are tuned for latency — see
|
||||
# docs/quantization_benchmark_results.md §5 (A100, BF16). Three flags carry the
|
||||
# win and MUST stay on; do not pass their negations via VLLM_EXTRA_ARGS:
|
||||
# --no-enforce-eager CUDA-graph capture. The single biggest lever: eager
|
||||
# mode (--enforce-eager) measured ~9.5x slower
|
||||
# (48.97s vs 5.15s/page). Set explicitly so a future
|
||||
# vLLM default flip can't silently disable graphs.
|
||||
# --enable-prefix-caching ~15% win; OCR prompts share a long fixed prefix.
|
||||
# chunked prefill (on by default) is REQUIRED — the qwen3_5 mamba/SSM cache
|
||||
# fails engine init with --no-enable-chunked-prefill.
|
||||
python3 -m vllm.entrypoints.openai.api_server \
|
||||
--host 127.0.0.1 \
|
||||
--port 8000 \
|
||||
--model "${SURYA_MODEL_CHECKPOINT}" \
|
||||
--served-model-name "${SURYA_MODEL_CHECKPOINT}" \
|
||||
--dtype "${VLLM_DTYPE}" \
|
||||
--max-model-len "${VLLM_MAX_MODEL_LEN}" \
|
||||
--max-num-seqs "${VLLM_MAX_NUM_SEQS}" \
|
||||
--max-num-batched-tokens "${VLLM_MAX_BATCHED_TOKENS}" \
|
||||
--gpu-memory-utilization "${VLLM_GPU_MEMORY_UTILIZATION}" \
|
||||
--no-enforce-eager \
|
||||
--enable-prefix-caching \
|
||||
--mm-processor-kwargs '{"min_pixels":3136,"max_pixels":6291456}' \
|
||||
${VLLM_EXTRA_ARGS:-} &
|
||||
VLLM_PID=$!
|
||||
|
||||
cleanup() {
|
||||
kill "${VLLM_PID}" 2>/dev/null || true
|
||||
}
|
||||
trap cleanup EXIT INT TERM
|
||||
|
||||
python3 - <<'PY'
|
||||
import os
|
||||
import time
|
||||
import urllib.request
|
||||
|
||||
base = os.environ.get("SURYA_INFERENCE_URL", "http://127.0.0.1:8000/v1")
|
||||
health = base[:-3] + "/health" if base.endswith("/v1") else base.rstrip("/") + "/health"
|
||||
deadline = time.time() + float(os.environ.get("SURYA_INFERENCE_STARTUP_TIMEOUT", "900"))
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
with urllib.request.urlopen(health, timeout=2) as response:
|
||||
if response.status == 200:
|
||||
print(f"vLLM health check passed: {health}", flush=True)
|
||||
raise SystemExit(0)
|
||||
except Exception:
|
||||
time.sleep(2)
|
||||
raise SystemExit(f"vLLM did not become healthy: {health}")
|
||||
PY
|
||||
|
||||
python3 api.py
|
||||
@@ -0,0 +1,3 @@
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
"""Pixel-content heuristics for detecting blank or near-uniform image regions.
|
||||
|
||||
Used by both the layout predictor (drop hallucinated layout blocks over empty
|
||||
space) and the recognition predictor (drop hallucinated text blocks from
|
||||
full-page OCR, decide whether an empty full-page output is a correct blank-page
|
||||
read or a failure).
|
||||
|
||||
Two signals, combined:
|
||||
* near-white fraction — most pixels have every RGB channel above a threshold
|
||||
* pixel-value standard deviation — the region is essentially one color
|
||||
(catches uniform-color fills that the white check misses)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
# Per-channel value at/above which a pixel is considered "near-white".
|
||||
# Tolerates the small noise typical of PDF renders at 96 DPI.
|
||||
BLANK_WHITE_THRESHOLD = 245
|
||||
# Fraction of pixels that must be near-white for a region to count as blank.
|
||||
BLANK_PIXEL_FRACTION = 0.99
|
||||
# Pixel-value std below which a region is "essentially one color" regardless
|
||||
# of what that color is (catches solid-fill rectangles, dark banners, etc.).
|
||||
UNIFORM_COLOR_STD = 8.0
|
||||
|
||||
|
||||
def near_white_fraction(
|
||||
image: Image.Image, white_threshold: int = BLANK_WHITE_THRESHOLD
|
||||
) -> float:
|
||||
"""Fraction of pixels where every RGB channel ≥ ``white_threshold``."""
|
||||
arr = np.asarray(image.convert("RGB"))
|
||||
if arr.size == 0:
|
||||
return 0.0
|
||||
return float(np.all(arr >= white_threshold, axis=-1).mean())
|
||||
|
||||
|
||||
def is_blank_region(
|
||||
image: Image.Image,
|
||||
*,
|
||||
white_threshold: int = BLANK_WHITE_THRESHOLD,
|
||||
blank_pixel_fraction: float = BLANK_PIXEL_FRACTION,
|
||||
uniform_color_std: float = UNIFORM_COLOR_STD,
|
||||
) -> bool:
|
||||
"""True iff the image is essentially blank — either mostly near-white or
|
||||
near-uniform color. Use this on a per-block crop or a whole page.
|
||||
|
||||
Returns False for empty (0-pixel) crops so callers don't accidentally
|
||||
treat a degenerate bbox as blank.
|
||||
"""
|
||||
arr = np.asarray(image.convert("RGB"))
|
||||
if arr.size == 0:
|
||||
return False
|
||||
if np.all(arr >= white_threshold, axis=-1).mean() > blank_pixel_fraction:
|
||||
return True
|
||||
# Per-channel std — a uniform solid color (e.g., red banner with RGB=(200,50,50))
|
||||
# has each channel constant across pixels, but mixing channels inflates the
|
||||
# aggregate std. Check each channel independently.
|
||||
per_channel_std = arr.reshape(-1, arr.shape[-1]).std(axis=0)
|
||||
if float(per_channel_std.max()) < uniform_color_std:
|
||||
return True
|
||||
return False
|
||||
@@ -0,0 +1,25 @@
|
||||
from typing import Optional, Any
|
||||
|
||||
import torch
|
||||
|
||||
from surya.settings import settings
|
||||
|
||||
|
||||
class ModelLoader:
|
||||
def __init__(self, checkpoint: Optional[str] = None):
|
||||
self.checkpoint = checkpoint
|
||||
|
||||
def model(
|
||||
self,
|
||||
device: torch.device | str | None = settings.TORCH_DEVICE_MODEL,
|
||||
dtype: Optional[torch.dtype | str] = settings.MODEL_DTYPE,
|
||||
attention_implementation: Optional[str] = None,
|
||||
) -> Any:
|
||||
raise NotImplementedError()
|
||||
|
||||
def processor(
|
||||
self,
|
||||
device: torch.device | str | None = settings.TORCH_DEVICE_MODEL,
|
||||
dtype: Optional[torch.dtype | str] = settings.MODEL_DTYPE,
|
||||
) -> Any:
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,122 @@
|
||||
import copy
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
from pydantic import BaseModel, field_validator, computed_field
|
||||
import numbers
|
||||
|
||||
|
||||
class PolygonBox(BaseModel):
|
||||
polygon: List[List[float]]
|
||||
confidence: Optional[float] = None
|
||||
|
||||
@field_validator("polygon", mode="before")
|
||||
@classmethod
|
||||
def convert_bbox_to_polygon(cls, value):
|
||||
if isinstance(value, (list, tuple)) and len(value) == 4:
|
||||
if all(isinstance(x, numbers.Number) for x in value):
|
||||
value = [float(v) for v in value]
|
||||
x_min, y_min, x_max, y_max = value
|
||||
polygon = [
|
||||
[x_min, y_min],
|
||||
[x_max, y_min],
|
||||
[x_max, y_max],
|
||||
[x_min, y_max],
|
||||
]
|
||||
return polygon
|
||||
elif all(
|
||||
isinstance(point, (list, tuple)) and len(point) == 2 for point in value
|
||||
):
|
||||
value = [[float(v) for v in point] for point in value]
|
||||
return value
|
||||
elif isinstance(value, np.ndarray):
|
||||
if value.shape == (4, 2):
|
||||
return value.tolist()
|
||||
|
||||
raise ValueError(
|
||||
f"Input must be either a bbox [x_min, y_min, x_max, y_max] or a polygon with 4 corners [(x,y), (x,y), (x,y), (x,y)]. All values must be numeric. You passed {value} of type {type(value)}. The first value is of type {type(value[0])}."
|
||||
)
|
||||
|
||||
@property
|
||||
def height(self):
|
||||
return self.bbox[3] - self.bbox[1]
|
||||
|
||||
@property
|
||||
def width(self):
|
||||
return self.bbox[2] - self.bbox[0]
|
||||
|
||||
@property
|
||||
def area(self):
|
||||
return self.width * self.height
|
||||
|
||||
@computed_field
|
||||
@property
|
||||
def bbox(self) -> List[float]:
|
||||
x_coords = [point[0] for point in self.polygon]
|
||||
y_coords = [point[1] for point in self.polygon]
|
||||
return [min(x_coords), min(y_coords), max(x_coords), max(y_coords)]
|
||||
|
||||
def rescale(self, processor_size, image_size):
|
||||
# Point is in x, y format
|
||||
page_width, page_height = processor_size
|
||||
|
||||
img_width, img_height = image_size
|
||||
width_scaler = img_width / page_width
|
||||
height_scaler = img_height / page_height
|
||||
|
||||
for corner in self.polygon:
|
||||
corner[0] = int(corner[0] * width_scaler)
|
||||
corner[1] = int(corner[1] * height_scaler)
|
||||
|
||||
def round(self, divisor):
|
||||
for corner in self.polygon:
|
||||
corner[0] = int(corner[0] / divisor) * divisor
|
||||
corner[1] = int(corner[1] / divisor) * divisor
|
||||
|
||||
def fit_to_bounds(self, bounds):
|
||||
new_corners = copy.deepcopy(self.polygon)
|
||||
for corner in new_corners:
|
||||
corner[0] = max(min(corner[0], bounds[2]), bounds[0])
|
||||
corner[1] = max(min(corner[1], bounds[3]), bounds[1])
|
||||
self.polygon = new_corners
|
||||
|
||||
def expand(self, x_margin: float, y_margin: float):
|
||||
new_polygon = []
|
||||
x_margin = x_margin * self.width
|
||||
y_margin = y_margin * self.height
|
||||
for idx, poly in enumerate(self.polygon):
|
||||
if idx == 0:
|
||||
new_polygon.append([int(poly[0] - x_margin), int(poly[1] - y_margin)])
|
||||
elif idx == 1:
|
||||
new_polygon.append([int(poly[0] + x_margin), int(poly[1] - y_margin)])
|
||||
elif idx == 2:
|
||||
new_polygon.append([int(poly[0] + x_margin), int(poly[1] + y_margin)])
|
||||
elif idx == 3:
|
||||
new_polygon.append([int(poly[0] - x_margin), int(poly[1] + y_margin)])
|
||||
self.polygon = new_polygon
|
||||
|
||||
def intersection_area(self, other, x_margin=0, y_margin=0):
|
||||
x_overlap = self.x_overlap(other, x_margin)
|
||||
y_overlap = self.y_overlap(other, y_margin)
|
||||
return x_overlap * y_overlap
|
||||
|
||||
def x_overlap(self, other, x_margin=0):
|
||||
return max(
|
||||
0,
|
||||
min(self.bbox[2] + x_margin, other.bbox[2] + x_margin)
|
||||
- max(self.bbox[0] - x_margin, other.bbox[0] - x_margin),
|
||||
)
|
||||
|
||||
def y_overlap(self, other, y_margin=0):
|
||||
return max(
|
||||
0,
|
||||
min(self.bbox[3] + y_margin, other.bbox[3] + y_margin)
|
||||
- max(self.bbox[1] - y_margin, other.bbox[1] - y_margin),
|
||||
)
|
||||
|
||||
@property
|
||||
def center(self):
|
||||
return [(self.bbox[0] + self.bbox[2]) / 2, (self.bbox[1] + self.bbox[3]) / 2]
|
||||
|
||||
def __hash__(self):
|
||||
return hash(tuple(self.bbox))
|
||||
@@ -0,0 +1,57 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from surya.common.load import ModelLoader
|
||||
from surya.settings import settings
|
||||
|
||||
|
||||
class BasePredictor:
|
||||
model_loader_cls = ModelLoader
|
||||
batch_size: Optional[int] = None
|
||||
default_batch_sizes = {"cpu": 1, "mps": 1, "cuda": 1}
|
||||
torch_dtype = settings.MODEL_DTYPE
|
||||
|
||||
@property
|
||||
def disable_tqdm(self) -> bool:
|
||||
return self._disable_tqdm
|
||||
|
||||
@disable_tqdm.setter
|
||||
def disable_tqdm(self, value: bool) -> None:
|
||||
self._disable_tqdm = bool(value)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
checkpoint: Optional[str] = None,
|
||||
device: torch.device | str | None = settings.TORCH_DEVICE_MODEL,
|
||||
dtype: Optional[torch.dtype | str] = None,
|
||||
attention_implementation: Optional[str] = None,
|
||||
):
|
||||
if dtype is None:
|
||||
dtype = self.torch_dtype
|
||||
|
||||
loader = self.model_loader_cls(checkpoint)
|
||||
self.model = loader.model(device, dtype, attention_implementation)
|
||||
self.processor = loader.processor()
|
||||
self._disable_tqdm = settings.DISABLE_TQDM
|
||||
|
||||
def to(self, device_dtype: torch.device | str | None = None):
|
||||
if hasattr(self, "model") and self.model:
|
||||
self.model.to(device_dtype)
|
||||
return
|
||||
# Predictors that don't own a torch model (e.g. VLM-backed predictors that
|
||||
# rely on an external server) treat .to() as a no-op.
|
||||
if hasattr(self, "manager") and self.manager is not None:
|
||||
return
|
||||
raise ValueError("Model not loaded")
|
||||
|
||||
def get_batch_size(self):
|
||||
batch_size = self.batch_size
|
||||
if batch_size is None:
|
||||
batch_size = self.default_batch_sizes["cpu"]
|
||||
if settings.TORCH_DEVICE_MODEL in self.default_batch_sizes:
|
||||
batch_size = self.default_batch_sizes[settings.TORCH_DEVICE_MODEL]
|
||||
return batch_size
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,22 @@
|
||||
from typing import Optional
|
||||
|
||||
from transformers import PreTrainedModel
|
||||
from transformers.utils import is_flash_attn_2_available
|
||||
|
||||
|
||||
class SuryaPreTrainedModel(PreTrainedModel):
|
||||
# No-op if we pass attention, so we can set attention however we want in the config
|
||||
def _check_and_adjust_attn_implementation(
|
||||
self, attn_implementation: Optional[str], **kwargs
|
||||
):
|
||||
if attn_implementation is None:
|
||||
try:
|
||||
self._sdpa_can_dispatch(True)
|
||||
attn_implementation = "sdpa"
|
||||
except (ValueError, ImportError):
|
||||
attn_implementation = "eager"
|
||||
|
||||
if self._supports_flash_attn and is_flash_attn_2_available():
|
||||
attn_implementation = "flash_attention_2"
|
||||
|
||||
return attn_implementation
|
||||
@@ -0,0 +1,182 @@
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
from tqdm import tqdm
|
||||
|
||||
from surya.logging import get_logger
|
||||
from surya.settings import settings
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
# Lock file expiration time in seconds (10 minutes)
|
||||
LOCK_EXPIRATION = 600
|
||||
|
||||
|
||||
def join_urls(url1: str, url2: str):
|
||||
url1 = url1.rstrip("/")
|
||||
url2 = url2.lstrip("/")
|
||||
return f"{url1}/{url2}"
|
||||
|
||||
|
||||
def get_model_name(pretrained_model_name_or_path: str):
|
||||
return pretrained_model_name_or_path.split("/")[0]
|
||||
|
||||
|
||||
def download_file(remote_path: str, local_path: str, chunk_size: int = 1024 * 1024):
|
||||
local_path = Path(local_path)
|
||||
try:
|
||||
response = requests.get(remote_path, stream=True, allow_redirects=True)
|
||||
response.raise_for_status() # Raise an exception for bad status codes
|
||||
|
||||
# Get file size from headers for progress bar
|
||||
total_size = int(response.headers.get('content-length', 0))
|
||||
|
||||
# Create progress bar with file name and size info
|
||||
filename = local_path.name
|
||||
pbar = tqdm(
|
||||
total=total_size,
|
||||
unit='B',
|
||||
unit_scale=True,
|
||||
unit_divisor=1024,
|
||||
desc=f"Downloading {filename}",
|
||||
miniters=1
|
||||
)
|
||||
|
||||
with open(local_path, "wb") as f:
|
||||
downloaded = 0
|
||||
for chunk in response.iter_content(chunk_size=chunk_size):
|
||||
if chunk:
|
||||
f.write(chunk)
|
||||
downloaded += len(chunk)
|
||||
pbar.update(len(chunk))
|
||||
|
||||
pbar.close()
|
||||
return local_path
|
||||
except Exception as e:
|
||||
if local_path.exists():
|
||||
local_path.unlink()
|
||||
logger.error(f"Download error for file {remote_path}: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
def check_manifest(local_dir: str):
|
||||
local_dir = Path(local_dir)
|
||||
manifest_path = local_dir / "manifest.json"
|
||||
if not os.path.exists(manifest_path):
|
||||
return False
|
||||
|
||||
try:
|
||||
with open(manifest_path, "r") as f:
|
||||
manifest = json.load(f)
|
||||
for file in manifest["files"]:
|
||||
if not os.path.exists(local_dir / file):
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def download_directory(remote_path: str, local_dir: str):
|
||||
model_name = get_model_name(remote_path)
|
||||
s3_url = join_urls(settings.S3_BASE_URL, remote_path)
|
||||
# Check to see if it's already downloaded
|
||||
model_exists = check_manifest(local_dir)
|
||||
if model_exists:
|
||||
return
|
||||
|
||||
# Use tempfile.TemporaryDirectory to automatically clean up
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
# Download the manifest file
|
||||
manifest_file = join_urls(s3_url, "manifest.json")
|
||||
manifest_path = os.path.join(temp_dir, "manifest.json")
|
||||
download_file(manifest_file, manifest_path)
|
||||
|
||||
# List and download all files
|
||||
with open(manifest_path, "r") as f:
|
||||
manifest = json.load(f)
|
||||
|
||||
pbar = tqdm(
|
||||
desc=f"Downloading {model_name} model to {local_dir}",
|
||||
total=len(manifest["files"]),
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(
|
||||
max_workers=settings.PARALLEL_DOWNLOAD_WORKERS
|
||||
) as executor:
|
||||
futures = []
|
||||
for file in manifest["files"]:
|
||||
remote_file = join_urls(s3_url, file)
|
||||
local_file = os.path.join(temp_dir, file)
|
||||
futures.append(executor.submit(download_file, remote_file, local_file))
|
||||
|
||||
for future in futures:
|
||||
future.result()
|
||||
pbar.update(1)
|
||||
|
||||
pbar.close()
|
||||
|
||||
# Move all files to new directory
|
||||
for file in os.listdir(temp_dir):
|
||||
shutil.move(os.path.join(temp_dir, file), local_dir)
|
||||
|
||||
|
||||
class S3DownloaderMixin:
|
||||
s3_prefix = "s3://"
|
||||
|
||||
@classmethod
|
||||
def get_local_path(cls, pretrained_model_name_or_path) -> str:
|
||||
if pretrained_model_name_or_path.startswith(cls.s3_prefix):
|
||||
pretrained_model_name_or_path = pretrained_model_name_or_path.replace(
|
||||
cls.s3_prefix, ""
|
||||
)
|
||||
cache_dir = settings.MODEL_CACHE_DIR
|
||||
local_path = os.path.join(cache_dir, pretrained_model_name_or_path)
|
||||
os.makedirs(local_path, exist_ok=True)
|
||||
else:
|
||||
local_path = ""
|
||||
return local_path
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
|
||||
# Allow loading models directly from the hub, or using s3
|
||||
if not pretrained_model_name_or_path.startswith(cls.s3_prefix):
|
||||
return super().from_pretrained(
|
||||
pretrained_model_name_or_path, *args, **kwargs
|
||||
)
|
||||
|
||||
local_path = cls.get_local_path(pretrained_model_name_or_path)
|
||||
pretrained_model_name_or_path = pretrained_model_name_or_path.replace(
|
||||
cls.s3_prefix, ""
|
||||
)
|
||||
|
||||
# Retry logic for downloading the model folder
|
||||
retries = 3
|
||||
delay = 5
|
||||
attempt = 0
|
||||
success = False
|
||||
while not success and attempt < retries:
|
||||
try:
|
||||
download_directory(pretrained_model_name_or_path, local_path)
|
||||
success = True # If download succeeded
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error downloading model from {pretrained_model_name_or_path}. Attempt {attempt + 1} of {retries}. Error: {e}"
|
||||
)
|
||||
attempt += 1
|
||||
if attempt < retries:
|
||||
logger.info(f"Retrying in {delay} seconds...")
|
||||
time.sleep(delay) # Wait before retrying
|
||||
else:
|
||||
logger.error(
|
||||
f"Failed to download {pretrained_model_name_or_path} after {retries} attempts."
|
||||
)
|
||||
raise e # Reraise exception after max retries
|
||||
|
||||
return super().from_pretrained(local_path, *args, **kwargs)
|
||||
@@ -0,0 +1,44 @@
|
||||
from typing import List
|
||||
|
||||
from surya.common.polygon import PolygonBox
|
||||
|
||||
|
||||
def clean_boxes(boxes: List[PolygonBox]) -> List[PolygonBox]:
|
||||
new_boxes = []
|
||||
for box_obj in boxes:
|
||||
xs = [point[0] for point in box_obj.polygon]
|
||||
ys = [point[1] for point in box_obj.polygon]
|
||||
if max(xs) == min(xs) or max(ys) == min(ys):
|
||||
continue
|
||||
|
||||
box = box_obj.bbox
|
||||
contained = False
|
||||
for other_box_obj in boxes:
|
||||
if other_box_obj.polygon == box_obj.polygon:
|
||||
continue
|
||||
|
||||
other_box = other_box_obj.bbox
|
||||
if box == other_box:
|
||||
continue
|
||||
if (
|
||||
box[0] >= other_box[0]
|
||||
and box[1] >= other_box[1]
|
||||
and box[2] <= other_box[2]
|
||||
and box[3] <= other_box[3]
|
||||
):
|
||||
contained = True
|
||||
break
|
||||
if not contained:
|
||||
new_boxes.append(box_obj)
|
||||
return new_boxes
|
||||
|
||||
|
||||
def expand_bbox(bbox, expansion_factor=0.01):
|
||||
expansion_low = 1 - expansion_factor
|
||||
expansion_high = 1 + expansion_factor
|
||||
return [
|
||||
bbox[0] * expansion_low,
|
||||
bbox[1] * expansion_low,
|
||||
bbox[2] * expansion_high,
|
||||
bbox[3] * expansion_high,
|
||||
]
|
||||
@@ -0,0 +1,66 @@
|
||||
from PIL import ImageDraw, ImageFont
|
||||
|
||||
from surya.debug.fonts import get_font_path
|
||||
from surya.debug.text import get_text_size
|
||||
|
||||
|
||||
def draw_bboxes_on_image(
|
||||
bboxes, image, labels=None, label_font_size=10, color: str | list = "red"
|
||||
):
|
||||
polys = []
|
||||
for bb in bboxes:
|
||||
# Clockwise polygon
|
||||
poly = [[bb[0], bb[1]], [bb[2], bb[1]], [bb[2], bb[3]], [bb[0], bb[3]]]
|
||||
polys.append(poly)
|
||||
|
||||
return draw_polys_on_image(
|
||||
polys, image, labels, label_font_size=label_font_size, color=color
|
||||
)
|
||||
|
||||
|
||||
def draw_polys_on_image(
|
||||
corners,
|
||||
image,
|
||||
labels=None,
|
||||
box_padding=-1,
|
||||
label_offset=1,
|
||||
label_font_size=10,
|
||||
color: str | list = "red",
|
||||
):
|
||||
draw = ImageDraw.Draw(image)
|
||||
font_path = get_font_path()
|
||||
label_font = ImageFont.truetype(font_path, label_font_size)
|
||||
|
||||
for i in range(len(corners)):
|
||||
poly = corners[i]
|
||||
poly = [(int(p[0]), int(p[1])) for p in poly]
|
||||
draw.polygon(
|
||||
poly, outline=color[i] if isinstance(color, list) else color, width=1
|
||||
)
|
||||
|
||||
if labels is not None:
|
||||
label = labels[i]
|
||||
text_position = (
|
||||
min([p[0] for p in poly]) + label_offset,
|
||||
min([p[1] for p in poly]) + label_offset,
|
||||
)
|
||||
text_size = get_text_size(label, label_font)
|
||||
box_position = (
|
||||
text_position[0] - box_padding + label_offset,
|
||||
text_position[1] - box_padding + label_offset,
|
||||
text_position[0] + text_size[0] + box_padding + label_offset,
|
||||
text_position[1] + text_size[1] + box_padding + label_offset,
|
||||
)
|
||||
try:
|
||||
draw.rectangle(box_position, fill="white")
|
||||
except Exception as e:
|
||||
print(f"Error drawing rectangle at {box_position}: {e}")
|
||||
continue
|
||||
draw.text(
|
||||
text_position,
|
||||
label,
|
||||
fill=color[i] if isinstance(color, list) else color,
|
||||
font=label_font,
|
||||
)
|
||||
|
||||
return image
|
||||
@@ -0,0 +1,24 @@
|
||||
from typing import List, Optional
|
||||
import os
|
||||
import requests
|
||||
|
||||
from surya.settings import settings
|
||||
|
||||
|
||||
def get_font_path(langs: Optional[List[str]] = None) -> str:
|
||||
font_path = settings.RECOGNITION_RENDER_FONTS["all"]
|
||||
if langs is not None:
|
||||
for k in settings.RECOGNITION_RENDER_FONTS:
|
||||
if k in langs and len(langs) == 1:
|
||||
font_path = settings.RECOGNITION_RENDER_FONTS[k]
|
||||
break
|
||||
|
||||
if not os.path.exists(font_path):
|
||||
os.makedirs(os.path.dirname(font_path), exist_ok=True)
|
||||
font_dl_path = f"{settings.RECOGNITION_FONT_DL_BASE}/{os.path.basename(font_path)}"
|
||||
with requests.get(font_dl_path, stream=True) as r, open(font_path, 'wb') as f:
|
||||
r.raise_for_status()
|
||||
for chunk in r.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
return font_path
|
||||
@@ -0,0 +1,64 @@
|
||||
<style>
|
||||
.katex-display-container {
|
||||
display: inline-block;
|
||||
max-width: 100%;
|
||||
overflow-x: auto;
|
||||
max-height: 100%;
|
||||
}
|
||||
|
||||
.katex-inline-container {
|
||||
display: inline-block;
|
||||
max-width: 100%;
|
||||
overflow-x: auto;
|
||||
max-height: 100%;
|
||||
}
|
||||
</style>
|
||||
<script src="https://cdn.jsdelivr.net/npm/katex@0.16.21/dist/katex.min.js" onload="setTimeout(function() {renderMath()})" async></script>
|
||||
<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.21/dist/katex.min.css">
|
||||
<script>
|
||||
function htmlUnescape(escapedText) {
|
||||
const htmlEntities = {
|
||||
'&': '&',
|
||||
'<': '<',
|
||||
'>': '>',
|
||||
'"': '"',
|
||||
''': "'",
|
||||
' ': ' '
|
||||
};
|
||||
|
||||
return escapedText.replace(/&|<|>|"|'| /g, match => htmlEntities[match]);
|
||||
}
|
||||
|
||||
const renderMath = (function() {
|
||||
try {
|
||||
const mathElements = document.querySelectorAll('math');
|
||||
|
||||
mathElements.forEach(function(element) {
|
||||
let mathContent = element.innerHTML.trim();
|
||||
mathContent = htmlUnescape(mathContent);
|
||||
const isDisplay = element.getAttribute('display') === 'block';
|
||||
|
||||
const container = document.createElement('span');
|
||||
container.className = isDisplay ? 'katex-display-container' : 'katex-inline-container';
|
||||
element.parentNode.insertBefore(container, element);
|
||||
|
||||
try {
|
||||
katex.render(mathContent, container, {
|
||||
displayMode: isDisplay,
|
||||
throwOnError: false
|
||||
});
|
||||
|
||||
} catch (err) {
|
||||
console.error('KaTeX rendering error:', err);
|
||||
container.textContent = mathContent; // Fallback to raw text
|
||||
}
|
||||
|
||||
element.parentNode.removeChild(element);
|
||||
});
|
||||
|
||||
console.log('Math rendering complete with', mathElements.length, 'expressions');
|
||||
} catch (err) {
|
||||
console.error('Error in renderMath function:', err);
|
||||
}
|
||||
});
|
||||
</script>
|
||||
@@ -0,0 +1,90 @@
|
||||
import html as htmllib
|
||||
import os.path
|
||||
import re
|
||||
|
||||
filepath = os.path.abspath(__file__)
|
||||
|
||||
def render_text_as_html(
|
||||
bboxes: list[list[int]],
|
||||
texts: list[str],
|
||||
image_size: tuple[int, int],
|
||||
base_font_size: int = 16,
|
||||
scaler: int = 2
|
||||
):
|
||||
katex_path = os.path.join(os.path.dirname(filepath), "katex.js")
|
||||
with open(katex_path, "r") as f:
|
||||
katex_script = f.read()
|
||||
|
||||
html_content = []
|
||||
image_size = tuple([int(s * scaler) for s in image_size])
|
||||
width, height = image_size
|
||||
|
||||
|
||||
html_content.append(f"""
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<style>
|
||||
body {{
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
width: {width}px;
|
||||
height: {height}px;
|
||||
position: relative;
|
||||
overflow: hidden;
|
||||
background: white;
|
||||
color: black;
|
||||
}}
|
||||
.text-box {{
|
||||
position: absolute;
|
||||
overflow: hidden;
|
||||
display: flex;
|
||||
justify-content: left;
|
||||
font-family: Arial, sans-serif;
|
||||
white-space: pre-wrap;
|
||||
}}
|
||||
.vertical-text {{
|
||||
writing-mode: vertical-rl; /* Top to bottom, right to left */
|
||||
}}
|
||||
</style>
|
||||
{katex_script}
|
||||
</head>
|
||||
<body>
|
||||
""")
|
||||
|
||||
for i, (bbox, text) in enumerate(zip(bboxes, texts)):
|
||||
bbox = bbox.copy()
|
||||
bbox = [int(bb * scaler) for bb in bbox]
|
||||
x1, y1, x2, y2 = bbox
|
||||
width = x2 - x1
|
||||
height = y2 - y1
|
||||
min_dim = min(width, height)
|
||||
|
||||
# Scale font size based on box height
|
||||
font_size = min(int(min_dim * 0.75), base_font_size)
|
||||
|
||||
# Create div with absolute positioning
|
||||
div_style = (
|
||||
f"left: {x1}px; "
|
||||
f"top: {y1}px; "
|
||||
f"width: {width}px; "
|
||||
f"height: {height}px; "
|
||||
f"font-size: {font_size}px;"
|
||||
)
|
||||
|
||||
class_ = "text-box"
|
||||
if height > width * 2:
|
||||
class_ += " vertical-text"
|
||||
|
||||
# Determine if content is HTML/MathML or plain text
|
||||
if "<" in text and ">" in text and re.search(r"<(html|math|div|sub|sup|i|u|mark|small|del|b|br|code)\b", text.lower()):
|
||||
# Content is already HTML/MathML, include as-is
|
||||
html_content.append(f'<span class="{class_}" id="box-{i}" style="{div_style}">{text}</span>')
|
||||
else:
|
||||
# Plain text, escape it
|
||||
escaped_text = htmllib.escape(text)
|
||||
html_content.append(f'<span class="{class_}" id="box-{i}" style="{div_style}">{escaped_text}</span>')
|
||||
|
||||
html_content.append("</body></html>")
|
||||
|
||||
return "\n".join(html_content), image_size
|
||||
@@ -0,0 +1,8 @@
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
|
||||
def get_text_size(text, font):
|
||||
im = Image.new(mode="P", size=(0, 0))
|
||||
draw = ImageDraw.Draw(im)
|
||||
_, _, width, height = draw.textbbox((0, 0), text=text, font=font)
|
||||
return width, height
|
||||
@@ -0,0 +1,147 @@
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import List, Generator, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from PIL import Image
|
||||
from tqdm import tqdm
|
||||
|
||||
from surya.common.predictor import BasePredictor
|
||||
|
||||
from surya.detection.loader import DetectionModelLoader
|
||||
from surya.detection.parallel import FakeExecutor
|
||||
from surya.detection.util import get_total_splits, split_image
|
||||
from surya.detection.schema import TextDetectionResult
|
||||
from surya.settings import settings
|
||||
from surya.detection.heatmap import parallel_get_boxes
|
||||
|
||||
|
||||
class DetectionPredictor(BasePredictor):
|
||||
model_loader_cls = DetectionModelLoader
|
||||
batch_size = settings.DETECTOR_BATCH_SIZE
|
||||
default_batch_sizes = {"cpu": 8, "mps": 8, "cuda": 36}
|
||||
|
||||
def __call__(
|
||||
self, images: List[Image.Image], batch_size=None, include_maps=False
|
||||
) -> List[TextDetectionResult]:
|
||||
detection_generator = self.batch_detection(images, batch_size=batch_size)
|
||||
|
||||
postprocessing_futures = []
|
||||
max_workers = min(settings.DETECTOR_POSTPROCESSING_CPU_WORKERS, len(images))
|
||||
parallelize = (
|
||||
not settings.IN_STREAMLIT
|
||||
and len(images) >= settings.DETECTOR_MIN_PARALLEL_THRESH
|
||||
)
|
||||
executor = ThreadPoolExecutor if parallelize else FakeExecutor
|
||||
with executor(max_workers=max_workers) as e:
|
||||
for preds, orig_sizes in detection_generator:
|
||||
for pred, orig_size in zip(preds, orig_sizes):
|
||||
postprocessing_futures.append(
|
||||
e.submit(parallel_get_boxes, pred, orig_size, include_maps)
|
||||
)
|
||||
|
||||
return [future.result() for future in postprocessing_futures]
|
||||
|
||||
def prepare_image(self, img):
|
||||
new_size = (self.processor.size["width"], self.processor.size["height"])
|
||||
|
||||
# This double resize actually necessary for downstream accuracy
|
||||
img.thumbnail(new_size, Image.Resampling.LANCZOS)
|
||||
img = img.resize(
|
||||
new_size, Image.Resampling.LANCZOS
|
||||
) # Stretch smaller dimension to fit new size
|
||||
|
||||
img = np.asarray(img, dtype=np.uint8)
|
||||
img = self.processor(img)["pixel_values"][0]
|
||||
img = torch.from_numpy(img)
|
||||
return img
|
||||
|
||||
def batch_detection(
|
||||
self, images: List, batch_size=None
|
||||
) -> Generator[Tuple[List[List[np.ndarray]], List[Tuple[int, int]]], None, None]:
|
||||
assert all([isinstance(image, Image.Image) for image in images])
|
||||
if batch_size is None:
|
||||
batch_size = self.get_batch_size()
|
||||
heatmap_count = self.model.config.num_labels
|
||||
|
||||
orig_sizes = [image.size for image in images]
|
||||
splits_per_image = [
|
||||
get_total_splits(size, self.processor.size["height"]) for size in orig_sizes
|
||||
]
|
||||
|
||||
batches = []
|
||||
current_batch_size = 0
|
||||
current_batch = []
|
||||
for i in range(len(images)):
|
||||
if current_batch_size + splits_per_image[i] > batch_size:
|
||||
if len(current_batch) > 0:
|
||||
batches.append(current_batch)
|
||||
current_batch = []
|
||||
current_batch_size = 0
|
||||
current_batch.append(i)
|
||||
current_batch_size += splits_per_image[i]
|
||||
|
||||
if len(current_batch) > 0:
|
||||
batches.append(current_batch)
|
||||
|
||||
for batch_idx in tqdm(
|
||||
range(len(batches)), desc="Detecting bboxes", disable=self.disable_tqdm
|
||||
):
|
||||
batch_image_idxs = batches[batch_idx]
|
||||
batch_images = [images[j].convert("RGB") for j in batch_image_idxs]
|
||||
|
||||
split_index = []
|
||||
split_heights = []
|
||||
image_splits = []
|
||||
for image_idx, image in enumerate(batch_images):
|
||||
image_parts, split_height = split_image(
|
||||
image, self.processor.size["height"]
|
||||
)
|
||||
image_splits.extend(image_parts)
|
||||
split_index.extend([image_idx] * len(image_parts))
|
||||
split_heights.extend(split_height)
|
||||
|
||||
image_splits = [self.prepare_image(image) for image in image_splits]
|
||||
# Batch images in dim 0
|
||||
batch = torch.stack(image_splits, dim=0).to(self.model.dtype)
|
||||
|
||||
with settings.INFERENCE_MODE():
|
||||
pred = self.model(pixel_values=batch.to(self.model.device))
|
||||
|
||||
logits = pred.logits
|
||||
correct_shape = [
|
||||
self.processor.size["height"],
|
||||
self.processor.size["width"],
|
||||
]
|
||||
current_shape = list(logits.shape[2:])
|
||||
if current_shape != correct_shape:
|
||||
logits = F.interpolate(
|
||||
logits, size=correct_shape, mode="bilinear", align_corners=False
|
||||
)
|
||||
|
||||
logits = logits.to(torch.float32).cpu().numpy()
|
||||
preds = []
|
||||
for i, (idx, height) in enumerate(zip(split_index, split_heights)):
|
||||
# If our current prediction length is below the image idx, that means we have a new image
|
||||
# Otherwise, we need to add to the current image
|
||||
if len(preds) <= idx:
|
||||
preds.append([logits[i][k] for k in range(heatmap_count)])
|
||||
else:
|
||||
heatmaps = preds[idx]
|
||||
pred_heatmaps = [logits[i][k] for k in range(heatmap_count)]
|
||||
|
||||
if height < self.processor.size["height"]:
|
||||
# Cut off padding to get original height
|
||||
pred_heatmaps = [
|
||||
pred_heatmap[:height, :] for pred_heatmap in pred_heatmaps
|
||||
]
|
||||
|
||||
for k in range(heatmap_count):
|
||||
heatmaps[k] = np.vstack([heatmaps[k], pred_heatmaps[k]])
|
||||
preds[idx] = heatmaps
|
||||
|
||||
yield preds, [orig_sizes[j] for j in batch_image_idxs]
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
@@ -0,0 +1,165 @@
|
||||
from typing import List
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from surya.common.util import clean_boxes
|
||||
from surya.detection import TextDetectionResult
|
||||
from surya.common.polygon import PolygonBox
|
||||
from surya.settings import settings
|
||||
|
||||
|
||||
def get_dynamic_thresholds(linemap, text_threshold, low_text, typical_top10_avg=0.7):
|
||||
# Find average intensity of top 10% pixels
|
||||
flat_map = linemap.ravel()
|
||||
top_10_count = int(len(flat_map) * 0.9)
|
||||
avg_intensity = np.mean(np.partition(flat_map, top_10_count)[top_10_count:])
|
||||
scaling_factor = np.clip(avg_intensity / typical_top10_avg, 0, 1) ** (1 / 2)
|
||||
|
||||
low_text = np.clip(low_text * scaling_factor, 0.1, 0.6)
|
||||
text_threshold = np.clip(text_threshold * scaling_factor, 0.15, 0.8)
|
||||
|
||||
return text_threshold, low_text
|
||||
|
||||
|
||||
def detect_boxes(linemap, text_threshold, low_text):
|
||||
# From CRAFT - https://github.com/clovaai/CRAFT-pytorch
|
||||
# Modified to return boxes and for speed, accuracy
|
||||
img_h, img_w = linemap.shape
|
||||
|
||||
text_threshold, low_text = get_dynamic_thresholds(linemap, text_threshold, low_text)
|
||||
|
||||
text_score_comb = (linemap > low_text).astype(np.uint8)
|
||||
label_count, labels, stats, centroids = cv2.connectedComponentsWithStats(
|
||||
text_score_comb, connectivity=4
|
||||
)
|
||||
|
||||
det = []
|
||||
confidences = []
|
||||
max_confidence = 0
|
||||
|
||||
for k in range(1, label_count):
|
||||
# size filtering
|
||||
size = stats[k, cv2.CC_STAT_AREA]
|
||||
if size < 10:
|
||||
continue
|
||||
|
||||
# make segmentation map
|
||||
x, y, w, h = stats[
|
||||
k,
|
||||
[cv2.CC_STAT_LEFT, cv2.CC_STAT_TOP, cv2.CC_STAT_WIDTH, cv2.CC_STAT_HEIGHT],
|
||||
]
|
||||
|
||||
try:
|
||||
niter = int(np.sqrt(min(w, h)))
|
||||
except ValueError:
|
||||
niter = 0
|
||||
|
||||
buffer = 1
|
||||
sx, sy = max(0, x - niter - buffer), max(0, y - niter - buffer)
|
||||
ex, ey = min(img_w, x + w + niter + buffer), min(img_h, y + h + niter + buffer)
|
||||
|
||||
mask = labels[sy:ey, sx:ex] == k
|
||||
selected_linemap = linemap[sy:ey, sx:ex][mask]
|
||||
if selected_linemap.size == 0:
|
||||
continue
|
||||
|
||||
line_max = np.max(selected_linemap)
|
||||
|
||||
# thresholding
|
||||
if line_max < text_threshold:
|
||||
continue
|
||||
|
||||
segmap = mask.astype(np.uint8)
|
||||
|
||||
ksize = buffer + niter
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (ksize, ksize))
|
||||
selected_segmap = cv2.dilate(segmap, kernel)
|
||||
|
||||
# make box
|
||||
y_inds, x_inds = np.nonzero(selected_segmap)
|
||||
x_inds += sx
|
||||
y_inds += sy
|
||||
np_contours = np.column_stack((x_inds, y_inds))
|
||||
rectangle = cv2.minAreaRect(np_contours)
|
||||
box = cv2.boxPoints(rectangle)
|
||||
|
||||
# align diamond-shape
|
||||
w, h = np.linalg.norm(box[0] - box[1]), np.linalg.norm(box[1] - box[2])
|
||||
box_ratio = max(w, h) / (min(w, h) + 1e-5)
|
||||
if abs(1 - box_ratio) <= 0.1:
|
||||
left, right = np_contours[:, 0].min(), np_contours[:, 0].max()
|
||||
top, bottom = np_contours[:, 1].min(), np_contours[:, 1].max()
|
||||
box = np.array(
|
||||
[[left, top], [right, top], [right, bottom], [left, bottom]],
|
||||
dtype=np.float32,
|
||||
)
|
||||
|
||||
# make clock-wise order
|
||||
startidx = box.sum(axis=1).argmin()
|
||||
box = np.roll(box, 4 - startidx, 0)
|
||||
|
||||
max_confidence = max(max_confidence, line_max)
|
||||
|
||||
confidences.append(line_max)
|
||||
det.append(box)
|
||||
|
||||
if max_confidence > 0:
|
||||
confidences = [c / max_confidence for c in confidences]
|
||||
return det, confidences
|
||||
|
||||
|
||||
def get_detected_boxes(textmap, text_threshold=None, low_text=None) -> List[PolygonBox]:
|
||||
if text_threshold is None:
|
||||
text_threshold = settings.DETECTOR_TEXT_THRESHOLD
|
||||
if low_text is None:
|
||||
low_text = settings.DETECTOR_BLANK_THRESHOLD
|
||||
|
||||
if textmap.dtype != np.float32:
|
||||
textmap = textmap.astype(np.float32)
|
||||
|
||||
boxes, confidences = detect_boxes(textmap, text_threshold, low_text)
|
||||
# From point form to box form
|
||||
return [
|
||||
PolygonBox(polygon=box, confidence=confidence)
|
||||
for box, confidence in zip(boxes, confidences)
|
||||
]
|
||||
|
||||
|
||||
def get_and_clean_boxes(
|
||||
textmap, processor_size, image_size, text_threshold=None, low_text=None
|
||||
) -> List[PolygonBox]:
|
||||
bboxes = get_detected_boxes(textmap, text_threshold, low_text)
|
||||
for bbox in bboxes:
|
||||
bbox.rescale(processor_size, image_size)
|
||||
bbox.fit_to_bounds([0, 0, image_size[0], image_size[1]])
|
||||
|
||||
bboxes = clean_boxes(bboxes)
|
||||
return bboxes
|
||||
|
||||
|
||||
def parallel_get_boxes(preds, orig_sizes, include_maps=False):
|
||||
heatmap, affinity_map = preds
|
||||
heat_img, aff_img = None, None
|
||||
|
||||
if include_maps:
|
||||
heat_img = Image.fromarray((heatmap * 255).astype(np.uint8))
|
||||
aff_img = Image.fromarray((affinity_map * 255).astype(np.uint8))
|
||||
heatmap_size = list(reversed(heatmap.shape))
|
||||
bboxes = get_and_clean_boxes(heatmap, heatmap_size, orig_sizes)
|
||||
for box in bboxes:
|
||||
# Skip for vertical boxes
|
||||
if box.height < 3 * box.width:
|
||||
box.expand(x_margin=0, y_margin=settings.DETECTOR_BOX_Y_EXPAND_MARGIN)
|
||||
box.fit_to_bounds(
|
||||
[0, 0, orig_sizes[0], orig_sizes[1]]
|
||||
) # Fix any bad expands
|
||||
|
||||
result = TextDetectionResult(
|
||||
bboxes=bboxes,
|
||||
heatmap=heat_img,
|
||||
affinity_map=aff_img,
|
||||
image_bbox=[0, 0, orig_sizes[0], orig_sizes[1]],
|
||||
)
|
||||
return result
|
||||
@@ -0,0 +1,53 @@
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from surya.common.load import ModelLoader
|
||||
from surya.detection.processor import SegformerImageProcessor
|
||||
|
||||
from surya.detection.model.config import EfficientViTConfig
|
||||
from surya.detection.model.encoderdecoder import EfficientViTForSemanticSegmentation
|
||||
from surya.logging import get_logger
|
||||
from surya.settings import settings
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
class DetectionModelLoader(ModelLoader):
|
||||
def __init__(self, checkpoint: Optional[str] = None):
|
||||
super().__init__(checkpoint)
|
||||
|
||||
if self.checkpoint is None:
|
||||
self.checkpoint = settings.DETECTOR_MODEL_CHECKPOINT
|
||||
|
||||
def model(
|
||||
self,
|
||||
device: Optional[torch.device | str] = None,
|
||||
dtype: Optional[torch.dtype | str] = None,
|
||||
attention_implementation: Optional[str] = None,
|
||||
) -> EfficientViTForSemanticSegmentation:
|
||||
if device is None:
|
||||
device = settings.TORCH_DEVICE_MODEL
|
||||
if dtype is None:
|
||||
dtype = settings.MODEL_DTYPE
|
||||
|
||||
config = EfficientViTConfig.from_pretrained(self.checkpoint)
|
||||
model = EfficientViTForSemanticSegmentation.from_pretrained(
|
||||
self.checkpoint,
|
||||
dtype=dtype,
|
||||
config=config,
|
||||
)
|
||||
model = model.to(device)
|
||||
model = model.eval()
|
||||
|
||||
logger.debug(
|
||||
f"Loaded detection model {self.checkpoint} from {EfficientViTForSemanticSegmentation.get_local_path(self.checkpoint)} onto device {device} with dtype {dtype}"
|
||||
)
|
||||
return model
|
||||
|
||||
def processor(
|
||||
self,
|
||||
device: Optional[torch.device | str] = None,
|
||||
dtype: Optional[torch.dtype | str] = None,
|
||||
) -> SegformerImageProcessor:
|
||||
return SegformerImageProcessor.from_pretrained(self.checkpoint)
|
||||
@@ -0,0 +1,53 @@
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from surya.common.s3 import S3DownloaderMixin
|
||||
|
||||
|
||||
class EfficientViTConfig(S3DownloaderMixin, PretrainedConfig):
|
||||
r"""
|
||||
```"""
|
||||
|
||||
model_type = "efficientvit"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_classes=2,
|
||||
num_channels=3,
|
||||
widths=(32, 64, 128, 256, 512),
|
||||
head_dim=32,
|
||||
num_stages=4,
|
||||
depths=(1, 1, 1, 6, 6),
|
||||
strides=(2, 2, 2, 2, 2),
|
||||
hidden_sizes=(32, 64, 160, 256),
|
||||
patch_size=(7, 7),
|
||||
hidden_dropout_prob=0.0,
|
||||
attention_probs_dropout_prob=0.0,
|
||||
classifier_dropout_prob=0.0,
|
||||
layer_norm_eps=1e-6,
|
||||
decoder_layer_hidden_size=128,
|
||||
decoder_hidden_size=512,
|
||||
semantic_loss_ignore_index=255,
|
||||
initializer_range=0.02,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
self.num_classes = num_classes
|
||||
self.widths = widths
|
||||
self.head_dim = head_dim
|
||||
|
||||
self.num_channels = num_channels
|
||||
self.num_stages = num_stages
|
||||
self.depths = depths
|
||||
self.strides = strides
|
||||
self.hidden_sizes = hidden_sizes
|
||||
self.patch_size = patch_size
|
||||
self.hidden_dropout_prob = hidden_dropout_prob
|
||||
self.attention_probs_dropout_prob = attention_probs_dropout_prob
|
||||
self.classifier_dropout_prob = classifier_dropout_prob
|
||||
self.layer_norm_eps = layer_norm_eps
|
||||
self.decoder_hidden_size = decoder_hidden_size
|
||||
self.decoder_layer_hidden_size = decoder_layer_hidden_size
|
||||
self.semantic_loss_ignore_index = semantic_loss_ignore_index
|
||||
|
||||
self.initializer_range = initializer_range
|
||||
@@ -0,0 +1,839 @@
|
||||
"""
|
||||
This is an implementation of efficientvit, with some modifications (decode head, etc).
|
||||
|
||||
Original paper at https://arxiv.org/abs/2205.14756
|
||||
|
||||
Code adapted from timm, https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/efficientvit_mit.py
|
||||
Original code (that timm adapted from) at https://github.com/mit-han-lab/efficientvit
|
||||
|
||||
License: Apache 2
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional, Union, Tuple, List, Any
|
||||
from functools import partial
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from transformers.modeling_outputs import SemanticSegmenterOutput
|
||||
|
||||
from surya.common.pretrained import SuryaPreTrainedModel
|
||||
from surya.common.s3 import S3DownloaderMixin
|
||||
from surya.detection.model.config import EfficientViTConfig
|
||||
|
||||
|
||||
def val2list(x: Union[List, Tuple, Any], repeat_time=1):
|
||||
if isinstance(x, (list, tuple)):
|
||||
return list(x)
|
||||
return [x for _ in range(repeat_time)]
|
||||
|
||||
|
||||
def val2tuple(x: Union[List, Tuple, Any], min_len: int = 1, idx_repeat: int = -1):
|
||||
# repeat elements if necessary
|
||||
x = val2list(x)
|
||||
if len(x) > 0:
|
||||
x[idx_repeat:idx_repeat] = [x[idx_repeat] for _ in range(min_len - len(x))]
|
||||
|
||||
return tuple(x)
|
||||
|
||||
|
||||
def get_same_padding(
|
||||
kernel_size: Union[int, Tuple[int, ...]],
|
||||
) -> Union[int, Tuple[int, ...]]:
|
||||
if isinstance(kernel_size, tuple):
|
||||
return tuple([get_same_padding(ks) for ks in kernel_size])
|
||||
else:
|
||||
assert kernel_size % 2 > 0, "kernel size should be odd number"
|
||||
return kernel_size // 2
|
||||
|
||||
|
||||
def get_padding(kernel_size: int, stride: int = 1, dilation: int = 1) -> int:
|
||||
padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
|
||||
return padding
|
||||
|
||||
|
||||
class ConvNormAct(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=False,
|
||||
dropout=0.0,
|
||||
norm_layer=nn.BatchNorm2d,
|
||||
act_layer=nn.ReLU,
|
||||
):
|
||||
super(ConvNormAct, self).__init__()
|
||||
self.dropout = nn.Dropout(dropout, inplace=False)
|
||||
padding = get_padding(kernel_size, stride, dilation)
|
||||
self.conv = nn.Conv2d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding=padding,
|
||||
)
|
||||
self.norm = (
|
||||
norm_layer(num_features=out_channels) if norm_layer else nn.Identity()
|
||||
)
|
||||
self.act = act_layer(inplace=True) if act_layer is not None else nn.Identity()
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.norm(x)
|
||||
x = self.act(x)
|
||||
return x
|
||||
|
||||
|
||||
class DSConv(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
use_bias=False,
|
||||
norm_layer=(nn.BatchNorm2d, nn.BatchNorm2d),
|
||||
act_layer=(nn.ReLU6, None),
|
||||
):
|
||||
super(DSConv, self).__init__()
|
||||
use_bias = val2tuple(use_bias, 2)
|
||||
norm_layer = val2tuple(norm_layer, 2)
|
||||
act_layer = val2tuple(act_layer, 2)
|
||||
|
||||
self.depth_conv = ConvNormAct(
|
||||
in_channels,
|
||||
in_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
groups=in_channels,
|
||||
norm_layer=norm_layer[0],
|
||||
act_layer=act_layer[0],
|
||||
bias=use_bias[0],
|
||||
)
|
||||
self.point_conv = ConvNormAct(
|
||||
in_channels,
|
||||
out_channels,
|
||||
1,
|
||||
norm_layer=norm_layer[1],
|
||||
act_layer=act_layer[1],
|
||||
bias=use_bias[1],
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.depth_conv(x)
|
||||
x = self.point_conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class ConvBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
mid_channels=None,
|
||||
expand_ratio=1,
|
||||
use_bias=False,
|
||||
norm_layer=(nn.BatchNorm2d, nn.BatchNorm2d),
|
||||
act_layer=(nn.ReLU6, None),
|
||||
):
|
||||
super(ConvBlock, self).__init__()
|
||||
use_bias = val2tuple(use_bias, 2)
|
||||
norm_layer = val2tuple(norm_layer, 2)
|
||||
act_layer = val2tuple(act_layer, 2)
|
||||
mid_channels = mid_channels or round(in_channels * expand_ratio)
|
||||
|
||||
self.conv1 = ConvNormAct(
|
||||
in_channels,
|
||||
mid_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
norm_layer=norm_layer[0],
|
||||
act_layer=act_layer[0],
|
||||
bias=use_bias[0],
|
||||
)
|
||||
self.conv2 = ConvNormAct(
|
||||
mid_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
1,
|
||||
norm_layer=norm_layer[1],
|
||||
act_layer=act_layer[1],
|
||||
bias=use_bias[1],
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.conv2(x)
|
||||
return x
|
||||
|
||||
|
||||
class MBConv(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
mid_channels=None,
|
||||
expand_ratio=6,
|
||||
use_bias=False,
|
||||
norm_layer=(nn.BatchNorm2d, nn.BatchNorm2d, nn.BatchNorm2d),
|
||||
act_layer=(nn.ReLU6, nn.ReLU6, None),
|
||||
):
|
||||
super(MBConv, self).__init__()
|
||||
use_bias = val2tuple(use_bias, 3)
|
||||
norm_layer = val2tuple(norm_layer, 3)
|
||||
act_layer = val2tuple(act_layer, 3)
|
||||
mid_channels = mid_channels or round(in_channels * expand_ratio)
|
||||
|
||||
self.inverted_conv = ConvNormAct(
|
||||
in_channels,
|
||||
mid_channels,
|
||||
1,
|
||||
stride=1,
|
||||
norm_layer=norm_layer[0],
|
||||
act_layer=act_layer[0],
|
||||
bias=use_bias[0],
|
||||
)
|
||||
self.depth_conv = ConvNormAct(
|
||||
mid_channels,
|
||||
mid_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
groups=mid_channels,
|
||||
norm_layer=norm_layer[1],
|
||||
act_layer=act_layer[1],
|
||||
bias=use_bias[1],
|
||||
)
|
||||
self.point_conv = ConvNormAct(
|
||||
mid_channels,
|
||||
out_channels,
|
||||
1,
|
||||
norm_layer=norm_layer[2],
|
||||
act_layer=act_layer[2],
|
||||
bias=use_bias[2],
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.inverted_conv(x)
|
||||
x = self.depth_conv(x)
|
||||
x = self.point_conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class FusedMBConv(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
mid_channels=None,
|
||||
expand_ratio=6,
|
||||
groups=1,
|
||||
use_bias=False,
|
||||
norm_layer=(nn.BatchNorm2d, nn.BatchNorm2d),
|
||||
act_layer=(nn.ReLU6, None),
|
||||
):
|
||||
super(FusedMBConv, self).__init__()
|
||||
use_bias = val2tuple(use_bias, 2)
|
||||
norm_layer = val2tuple(norm_layer, 2)
|
||||
act_layer = val2tuple(act_layer, 2)
|
||||
mid_channels = mid_channels or round(in_channels * expand_ratio)
|
||||
|
||||
self.spatial_conv = ConvNormAct(
|
||||
in_channels,
|
||||
mid_channels,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
groups=groups,
|
||||
norm_layer=norm_layer[0],
|
||||
act_layer=act_layer[0],
|
||||
bias=use_bias[0],
|
||||
)
|
||||
self.point_conv = ConvNormAct(
|
||||
mid_channels,
|
||||
out_channels,
|
||||
1,
|
||||
norm_layer=norm_layer[1],
|
||||
act_layer=act_layer[1],
|
||||
bias=use_bias[1],
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.spatial_conv(x)
|
||||
x = self.point_conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class LiteMLA(nn.Module):
|
||||
"""Lightweight multi-scale linear attention"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
heads: Union[int, None] = None,
|
||||
heads_ratio: float = 1.0,
|
||||
dim=8,
|
||||
use_bias=False,
|
||||
norm_layer=(None, nn.BatchNorm2d),
|
||||
act_layer=(None, None),
|
||||
kernel_func=nn.ReLU,
|
||||
scales=(5,),
|
||||
eps=1e-5,
|
||||
):
|
||||
super(LiteMLA, self).__init__()
|
||||
self.eps = eps
|
||||
heads = heads or int(in_channels // dim * heads_ratio)
|
||||
total_dim = heads * dim
|
||||
use_bias = val2tuple(use_bias, 2)
|
||||
norm_layer = val2tuple(norm_layer, 2)
|
||||
act_layer = val2tuple(act_layer, 2)
|
||||
|
||||
self.dim = dim
|
||||
self.qkv = ConvNormAct(
|
||||
in_channels,
|
||||
3 * total_dim,
|
||||
1,
|
||||
bias=use_bias[0],
|
||||
norm_layer=norm_layer[0],
|
||||
act_layer=act_layer[0],
|
||||
)
|
||||
self.aggreg = nn.ModuleList(
|
||||
[
|
||||
nn.Sequential(
|
||||
nn.Conv2d(
|
||||
3 * total_dim,
|
||||
3 * total_dim,
|
||||
scale,
|
||||
padding=get_same_padding(scale),
|
||||
groups=3 * total_dim,
|
||||
bias=use_bias[0],
|
||||
),
|
||||
nn.Conv2d(
|
||||
3 * total_dim,
|
||||
3 * total_dim,
|
||||
1,
|
||||
groups=3 * heads,
|
||||
bias=use_bias[0],
|
||||
),
|
||||
)
|
||||
for scale in scales
|
||||
]
|
||||
)
|
||||
self.kernel_func = kernel_func(inplace=False)
|
||||
|
||||
self.proj = ConvNormAct(
|
||||
total_dim * (1 + len(scales)),
|
||||
out_channels,
|
||||
1,
|
||||
bias=use_bias[1],
|
||||
norm_layer=norm_layer[1],
|
||||
act_layer=act_layer[1],
|
||||
)
|
||||
|
||||
def _attn(self, q, k, v):
|
||||
dtype = v.dtype
|
||||
q, k, v = q.float(), k.float(), v.float()
|
||||
kv = k.transpose(-1, -2) @ v
|
||||
out = q @ kv
|
||||
out = out[..., :-1] / (out[..., -1:] + self.eps)
|
||||
return out.to(dtype)
|
||||
|
||||
def forward(self, x):
|
||||
# Shape is B, C, H, W
|
||||
B, _, H, W = x.shape
|
||||
|
||||
# generate multi-scale q, k, v
|
||||
qkv = self.qkv(x)
|
||||
multi_scale_qkv = [qkv]
|
||||
for op in self.aggreg:
|
||||
multi_scale_qkv.append(op(qkv))
|
||||
multi_scale_qkv = torch.cat(multi_scale_qkv, dim=1)
|
||||
multi_scale_qkv = multi_scale_qkv.reshape(B, -1, 3 * self.dim, H * W).transpose(
|
||||
-1, -2
|
||||
)
|
||||
# Shape for each is B, C, HW, head_dim
|
||||
q, k, v = multi_scale_qkv.chunk(3, dim=-1)
|
||||
|
||||
# lightweight global attention
|
||||
q = self.kernel_func(q)
|
||||
k = self.kernel_func(k)
|
||||
v = F.pad(v, (0, 1), mode="constant", value=1.0)
|
||||
|
||||
out = self._attn(q, k, v)
|
||||
|
||||
# final projection
|
||||
out = out.transpose(-1, -2).reshape(B, -1, H, W)
|
||||
out = self.proj(out)
|
||||
return out
|
||||
|
||||
|
||||
class EfficientVitBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
heads_ratio=1.0,
|
||||
head_dim=32,
|
||||
expand_ratio=4,
|
||||
norm_layer=nn.BatchNorm2d,
|
||||
act_layer=nn.Hardswish,
|
||||
):
|
||||
super(EfficientVitBlock, self).__init__()
|
||||
self.context_module = ResidualBlock(
|
||||
LiteMLA(
|
||||
in_channels=in_channels,
|
||||
out_channels=in_channels,
|
||||
heads_ratio=heads_ratio,
|
||||
dim=head_dim,
|
||||
norm_layer=(None, norm_layer),
|
||||
),
|
||||
nn.Identity(),
|
||||
)
|
||||
self.local_module = ResidualBlock(
|
||||
MBConv(
|
||||
in_channels=in_channels,
|
||||
out_channels=in_channels,
|
||||
expand_ratio=expand_ratio,
|
||||
use_bias=(True, True, False),
|
||||
norm_layer=(None, None, norm_layer),
|
||||
act_layer=(act_layer, act_layer, None),
|
||||
),
|
||||
nn.Identity(),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.context_module(x)
|
||||
x = self.local_module(x)
|
||||
return x
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
main: Optional[nn.Module],
|
||||
shortcut: Optional[nn.Module] = None,
|
||||
pre_norm: Optional[nn.Module] = None,
|
||||
):
|
||||
super(ResidualBlock, self).__init__()
|
||||
self.pre_norm = pre_norm if pre_norm is not None else nn.Identity()
|
||||
self.main = main
|
||||
self.shortcut = shortcut
|
||||
|
||||
def forward(self, x):
|
||||
res = self.main(self.pre_norm(x))
|
||||
if self.shortcut is not None:
|
||||
res = res + self.shortcut(x)
|
||||
return res
|
||||
|
||||
|
||||
def build_local_block(
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
stride: int,
|
||||
kernel_size: int,
|
||||
expand_ratio: float,
|
||||
norm_layer: str,
|
||||
act_layer: str,
|
||||
fewer_norm: bool = False,
|
||||
block_type: str = "default",
|
||||
):
|
||||
assert block_type in ["default", "large", "fused"]
|
||||
if expand_ratio == 1:
|
||||
if block_type == "default":
|
||||
block = DSConv(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
stride=stride,
|
||||
kernel_size=kernel_size,
|
||||
use_bias=(True, False) if fewer_norm else False,
|
||||
norm_layer=(None, norm_layer) if fewer_norm else norm_layer,
|
||||
act_layer=(act_layer, None),
|
||||
)
|
||||
else:
|
||||
block = ConvBlock(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
stride=stride,
|
||||
kernel_size=kernel_size,
|
||||
use_bias=(True, False) if fewer_norm else False,
|
||||
norm_layer=(None, norm_layer) if fewer_norm else norm_layer,
|
||||
act_layer=(act_layer, None),
|
||||
)
|
||||
else:
|
||||
if block_type == "default":
|
||||
block = MBConv(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
stride=stride,
|
||||
kernel_size=kernel_size,
|
||||
expand_ratio=expand_ratio,
|
||||
use_bias=(True, True, False) if fewer_norm else False,
|
||||
norm_layer=(None, None, norm_layer) if fewer_norm else norm_layer,
|
||||
act_layer=(act_layer, act_layer, None),
|
||||
)
|
||||
else:
|
||||
block = FusedMBConv(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
stride=stride,
|
||||
kernel_size=kernel_size,
|
||||
expand_ratio=expand_ratio,
|
||||
use_bias=(True, False) if fewer_norm else False,
|
||||
norm_layer=(None, norm_layer) if fewer_norm else norm_layer,
|
||||
act_layer=(act_layer, None),
|
||||
)
|
||||
return block
|
||||
|
||||
|
||||
class Stem(nn.Sequential):
|
||||
def __init__(
|
||||
self,
|
||||
in_chs,
|
||||
out_chs,
|
||||
depth,
|
||||
stride,
|
||||
norm_layer,
|
||||
act_layer,
|
||||
block_type="default",
|
||||
):
|
||||
super().__init__()
|
||||
self.stride = stride
|
||||
|
||||
self.add_module(
|
||||
"in_conv",
|
||||
ConvNormAct(
|
||||
in_chs,
|
||||
out_chs,
|
||||
kernel_size=stride + 1,
|
||||
stride=stride,
|
||||
norm_layer=norm_layer,
|
||||
act_layer=act_layer,
|
||||
),
|
||||
)
|
||||
stem_block = 0
|
||||
for _ in range(depth):
|
||||
self.add_module(
|
||||
f"res{stem_block}",
|
||||
ResidualBlock(
|
||||
build_local_block(
|
||||
in_channels=out_chs,
|
||||
out_channels=out_chs,
|
||||
stride=1,
|
||||
kernel_size=3,
|
||||
expand_ratio=1,
|
||||
norm_layer=norm_layer,
|
||||
act_layer=act_layer,
|
||||
block_type=block_type,
|
||||
),
|
||||
nn.Identity(),
|
||||
),
|
||||
)
|
||||
stem_block += 1
|
||||
|
||||
|
||||
class EfficientVitLargeStage(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_chs,
|
||||
out_chs,
|
||||
depth,
|
||||
stride,
|
||||
norm_layer,
|
||||
act_layer,
|
||||
head_dim,
|
||||
vit_stage=False,
|
||||
fewer_norm=False,
|
||||
):
|
||||
super(EfficientVitLargeStage, self).__init__()
|
||||
blocks = [
|
||||
ResidualBlock(
|
||||
build_local_block(
|
||||
in_channels=in_chs,
|
||||
out_channels=out_chs,
|
||||
stride=stride,
|
||||
kernel_size=stride + 1,
|
||||
expand_ratio=24 if vit_stage else 16,
|
||||
norm_layer=norm_layer,
|
||||
act_layer=act_layer,
|
||||
fewer_norm=vit_stage or fewer_norm,
|
||||
block_type="default" if fewer_norm else "fused",
|
||||
),
|
||||
None,
|
||||
)
|
||||
]
|
||||
in_chs = out_chs
|
||||
|
||||
if vit_stage:
|
||||
# for stage 4
|
||||
for _ in range(depth):
|
||||
blocks.append(
|
||||
EfficientVitBlock(
|
||||
in_channels=in_chs,
|
||||
head_dim=head_dim,
|
||||
expand_ratio=6,
|
||||
norm_layer=norm_layer,
|
||||
act_layer=act_layer,
|
||||
)
|
||||
)
|
||||
else:
|
||||
# for stage 1, 2, 3
|
||||
for i in range(depth):
|
||||
blocks.append(
|
||||
ResidualBlock(
|
||||
build_local_block(
|
||||
in_channels=in_chs,
|
||||
out_channels=out_chs,
|
||||
stride=1,
|
||||
kernel_size=3,
|
||||
expand_ratio=4,
|
||||
norm_layer=norm_layer,
|
||||
act_layer=act_layer,
|
||||
fewer_norm=fewer_norm,
|
||||
block_type="default" if fewer_norm else "fused",
|
||||
),
|
||||
nn.Identity(),
|
||||
)
|
||||
)
|
||||
|
||||
self.blocks = nn.Sequential(*blocks)
|
||||
|
||||
def forward(self, x):
|
||||
return self.blocks(x)
|
||||
|
||||
|
||||
class EfficientVitLarge(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
config: EfficientViTConfig,
|
||||
norm_layer=nn.BatchNorm2d,
|
||||
act_layer=nn.Hardswish,
|
||||
):
|
||||
super(EfficientVitLarge, self).__init__()
|
||||
self.grad_checkpointing = False
|
||||
self.num_classes = config.num_classes
|
||||
self.norm_eps = config.layer_norm_eps
|
||||
norm_layer = partial(norm_layer, eps=self.norm_eps)
|
||||
|
||||
# input stem
|
||||
self.stem = Stem(
|
||||
config.num_channels,
|
||||
config.widths[0],
|
||||
config.depths[0],
|
||||
config.strides[0],
|
||||
norm_layer,
|
||||
act_layer,
|
||||
block_type="large",
|
||||
)
|
||||
stride = config.strides[0]
|
||||
|
||||
# stages
|
||||
self.feature_info = []
|
||||
self.stages = nn.Sequential()
|
||||
in_channels = config.widths[0]
|
||||
for i, (w, d, s) in enumerate(
|
||||
zip(config.widths[1:], config.depths[1:], config.strides[1:])
|
||||
):
|
||||
self.stages.append(
|
||||
EfficientVitLargeStage(
|
||||
in_channels,
|
||||
w,
|
||||
depth=d,
|
||||
stride=s,
|
||||
norm_layer=norm_layer,
|
||||
act_layer=act_layer,
|
||||
head_dim=config.head_dim,
|
||||
vit_stage=i >= 3,
|
||||
fewer_norm=i >= 2,
|
||||
)
|
||||
)
|
||||
stride *= s
|
||||
in_channels = w
|
||||
self.feature_info += [
|
||||
dict(num_chs=in_channels, reduction=stride, module=f"stages.{i}")
|
||||
]
|
||||
|
||||
self.num_features = in_channels
|
||||
|
||||
@torch.jit.ignore
|
||||
def set_grad_checkpointing(self, enable=True):
|
||||
self.grad_checkpointing = enable
|
||||
|
||||
def forward(self, x):
|
||||
x = self.stem(x)
|
||||
encoder_hidden_states = []
|
||||
for i, module in enumerate(self.stages):
|
||||
x = module(x)
|
||||
encoder_hidden_states.append(x)
|
||||
|
||||
return encoder_hidden_states
|
||||
|
||||
|
||||
class EfficientViTPreTrainedModel(SuryaPreTrainedModel):
|
||||
"""
|
||||
An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
|
||||
models.
|
||||
"""
|
||||
|
||||
config_class = EfficientViTConfig
|
||||
base_model_prefix = "efficientvit"
|
||||
main_input_name = "pixel_values"
|
||||
|
||||
def _init_weights(self, module):
|
||||
"""Initialize the weights"""
|
||||
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
||||
# Slightly different from the TF version which uses truncated_normal for initialization
|
||||
# cf https://github.com/pytorch/pytorch/pull/5617
|
||||
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
|
||||
if module.bias is not None:
|
||||
module.bias.data.zero_()
|
||||
elif isinstance(module, nn.Embedding):
|
||||
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
|
||||
if module.padding_idx is not None:
|
||||
module.weight.data[module.padding_idx].zero_()
|
||||
elif isinstance(module, nn.LayerNorm):
|
||||
module.bias.data.zero_()
|
||||
module.weight.data.fill_(1.0)
|
||||
|
||||
|
||||
class DecodeMLP(nn.Module):
|
||||
def __init__(self, input_dim, output_dim):
|
||||
super().__init__()
|
||||
self.proj = nn.Linear(input_dim, output_dim)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor):
|
||||
# Input is B, C, H, W
|
||||
hidden_states = hidden_states.flatten(2).transpose(1, 2)
|
||||
# Output is B, HW, C
|
||||
hidden_states = self.proj(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class DecodeHead(EfficientViTPreTrainedModel):
|
||||
def __init__(self, config: EfficientViTConfig):
|
||||
super().__init__(config)
|
||||
|
||||
# linear layers which will unify the channel dimension of each of the encoder blocks to the same config.decoder_hidden_size
|
||||
mlps = []
|
||||
for width in config.widths[1:]:
|
||||
mlp = DecodeMLP(
|
||||
input_dim=width, output_dim=config.decoder_layer_hidden_size
|
||||
)
|
||||
mlps.append(mlp)
|
||||
self.linear_c = nn.ModuleList(mlps)
|
||||
|
||||
# the following 3 layers implement the ConvModule of the original implementation
|
||||
self.linear_fuse = nn.Conv2d(
|
||||
in_channels=config.decoder_layer_hidden_size * config.num_stages,
|
||||
out_channels=config.decoder_hidden_size,
|
||||
kernel_size=1,
|
||||
bias=False,
|
||||
)
|
||||
self.batch_norm = nn.BatchNorm2d(config.decoder_hidden_size)
|
||||
self.activation = nn.ReLU()
|
||||
|
||||
self.dropout = nn.Dropout(config.classifier_dropout_prob)
|
||||
self.classifier = nn.Conv2d(
|
||||
config.decoder_hidden_size, config.num_labels, kernel_size=1
|
||||
)
|
||||
|
||||
self.config = config
|
||||
|
||||
def forward(self, encoder_hidden_states: torch.FloatTensor) -> torch.Tensor:
|
||||
batch_size = encoder_hidden_states[-1].shape[0]
|
||||
|
||||
all_hidden_states = ()
|
||||
for encoder_hidden_state, mlp in zip(encoder_hidden_states, self.linear_c):
|
||||
height, width = encoder_hidden_state.shape[2], encoder_hidden_state.shape[3]
|
||||
encoder_hidden_state = mlp(encoder_hidden_state) # Output is B, HW, C
|
||||
# Permute to B, C, HW
|
||||
encoder_hidden_state = encoder_hidden_state.permute(0, 2, 1)
|
||||
encoder_hidden_state = encoder_hidden_state.reshape(
|
||||
batch_size, -1, height, width
|
||||
)
|
||||
# upsample
|
||||
encoder_hidden_state = nn.functional.interpolate(
|
||||
encoder_hidden_state,
|
||||
size=encoder_hidden_states[0].size()[2:],
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
)
|
||||
all_hidden_states += (encoder_hidden_state,)
|
||||
|
||||
hidden_states = self.linear_fuse(torch.cat(all_hidden_states[::-1], dim=1))
|
||||
hidden_states = self.batch_norm(hidden_states)
|
||||
hidden_states = self.activation(hidden_states)
|
||||
|
||||
# logits are of shape (batch_size, num_labels, height/4, width/4)
|
||||
logits = self.classifier(hidden_states)
|
||||
|
||||
return logits
|
||||
|
||||
|
||||
class EfficientViTForSemanticSegmentation(
|
||||
S3DownloaderMixin, EfficientViTPreTrainedModel
|
||||
):
|
||||
def __init__(self, config, **kwargs):
|
||||
super().__init__(config)
|
||||
self.vit = EfficientVitLarge(config)
|
||||
self.decode_head = DecodeHead(config)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def forward(
|
||||
self, pixel_values: torch.FloatTensor
|
||||
) -> Union[Tuple, SemanticSegmenterOutput]:
|
||||
# Pixel values should be B,C,H,W
|
||||
encoder_hidden_states = self.vit(
|
||||
pixel_values,
|
||||
)
|
||||
|
||||
logits = self.decode_head(encoder_hidden_states)
|
||||
|
||||
# Apply sigmoid to get 0-1 output
|
||||
logits = torch.special.expit(logits)
|
||||
|
||||
return SemanticSegmenterOutput(
|
||||
loss=None, logits=logits, hidden_states=encoder_hidden_states
|
||||
)
|
||||
|
||||
|
||||
class EfficientViTForSemanticLayoutSegmentation(EfficientViTPreTrainedModel):
|
||||
def __init__(self, config, **kwargs):
|
||||
super().__init__(config, **kwargs)
|
||||
self.vit = EfficientVitLarge(config)
|
||||
self.decode_head = DecodeHead(config)
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def forward(
|
||||
self, pixel_values: torch.FloatTensor
|
||||
) -> Union[Tuple, SemanticSegmenterOutput]:
|
||||
# Pixel values should be B,C,H,W
|
||||
encoder_hidden_states = self.vit(
|
||||
pixel_values,
|
||||
)
|
||||
|
||||
logits = self.decode_head(encoder_hidden_states)
|
||||
|
||||
# Apply sigmoid to get 0-1 output
|
||||
logits = torch.special.expit(logits)
|
||||
|
||||
return SemanticSegmenterOutput(
|
||||
loss=None, logits=logits, hidden_states=encoder_hidden_states
|
||||
)
|
||||
@@ -0,0 +1,19 @@
|
||||
class FakeFuture:
|
||||
def __init__(self, func, *args, **kwargs):
|
||||
self._result = func(*args, **kwargs)
|
||||
|
||||
def result(self):
|
||||
return self._result
|
||||
|
||||
class FakeExecutor:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *excinfo):
|
||||
pass
|
||||
|
||||
def submit(self, fn, *args, **kwargs):
|
||||
return FakeFuture(fn, *args, **kwargs)
|
||||
@@ -0,0 +1,317 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Modified image processor class for Segformer based on transformers"""
|
||||
|
||||
import warnings
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from transformers.image_processing_utils import (
|
||||
BaseImageProcessor,
|
||||
BatchFeature,
|
||||
get_size_dict,
|
||||
)
|
||||
from transformers.image_transforms import to_channel_dimension_format
|
||||
from transformers.image_utils import (
|
||||
IMAGENET_DEFAULT_MEAN,
|
||||
IMAGENET_DEFAULT_STD,
|
||||
ChannelDimension,
|
||||
ImageInput,
|
||||
PILImageResampling,
|
||||
infer_channel_dimension_format,
|
||||
make_list_of_images,
|
||||
)
|
||||
from transformers.utils import TensorType
|
||||
|
||||
|
||||
import PIL.Image
|
||||
|
||||
from surya.common.s3 import S3DownloaderMixin
|
||||
|
||||
|
||||
class SegformerImageProcessor(S3DownloaderMixin, BaseImageProcessor):
|
||||
r"""
|
||||
Constructs a Segformer image processor.
|
||||
|
||||
Args:
|
||||
do_resize (`bool`, *optional*, defaults to `True`):
|
||||
Whether to resize the image's (height, width) dimensions to the specified `(size["height"],
|
||||
size["width"])`. Can be overridden by the `do_resize` parameter in the `preprocess` method.
|
||||
size (`Dict[str, int]` *optional*, defaults to `{"height": 512, "width": 512}`):
|
||||
Size of the output image after resizing. Can be overridden by the `size` parameter in the `preprocess`
|
||||
method.
|
||||
resample (`PILImageResampling`, *optional*, defaults to `Resampling.BILINEAR`):
|
||||
Resampling filter to use if resizing the image. Can be overridden by the `resample` parameter in the
|
||||
`preprocess` method.
|
||||
do_rescale (`bool`, *optional*, defaults to `True`):
|
||||
Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the `do_rescale`
|
||||
parameter in the `preprocess` method.
|
||||
rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
|
||||
Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`
|
||||
method.
|
||||
do_normalize (`bool`, *optional*, defaults to `True`):
|
||||
Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`
|
||||
method.
|
||||
image_mean (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`):
|
||||
Mean to use if normalizing the image. This is a float or list of floats the length of the number of
|
||||
channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
|
||||
image_std (`float` or `List[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`):
|
||||
Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
|
||||
number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
|
||||
do_reduce_labels (`bool`, *optional*, defaults to `False`):
|
||||
Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0 is
|
||||
used for background, and background itself is not included in all classes of a dataset (e.g. ADE20k). The
|
||||
background label will be replaced by 255. Can be overridden by the `do_reduce_labels` parameter in the
|
||||
`preprocess` method.
|
||||
"""
|
||||
|
||||
model_input_names = ["pixel_values"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
do_resize: bool = True,
|
||||
size: Dict[str, int] = None,
|
||||
resample: PILImageResampling = PILImageResampling.BILINEAR,
|
||||
do_rescale: bool = True,
|
||||
rescale_factor: Union[int, float] = 1 / 255,
|
||||
do_normalize: bool = True,
|
||||
image_mean: Optional[Union[float, List[float]]] = None,
|
||||
image_std: Optional[Union[float, List[float]]] = None,
|
||||
do_reduce_labels: bool = False,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
if "reduce_labels" in kwargs:
|
||||
warnings.warn(
|
||||
"The `reduce_labels` parameter is deprecated and will be removed in a future version. Please use "
|
||||
"`do_reduce_labels` instead.",
|
||||
FutureWarning,
|
||||
)
|
||||
do_reduce_labels = kwargs.pop("reduce_labels")
|
||||
|
||||
super().__init__(**kwargs)
|
||||
size = size if size is not None else {"height": 512, "width": 512}
|
||||
size = get_size_dict(size)
|
||||
self.do_resize = do_resize
|
||||
self.size = size
|
||||
self.resample = resample
|
||||
self.do_rescale = do_rescale
|
||||
self.rescale_factor = rescale_factor
|
||||
self.do_normalize = do_normalize
|
||||
self.image_mean = (
|
||||
image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
|
||||
)
|
||||
self.image_std = image_std if image_std is not None else IMAGENET_DEFAULT_STD
|
||||
self.do_reduce_labels = do_reduce_labels
|
||||
self._valid_processor_keys = [
|
||||
"images",
|
||||
"segmentation_maps",
|
||||
"do_resize",
|
||||
"size",
|
||||
"resample",
|
||||
"do_rescale",
|
||||
"rescale_factor",
|
||||
"do_normalize",
|
||||
"image_mean",
|
||||
"image_std",
|
||||
"do_reduce_labels",
|
||||
"return_tensors",
|
||||
"data_format",
|
||||
"input_data_format",
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, image_processor_dict: Dict[str, Any], **kwargs):
|
||||
"""
|
||||
Overrides the `from_dict` method from the base class to make sure `do_reduce_labels` is updated if image
|
||||
processor is created using from_dict and kwargs e.g. `SegformerImageProcessor.from_pretrained(checkpoint,
|
||||
reduce_labels=True)`
|
||||
"""
|
||||
image_processor_dict = image_processor_dict.copy()
|
||||
if "reduce_labels" in kwargs:
|
||||
image_processor_dict["reduce_labels"] = kwargs.pop("reduce_labels")
|
||||
return super().from_dict(image_processor_dict, **kwargs)
|
||||
|
||||
def _preprocess(
|
||||
self,
|
||||
image: ImageInput,
|
||||
do_resize: bool,
|
||||
do_rescale: bool,
|
||||
do_normalize: bool,
|
||||
size: Optional[Dict[str, int]] = None,
|
||||
resample: PILImageResampling = None,
|
||||
rescale_factor: Optional[float] = None,
|
||||
image_mean: Optional[Union[float, List[float]]] = None,
|
||||
image_std: Optional[Union[float, List[float]]] = None,
|
||||
input_data_format: Optional[Union[str, ChannelDimension]] = None,
|
||||
):
|
||||
if do_rescale:
|
||||
image = self.rescale(
|
||||
image=image, scale=rescale_factor, input_data_format=input_data_format
|
||||
)
|
||||
|
||||
if do_normalize:
|
||||
image = self.normalize(
|
||||
image=image,
|
||||
mean=image_mean,
|
||||
std=image_std,
|
||||
input_data_format=input_data_format,
|
||||
)
|
||||
|
||||
return image
|
||||
|
||||
def _preprocess_image(
|
||||
self,
|
||||
image: ImageInput,
|
||||
do_resize: bool = None,
|
||||
size: Dict[str, int] = None,
|
||||
resample: PILImageResampling = None,
|
||||
do_rescale: bool = None,
|
||||
rescale_factor: float = None,
|
||||
do_normalize: bool = None,
|
||||
image_mean: Optional[Union[float, List[float]]] = None,
|
||||
image_std: Optional[Union[float, List[float]]] = None,
|
||||
data_format: Optional[Union[str, ChannelDimension]] = None,
|
||||
input_data_format: Optional[Union[str, ChannelDimension]] = None,
|
||||
) -> np.ndarray:
|
||||
"""Preprocesses a single image."""
|
||||
# All transformations expect numpy arrays.
|
||||
if input_data_format is None:
|
||||
input_data_format = infer_channel_dimension_format(image)
|
||||
|
||||
image = self._preprocess(
|
||||
image=image,
|
||||
do_resize=do_resize,
|
||||
size=size,
|
||||
resample=resample,
|
||||
do_rescale=do_rescale,
|
||||
rescale_factor=rescale_factor,
|
||||
do_normalize=do_normalize,
|
||||
image_mean=image_mean,
|
||||
image_std=image_std,
|
||||
input_data_format=input_data_format,
|
||||
)
|
||||
if data_format is not None:
|
||||
image = to_channel_dimension_format(
|
||||
image, data_format, input_channel_dim=input_data_format
|
||||
)
|
||||
return image
|
||||
|
||||
def __call__(self, images, segmentation_maps=None, **kwargs):
|
||||
"""
|
||||
Preprocesses a batch of images and optionally segmentation maps.
|
||||
|
||||
Overrides the `__call__` method of the `Preprocessor` class so that both images and segmentation maps can be
|
||||
passed in as positional arguments.
|
||||
"""
|
||||
return super().__call__(images, segmentation_maps=segmentation_maps, **kwargs)
|
||||
|
||||
def preprocess(
|
||||
self,
|
||||
images: ImageInput,
|
||||
segmentation_maps: Optional[ImageInput] = None,
|
||||
do_resize: Optional[bool] = None,
|
||||
size: Optional[Dict[str, int]] = None,
|
||||
resample: PILImageResampling = None,
|
||||
do_rescale: Optional[bool] = None,
|
||||
rescale_factor: Optional[float] = None,
|
||||
do_normalize: Optional[bool] = None,
|
||||
image_mean: Optional[Union[float, List[float]]] = None,
|
||||
image_std: Optional[Union[float, List[float]]] = None,
|
||||
do_reduce_labels: Optional[bool] = None,
|
||||
return_tensors: Optional[Union[str, TensorType]] = None,
|
||||
data_format: ChannelDimension = ChannelDimension.FIRST,
|
||||
input_data_format: Optional[Union[str, ChannelDimension]] = None,
|
||||
**kwargs,
|
||||
) -> PIL.Image.Image:
|
||||
"""
|
||||
Preprocess an image or batch of images.
|
||||
|
||||
Args:
|
||||
images (`ImageInput`):
|
||||
Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
|
||||
passing in images with pixel values between 0 and 1, set `do_rescale=False`.
|
||||
segmentation_maps (`ImageInput`, *optional*):
|
||||
Segmentation map to preprocess.
|
||||
do_resize (`bool`, *optional*, defaults to `self.do_resize`):
|
||||
Whether to resize the image.
|
||||
size (`Dict[str, int]`, *optional*, defaults to `self.size`):
|
||||
Size of the image after `resize` is applied.
|
||||
resample (`int`, *optional*, defaults to `self.resample`):
|
||||
Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`, Only
|
||||
has an effect if `do_resize` is set to `True`.
|
||||
do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
|
||||
Whether to rescale the image values between [0 - 1].
|
||||
rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
|
||||
Rescale factor to rescale the image by if `do_rescale` is set to `True`.
|
||||
do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
|
||||
Whether to normalize the image.
|
||||
image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
|
||||
Image mean.
|
||||
image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
|
||||
Image standard deviation.
|
||||
do_reduce_labels (`bool`, *optional*, defaults to `self.do_reduce_labels`):
|
||||
Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0
|
||||
is used for background, and background itself is not included in all classes of a dataset (e.g.
|
||||
ADE20k). The background label will be replaced by 255.
|
||||
return_tensors (`str` or `TensorType`, *optional*):
|
||||
The type of tensors to return. Can be one of:
|
||||
- Unset: Return a list of `np.ndarray`.
|
||||
- `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
|
||||
- `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
|
||||
- `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
|
||||
- `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
|
||||
data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
|
||||
The channel dimension format for the output image. Can be one of:
|
||||
- `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
|
||||
- `ChannelDimension.LAST`: image in (height, width, num_channels) format.
|
||||
input_data_format (`ChannelDimension` or `str`, *optional*):
|
||||
The channel dimension format for the input image. If unset, the channel dimension format is inferred
|
||||
from the input image. Can be one of:
|
||||
- `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
|
||||
- `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
|
||||
- `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
|
||||
"""
|
||||
do_resize = do_resize if do_resize is not None else self.do_resize
|
||||
do_rescale = do_rescale if do_rescale is not None else self.do_rescale
|
||||
do_normalize = do_normalize if do_normalize is not None else self.do_normalize
|
||||
resample = resample if resample is not None else self.resample
|
||||
size = size if size is not None else self.size
|
||||
rescale_factor = (
|
||||
rescale_factor if rescale_factor is not None else self.rescale_factor
|
||||
)
|
||||
image_mean = image_mean if image_mean is not None else self.image_mean
|
||||
image_std = image_std if image_std is not None else self.image_std
|
||||
|
||||
images = make_list_of_images(images)
|
||||
images = [
|
||||
self._preprocess_image(
|
||||
image=img,
|
||||
do_resize=do_resize,
|
||||
resample=resample,
|
||||
size=size,
|
||||
do_rescale=do_rescale,
|
||||
rescale_factor=rescale_factor,
|
||||
do_normalize=do_normalize,
|
||||
image_mean=image_mean,
|
||||
image_std=image_std,
|
||||
data_format=data_format,
|
||||
input_data_format=input_data_format,
|
||||
)
|
||||
for img in images
|
||||
]
|
||||
|
||||
data = {"pixel_values": images}
|
||||
return BatchFeature(data=data, tensor_type=return_tensors)
|
||||
@@ -0,0 +1,12 @@
|
||||
from typing import List, Optional, Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from surya.common.polygon import PolygonBox
|
||||
|
||||
|
||||
class TextDetectionResult(BaseModel):
|
||||
bboxes: List[PolygonBox]
|
||||
heatmap: Optional[Any]
|
||||
affinity_map: Optional[Any]
|
||||
image_bbox: List[float]
|
||||
@@ -0,0 +1,36 @@
|
||||
import math
|
||||
from PIL import ImageOps
|
||||
|
||||
from surya.settings import settings
|
||||
|
||||
|
||||
def get_total_splits(image_size, height):
|
||||
img_height = list(image_size)[1]
|
||||
max_height = settings.DETECTOR_IMAGE_CHUNK_HEIGHT
|
||||
if img_height > max_height:
|
||||
num_splits = math.ceil(img_height / height)
|
||||
return num_splits
|
||||
return 1
|
||||
|
||||
|
||||
def split_image(img, height):
|
||||
# This will not modify/return the original image - it will either crop, or copy the image
|
||||
img_height = list(img.size)[1]
|
||||
max_height = settings.DETECTOR_IMAGE_CHUNK_HEIGHT
|
||||
if img_height > max_height:
|
||||
num_splits = math.ceil(img_height / height)
|
||||
splits = []
|
||||
split_heights = []
|
||||
for i in range(num_splits):
|
||||
top = i * height
|
||||
bottom = (i + 1) * height
|
||||
if bottom > img_height:
|
||||
bottom = img_height
|
||||
cropped = img.crop((0, top, img.size[0], bottom))
|
||||
chunk_height = bottom - top
|
||||
if chunk_height < height:
|
||||
cropped = ImageOps.pad(cropped, (img.size[0], height), color=255, centering=(0, 0))
|
||||
splits.append(cropped)
|
||||
split_heights.append(chunk_height)
|
||||
return splits, split_heights
|
||||
return [img.copy()], [img_height]
|
||||
@@ -0,0 +1,97 @@
|
||||
import json
|
||||
import logging
|
||||
import logging.config
|
||||
import time
|
||||
import uuid
|
||||
|
||||
import yaml
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from surya.endpoint.legacy import router as legacy_router
|
||||
from surya.endpoint.openai import router as openai_router
|
||||
from surya.endpoint.service import _error_response
|
||||
|
||||
with open("logger.yaml", "r", encoding="utf-8") as f:
|
||||
config = yaml.safe_load(f.read())
|
||||
logging.config.dictConfig(config)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
app = FastAPI(
|
||||
title="OCR SERVICE API SERVICE",
|
||||
version="1.0",
|
||||
docs_url="/v1/api/ai/swagger",
|
||||
openapi_url="/v1/api/ai/openapi.json",
|
||||
)
|
||||
|
||||
|
||||
@app.exception_handler(RequestValidationError)
|
||||
async def validation_exception_handler(request: Request, exc: RequestValidationError):
|
||||
request_id = getattr(request.state, "request_id", "-")
|
||||
logger.warning(
|
||||
"request_validation_failed request_id=%s method=%s path=%s errors=%s",
|
||||
request_id,
|
||||
request.method,
|
||||
request.url.path,
|
||||
exc.errors(),
|
||||
)
|
||||
return JSONResponse(
|
||||
status_code=200,
|
||||
content={
|
||||
"data": [],
|
||||
"message": json.dumps({"error info": exc.errors()}, default=str),
|
||||
"code": 422,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@app.middleware("http")
|
||||
async def allow_openapi_unauthorized(request: Request, call_next):
|
||||
if request.url.path == "/v1/api/ai/openapi.json":
|
||||
response = await call_next(request)
|
||||
return response
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
@app.middleware("http")
|
||||
async def log_request_response(request, call_next):
|
||||
request_id = uuid.uuid4().hex
|
||||
request.state.request_id = request_id
|
||||
start = time.perf_counter()
|
||||
client = request.client.host if request.client else "-"
|
||||
logger.info(
|
||||
"request_start request_id=%s method=%s path=%s client=%s",
|
||||
request_id,
|
||||
request.method,
|
||||
request.url.path,
|
||||
client,
|
||||
)
|
||||
try:
|
||||
response = await call_next(request)
|
||||
except Exception as e:
|
||||
duration_ms = (time.perf_counter() - start) * 1000
|
||||
logger.exception(
|
||||
"request_unhandled_exception request_id=%s method=%s path=%s duration_ms=%.2f",
|
||||
request_id,
|
||||
request.method,
|
||||
request.url.path,
|
||||
duration_ms,
|
||||
exc_info=True,
|
||||
)
|
||||
return JSONResponse(status_code=500, content=_error_response(e))
|
||||
duration_ms = (time.perf_counter() - start) * 1000
|
||||
logger.info(
|
||||
"request_end request_id=%s method=%s path=%s status_code=%s duration_ms=%.2f",
|
||||
request_id,
|
||||
request.method,
|
||||
request.url.path,
|
||||
response.status_code,
|
||||
duration_ms,
|
||||
)
|
||||
return response
|
||||
|
||||
|
||||
app.include_router(legacy_router)
|
||||
app.include_router(openai_router)
|
||||
@@ -0,0 +1,157 @@
|
||||
import io
|
||||
import time
|
||||
|
||||
from fastapi import APIRouter, File, Request, UploadFile
|
||||
from PIL import Image
|
||||
|
||||
from tools import (
|
||||
ocr,
|
||||
text_detection,
|
||||
layout_detection,
|
||||
table_recognition,
|
||||
extract_text_from_image,
|
||||
)
|
||||
from vllm_batcher import vllm_ocr_batcher
|
||||
from vllm_tools import vllm_backend_info
|
||||
from surya.endpoint.schemas import ApiResponse, Info
|
||||
from surya.endpoint.service import (
|
||||
_dump_model_or_list,
|
||||
_error_response,
|
||||
_load_base64_image,
|
||||
_log_exception,
|
||||
_ocr_response_data,
|
||||
_success_response,
|
||||
logger,
|
||||
)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.post("/home")
|
||||
def home():
|
||||
return "<h1>Welcome to SURYA OCR API!</h1>"
|
||||
|
||||
|
||||
@router.post("/v1/api/ai/suya_ocr", include_in_schema=False, response_model=ApiResponse)
|
||||
@router.post("/v1/api/ai/suya_ocr/", response_model=ApiResponse)
|
||||
def run_ocr(p: Info, request: Request):
|
||||
try:
|
||||
pil_image, pil_image_highres = _load_base64_image(p)
|
||||
rec_img, pred, box_img = ocr(
|
||||
pil_image,
|
||||
pil_image_highres,
|
||||
p.skip_text_detection,
|
||||
p.recognize_math,
|
||||
with_bboxes=p.ocr_with_boxes,
|
||||
)
|
||||
return _success_response(_ocr_response_data(pred))
|
||||
except Exception as e:
|
||||
results = _error_response(e)
|
||||
_log_exception("suya_ocr", request, e, results)
|
||||
return results
|
||||
|
||||
|
||||
@router.get("/v1/api/ai/suya_ocr_vllm/health", response_model=ApiResponse)
|
||||
async def suya_ocr_vllm_health():
|
||||
return _success_response(vllm_backend_info())
|
||||
|
||||
|
||||
@router.post("/v1/api/ai/suya_ocr_vllm", include_in_schema=False, response_model=ApiResponse)
|
||||
@router.post("/v1/api/ai/suya_ocr_vllm/", response_model=ApiResponse)
|
||||
def run_ocr_vllm(p: Info, request: Request):
|
||||
try:
|
||||
pil_image, pil_image_highres = _load_base64_image(p)
|
||||
start = time.perf_counter()
|
||||
_, pred, _ = vllm_ocr_batcher.submit(
|
||||
pil_image,
|
||||
pil_image_highres,
|
||||
skip_text_detection=p.skip_text_detection,
|
||||
recognize_math=p.recognize_math,
|
||||
with_bboxes=False,
|
||||
request_id=getattr(request.state, "request_id", "-"),
|
||||
)
|
||||
logger.info(
|
||||
"api_vllm_submit_wait_complete request_id=%s duration_ms=%.2f",
|
||||
getattr(request.state, "request_id", "-"),
|
||||
(time.perf_counter() - start) * 1000,
|
||||
)
|
||||
return _success_response(_ocr_response_data(pred))
|
||||
except Exception as e:
|
||||
results = _error_response(e)
|
||||
_log_exception("suya_ocr_vllm", request, e, results)
|
||||
return results
|
||||
|
||||
|
||||
@router.post("/v1/api/ai/suya_text_det", include_in_schema=False, response_model=ApiResponse)
|
||||
@router.post("/v1/api/ai/suya_text_det/", response_model=ApiResponse)
|
||||
def run_text_det(p: Info, request: Request):
|
||||
try:
|
||||
pil_image, _ = _load_base64_image(p)
|
||||
det_img, text_pred = text_detection(pil_image)
|
||||
text_lines = text_pred.model_dump(exclude=["heatmap", "affinity_map"])
|
||||
return _success_response({"text_lines": text_lines})
|
||||
except Exception as e:
|
||||
results = _error_response(e)
|
||||
_log_exception("suya_text_det", request, e, results)
|
||||
return results
|
||||
|
||||
|
||||
@router.post("/v1/api/ai/suya_layout_det", include_in_schema=False, response_model=ApiResponse)
|
||||
@router.post("/v1/api/ai/suya_layout_det/", response_model=ApiResponse)
|
||||
def run_layout_det(p: Info, request: Request):
|
||||
try:
|
||||
pil_image, _ = _load_base64_image(p)
|
||||
layout_img, pred = layout_detection(pil_image)
|
||||
text_lines = pred.model_dump(exclude=["segmentation_map"])
|
||||
return _success_response({"text_lines": text_lines})
|
||||
except Exception as e:
|
||||
results = _error_response(e)
|
||||
_log_exception("suya_layout_det", request, e, results)
|
||||
return results
|
||||
|
||||
|
||||
@router.post("/v1/api/ai/suya_table_rec", include_in_schema=False, response_model=ApiResponse)
|
||||
@router.post("/v1/api/ai/suya_table_rec/", response_model=ApiResponse)
|
||||
def run_table_rec(p: Info, request: Request):
|
||||
try:
|
||||
pil_image, pil_image_highres = _load_base64_image(p)
|
||||
|
||||
table_img, pred = table_recognition(
|
||||
pil_image, pil_image_highres, p.skip_table_detection
|
||||
)
|
||||
|
||||
text_json = _dump_model_or_list(pred)
|
||||
text_lines = ""
|
||||
if isinstance(pred, list):
|
||||
text_lines = "\n".join(
|
||||
[item.html for item in pred if getattr(item, "html", None)]
|
||||
)
|
||||
elif hasattr(pred, "text_lines"):
|
||||
text_lines = "\n".join([p.text for p in pred.text_lines])
|
||||
|
||||
return _success_response({"ocr_text_json": text_json, "text_lines": text_lines})
|
||||
except Exception as e:
|
||||
results = _error_response(e)
|
||||
_log_exception("suya_table_rec", request, e, results)
|
||||
return results
|
||||
|
||||
|
||||
@router.post("/image2text", response_model=ApiResponse)
|
||||
async def image_to_text(request: Request, file: UploadFile = File(...)):
|
||||
try:
|
||||
contents = await file.read()
|
||||
image = Image.open(io.BytesIO(contents))
|
||||
logger.info(
|
||||
"image2text_upload request_id=%s filename=%s content_type=%s size_bytes=%s image_size=%s",
|
||||
getattr(request.state, "request_id", "-"),
|
||||
file.filename,
|
||||
file.content_type,
|
||||
len(contents),
|
||||
image.size,
|
||||
)
|
||||
full_text = extract_text_from_image(image)
|
||||
return _success_response(full_text)
|
||||
except Exception as e:
|
||||
results = _error_response(e)
|
||||
_log_exception("image2text", request, e, results)
|
||||
return results
|
||||
@@ -0,0 +1,108 @@
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Literal, Optional, Union
|
||||
|
||||
from fastapi import APIRouter, Request
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from surya.endpoint import service
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
class _ImageUrl(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
url: str
|
||||
|
||||
|
||||
class _ContentPart(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
type: str
|
||||
image_url: Optional[_ImageUrl] = None
|
||||
text: Optional[str] = None
|
||||
|
||||
|
||||
class _ChatMessage(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
role: str
|
||||
content: Union[str, List[_ContentPart]]
|
||||
|
||||
|
||||
class ChatCompletionRequest(BaseModel):
|
||||
# The OpenAI SDK merges `extra_body={...}` into the TOP LEVEL of the request
|
||||
# JSON, so the OCR control flags live here rather than nested under a key.
|
||||
model_config = ConfigDict(extra="allow")
|
||||
model: str = "surya-ocr"
|
||||
messages: List[_ChatMessage]
|
||||
stream: Optional[bool] = False
|
||||
skip_text_detection: bool = False
|
||||
recognize_math: bool = False
|
||||
skip_table_detection: bool = False
|
||||
ocr_with_boxes: bool = True
|
||||
mode: Literal["block", "full_page", "table"] = "block"
|
||||
|
||||
|
||||
def _openai_error(message: str, err_type: str, status_code: int) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
status_code=status_code,
|
||||
content={"error": {"message": message, "type": err_type, "code": None}},
|
||||
)
|
||||
|
||||
|
||||
def _build_chat_completion(
|
||||
content: str, model: str, ocr_json: Any, elapsed: Optional[float]
|
||||
) -> Dict[str, Any]:
|
||||
# Token accounting happens inside vLLM and is not surfaced to this layer, so
|
||||
# usage counts are reported as 0 (prompt_tokens intentionally 0).
|
||||
return {
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion",
|
||||
"created": int(time.time()),
|
||||
"model": model,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"message": {"role": "assistant", "content": content},
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
|
||||
"surya": {"ocr_text_json": ocr_json, "elapsed_seconds": elapsed},
|
||||
}
|
||||
|
||||
|
||||
@router.post("/v1/chat/completions")
|
||||
def chat_completions(req: ChatCompletionRequest, request: Request):
|
||||
request_id = getattr(request.state, "request_id", "-")
|
||||
if req.stream:
|
||||
return _openai_error(
|
||||
"streaming is not supported by this OCR endpoint",
|
||||
"invalid_request_error",
|
||||
400,
|
||||
)
|
||||
try:
|
||||
image = service.extract_last_image(req.messages)
|
||||
except service.OcrRequestError as exc:
|
||||
return _openai_error(str(exc), "invalid_request_error", 400)
|
||||
|
||||
try:
|
||||
if req.mode == "table":
|
||||
data = service.table_image(image, req.skip_table_detection)
|
||||
else:
|
||||
skip = req.skip_text_detection or req.mode == "full_page"
|
||||
data = service.ocr_via_batcher(
|
||||
image,
|
||||
skip_text_detection=skip,
|
||||
recognize_math=req.recognize_math,
|
||||
request_id=request_id,
|
||||
)
|
||||
except service.OcrRequestError as exc:
|
||||
return _openai_error(str(exc), "invalid_request_error", 400)
|
||||
except Exception as exc:
|
||||
service._log_exception("chat_completions", request, exc, {})
|
||||
return _openai_error(str(exc), "internal_error", 500)
|
||||
|
||||
content = data.get("text_lines", "") or ""
|
||||
ocr_json = data.get("ocr_text_json") if req.ocr_with_boxes else None
|
||||
return _build_chat_completion(content, req.model, ocr_json, data.get("elapsed_seconds"))
|
||||
@@ -0,0 +1,46 @@
|
||||
import base64
|
||||
import binascii
|
||||
from typing import Any, Literal, Optional
|
||||
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
Field,
|
||||
PrivateAttr,
|
||||
StrictBool,
|
||||
field_validator,
|
||||
model_validator,
|
||||
)
|
||||
|
||||
|
||||
class Info(BaseModel):
|
||||
_decoded_file: bytes = PrivateAttr(default=b"")
|
||||
|
||||
file: str # base64 string
|
||||
type: Literal["png", "jpg", "jpeg", "gif"] # image type like 'png', 'jpg'
|
||||
skip_text_detection: Optional[StrictBool] = Field(False)
|
||||
skip_table_detection: Optional[StrictBool] = Field(False)
|
||||
recognize_math: Optional[StrictBool] = Field(False)
|
||||
ocr_with_boxes: Optional[StrictBool] = Field(True)
|
||||
|
||||
@field_validator("type", mode="before")
|
||||
@classmethod
|
||||
def normalize_type(cls, value: str) -> str:
|
||||
if not isinstance(value, str):
|
||||
raise ValueError("type must be a string")
|
||||
return value.lower()
|
||||
|
||||
@model_validator(mode="after")
|
||||
def decode_base64_file(self):
|
||||
if not self.file:
|
||||
raise ValueError("file must be a non-empty base64 string")
|
||||
try:
|
||||
self._decoded_file = base64.b64decode(self.file, validate=True)
|
||||
except binascii.Error as exc:
|
||||
raise ValueError("file must be valid base64") from exc
|
||||
return self
|
||||
|
||||
|
||||
class ApiResponse(BaseModel):
|
||||
data: Any
|
||||
message: str
|
||||
code: int
|
||||
@@ -0,0 +1,166 @@
|
||||
import base64
|
||||
import binascii
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import traceback
|
||||
from typing import Any, Dict
|
||||
|
||||
import requests
|
||||
from fastapi import Request
|
||||
from PIL import Image
|
||||
|
||||
import tools
|
||||
from vllm_batcher import vllm_ocr_batcher
|
||||
from vllm_tools import page_ocr_to_response
|
||||
|
||||
logger = logging.getLogger("surya.endpoint")
|
||||
|
||||
|
||||
class OcrRequestError(Exception):
|
||||
"""Client-side error in an OCR request. Maps to HTTP 400."""
|
||||
|
||||
|
||||
# ---- helpers moved verbatim from api.py ------------------------------------
|
||||
|
||||
def _load_base64_image(p, copy_highres: bool = False):
|
||||
start = time.perf_counter()
|
||||
filetype = p.type.lower()
|
||||
allowed_types = ["png", "jpg", "jpeg", "gif"]
|
||||
if filetype not in allowed_types:
|
||||
raise ValueError("Unsupported file type")
|
||||
pil_image = Image.open(io.BytesIO(p._decoded_file)).convert("RGB")
|
||||
logger.info(
|
||||
"api_decode_image duration_ms=%.2f size_bytes=%s image_size=%s copy_highres=%s",
|
||||
(time.perf_counter() - start) * 1000,
|
||||
len(p._decoded_file),
|
||||
pil_image.size,
|
||||
copy_highres,
|
||||
)
|
||||
return pil_image, pil_image.copy() if copy_highres else pil_image
|
||||
|
||||
|
||||
def _error_response(e: Exception):
|
||||
tb = traceback.extract_tb(e.__traceback__)
|
||||
frame = tb[-1] if tb else None
|
||||
return {
|
||||
"data": [],
|
||||
"message": json.dumps(
|
||||
{
|
||||
"error info": str(e),
|
||||
"error at": frame.filename if frame else "unknown",
|
||||
"errort line at": frame.lineno if frame else 0,
|
||||
}
|
||||
),
|
||||
"code": 500,
|
||||
}
|
||||
|
||||
|
||||
def _success_response(data: Any) -> Dict[str, Any]:
|
||||
return {"data": data, "message": "success", "code": 200}
|
||||
|
||||
|
||||
def _log_exception(endpoint: str, request: Request, e: Exception, results: Dict[str, Any]):
|
||||
request_id = getattr(request.state, "request_id", "-")
|
||||
logger.exception(
|
||||
"endpoint_failed request_id=%s endpoint=%s code=%s message=%s",
|
||||
request_id,
|
||||
endpoint,
|
||||
results.get("code"),
|
||||
results.get("message"),
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
|
||||
def _ocr_response_data(pred):
|
||||
if hasattr(pred, "text_lines"):
|
||||
text_json = pred.model_dump()
|
||||
text_lines = "\n".join([p.text for p in pred.text_lines])
|
||||
return {"ocr_text_json": text_json, "text_lines": text_lines}
|
||||
return page_ocr_to_response(pred)
|
||||
|
||||
|
||||
def _dump_model_or_list(value):
|
||||
if hasattr(value, "model_dump"):
|
||||
return value.model_dump()
|
||||
if isinstance(value, list):
|
||||
return [
|
||||
item.model_dump() if hasattr(item, "model_dump") else item
|
||||
for item in value
|
||||
]
|
||||
return value
|
||||
|
||||
|
||||
# ---- new helpers for the OpenAI endpoint -----------------------------------
|
||||
|
||||
def load_image_from_url(url: str) -> Image.Image:
|
||||
"""Decode an OpenAI `image_url` (data: base64 URL or http(s): URL) to RGB."""
|
||||
if url.startswith("data:"):
|
||||
header, _, b64 = url.partition(",")
|
||||
if "base64" not in header or not b64:
|
||||
raise OcrRequestError("image_url data URL must be base64-encoded")
|
||||
try:
|
||||
raw = base64.b64decode(b64, validate=True)
|
||||
except binascii.Error as exc:
|
||||
raise OcrRequestError("image_url contains invalid base64") from exc
|
||||
elif url.startswith("http://") or url.startswith("https://"):
|
||||
try:
|
||||
resp = requests.get(url, timeout=10)
|
||||
resp.raise_for_status()
|
||||
except Exception as exc:
|
||||
raise OcrRequestError(f"could not fetch image_url: {exc}") from exc
|
||||
raw = resp.content
|
||||
else:
|
||||
raise OcrRequestError("image_url must be a data: or http(s): URL")
|
||||
try:
|
||||
return Image.open(io.BytesIO(raw)).convert("RGB")
|
||||
except Exception as exc:
|
||||
raise OcrRequestError(f"could not decode image: {exc}") from exc
|
||||
|
||||
|
||||
def extract_last_image(messages) -> Image.Image:
|
||||
"""Return the last image found across the chat messages' content parts."""
|
||||
url = None
|
||||
for msg in messages:
|
||||
content = msg.content
|
||||
if isinstance(content, list):
|
||||
for part in content:
|
||||
if getattr(part, "type", None) == "image_url" and part.image_url:
|
||||
url = part.image_url.url
|
||||
if url is None:
|
||||
raise OcrRequestError("no image_url content part found in messages")
|
||||
return load_image_from_url(url)
|
||||
|
||||
|
||||
def ocr_via_batcher(
|
||||
image: Image.Image,
|
||||
*,
|
||||
skip_text_detection: bool,
|
||||
recognize_math: bool,
|
||||
request_id: str,
|
||||
) -> Dict[str, Any]:
|
||||
"""Run page OCR through the shared request batcher. The annotated image is
|
||||
discarded, so we request with_bboxes=False to maximise batch coalescing with
|
||||
the primary /suya_ocr_vllm endpoint."""
|
||||
_, pred, _ = vllm_ocr_batcher.submit(
|
||||
image,
|
||||
image,
|
||||
skip_text_detection=skip_text_detection,
|
||||
recognize_math=recognize_math,
|
||||
with_bboxes=False,
|
||||
request_id=request_id,
|
||||
)
|
||||
return _ocr_response_data(pred)
|
||||
|
||||
|
||||
def table_image(image: Image.Image, skip_table_detection: bool) -> Dict[str, Any]:
|
||||
"""Run table-structure recognition; mirrors the /suya_table_rec response shape."""
|
||||
_, pred = tools.table_recognition(image, image.copy(), skip_table_detection)
|
||||
text_json = _dump_model_or_list(pred)
|
||||
text_lines = ""
|
||||
if isinstance(pred, list):
|
||||
text_lines = "\n".join([item.html for item in pred if getattr(item, "html", None)])
|
||||
elif hasattr(pred, "text_lines"):
|
||||
text_lines = "\n".join([line.text for line in pred.text_lines])
|
||||
return {"ocr_text_json": text_json, "text_lines": text_lines, "elapsed_seconds": None}
|
||||
@@ -0,0 +1,116 @@
|
||||
"""Surya inference manager.
|
||||
|
||||
One process owns one SuryaInferenceManager. The manager wraps a single backend
|
||||
(vllm | llamacpp) which speaks OpenAI-compatible chat completions.
|
||||
|
||||
Predictors take the manager via explicit injection at construction time.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from typing import List, Optional
|
||||
|
||||
from surya.inference.backends.base import Backend
|
||||
from surya.inference.schema import BatchInputItem, BatchOutputItem
|
||||
from surya.logging import get_logger
|
||||
from surya.settings import settings
|
||||
from surya.timing import timing_span
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def _has_nvidia_gpu() -> bool:
|
||||
"""True if an NVIDIA GPU is present on this host.
|
||||
|
||||
We deliberately do *not* rely solely on ``torch.cuda.is_available()``:
|
||||
the installed torch wheel's CUDA build can be newer than the host driver
|
||||
(PyPI's default wheel tracks the latest CUDA), in which case torch reports
|
||||
no CUDA even on a perfectly good GPU box. That would silently route us to
|
||||
the CPU llama.cpp backend on a machine that should be running vllm. So we
|
||||
take torch's word when it *does* see CUDA, and otherwise fall back to
|
||||
probing for the GPU directly via ``nvidia-smi``.
|
||||
"""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Instant, load-independent check: the NVIDIA device node only exists when
|
||||
# a GPU + driver are present. Preferred over nvidia-smi because nvidia-smi
|
||||
# can block for several seconds on a GPU under heavy load, which would race
|
||||
# a timeout and falsely report "no GPU".
|
||||
if os.path.exists("/dev/nvidia0"):
|
||||
return True
|
||||
|
||||
nvidia_smi = shutil.which("nvidia-smi")
|
||||
if not nvidia_smi:
|
||||
return False
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[nvidia_smi, "-L"], capture_output=True, text=True, timeout=15
|
||||
)
|
||||
return result.returncode == 0 and "GPU" in result.stdout
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _autodetect_backend() -> str:
|
||||
if settings.SURYA_INFERENCE_BACKEND:
|
||||
return settings.SURYA_INFERENCE_BACKEND
|
||||
# NVIDIA GPU → vllm, mps/cpu → llamacpp
|
||||
if _has_nvidia_gpu():
|
||||
return "vllm"
|
||||
return "llamacpp"
|
||||
|
||||
|
||||
def _build_backend(method: str) -> Backend:
|
||||
method = method.lower()
|
||||
if method == "vllm":
|
||||
from surya.inference.backends.vllm import VllmBackend
|
||||
|
||||
return VllmBackend()
|
||||
if method == "llamacpp":
|
||||
from surya.inference.backends.llamacpp import LlamaCppBackend
|
||||
return LlamaCppBackend()
|
||||
raise ValueError(
|
||||
f"Unknown inference backend {method!r}. Supported: 'vllm', 'llamacpp'."
|
||||
)
|
||||
|
||||
|
||||
class SuryaInferenceManager:
|
||||
"""Single entry point for VLM inference. Construct once per process."""
|
||||
|
||||
def __init__(self, method: Optional[str] = None, lazy: bool = True):
|
||||
self.method = method or _autodetect_backend()
|
||||
self.backend: Backend = _build_backend(self.method)
|
||||
if not lazy:
|
||||
self.backend.start()
|
||||
|
||||
def start(self) -> None:
|
||||
self.backend.start()
|
||||
|
||||
def stop(self) -> None:
|
||||
self.backend.stop()
|
||||
|
||||
def generate(self, batch: List[BatchInputItem]) -> List[BatchOutputItem]:
|
||||
with timing_span("surya_manager_generate", backend=self.method, item_count=len(batch)):
|
||||
return self.backend.generate(batch)
|
||||
|
||||
|
||||
# Module-level lazy singleton for callers that don't want explicit construction
|
||||
# (notebooks, ad-hoc scripts). Surya's own models.py and marker should use
|
||||
# explicit construction.
|
||||
_default_manager: Optional[SuryaInferenceManager] = None
|
||||
|
||||
|
||||
def get_default_manager() -> SuryaInferenceManager:
|
||||
global _default_manager
|
||||
if _default_manager is None:
|
||||
_default_manager = SuryaInferenceManager()
|
||||
return _default_manager
|
||||
@@ -0,0 +1,2 @@
|
||||
from surya.inference.backends.base import Backend as Backend
|
||||
from surya.inference.backends.base import ServerHandle as ServerHandle
|
||||
@@ -0,0 +1,31 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import List
|
||||
|
||||
|
||||
from surya.inference.schema import BatchInputItem, BatchOutputItem
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServerHandle:
|
||||
base_url: str # e.g. "http://127.0.0.1:8765/v1"
|
||||
model_name: str # what gets passed in OpenAI `model` field
|
||||
spawned_by_us: bool # if True, we manage atexit cleanup
|
||||
|
||||
|
||||
class Backend:
|
||||
"""Abstract backend. Concrete backends own server lifecycle + generation."""
|
||||
|
||||
name: str # "vllm" | "llamacpp"
|
||||
|
||||
def start(self) -> ServerHandle:
|
||||
"""Idempotent: probe → attach if alive, else spawn. Returns handle."""
|
||||
raise NotImplementedError
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Stop the server if we spawned it."""
|
||||
raise NotImplementedError
|
||||
|
||||
def generate(self, batch: List[BatchInputItem]) -> List[BatchOutputItem]:
|
||||
raise NotImplementedError
|
||||
@@ -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,
|
||||
)
|
||||
@@ -0,0 +1,257 @@
|
||||
"""Shared OpenAI-compatible chat completions client. Used by vllm + llama.cpp.
|
||||
|
||||
Both servers expose `/v1/chat/completions` with the same request/response shape,
|
||||
so this module is the single point of HTTP contact for both backends.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
import math
|
||||
import os
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import List, Optional
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from surya.inference.prompts import PROMPT_MAPPING
|
||||
from surya.inference.schema import (
|
||||
BatchInputItem,
|
||||
BatchOutputItem,
|
||||
GenerationResult,
|
||||
)
|
||||
from surya.inference.util import detect_repeat_token, scale_to_fit
|
||||
from surya.logging import get_logger
|
||||
from surya.timing import get_current_timing, set_current_timing, timing_span
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def resolve_max_workers(batch_len: int, max_inflight: int) -> int:
|
||||
"""Concurrent HTTP workers for a batch: as many as the batch needs, capped
|
||||
by max_inflight so we keep vLLM's sequence slots full without over-queueing."""
|
||||
return max(1, min(batch_len, max_inflight))
|
||||
|
||||
|
||||
def encode_image_b64(image: Image.Image) -> tuple[str, str]:
|
||||
image_format = os.getenv("SUYA_VLLM_IMAGE_FORMAT", "JPEG").upper()
|
||||
if image_format not in {"JPEG", "PNG"}:
|
||||
raise ValueError("SUYA_VLLM_IMAGE_FORMAT must be JPEG or PNG")
|
||||
|
||||
buf = io.BytesIO()
|
||||
if image_format == "JPEG":
|
||||
quality = int(os.getenv("SUYA_VLLM_JPEG_QUALITY", "92"))
|
||||
image.save(buf, format="JPEG", quality=quality, subsampling=0)
|
||||
mime_type = "image/jpeg"
|
||||
else:
|
||||
image.save(buf, format="PNG")
|
||||
mime_type = "image/png"
|
||||
view = buf.getbuffer()
|
||||
try:
|
||||
return base64.b64encode(view).decode("ascii"), mime_type
|
||||
finally:
|
||||
view.release()
|
||||
|
||||
|
||||
def _build_messages(image: Image.Image, prompt: str):
|
||||
with timing_span("openai_encode_image_b64", image_size=image.size):
|
||||
image_b64, mime_type = encode_image_b64(image)
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:{mime_type};base64,{image_b64}"},
|
||||
},
|
||||
{"type": "text", "text": prompt},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _mean_token_prob(logprobs_content) -> Optional[float]:
|
||||
if not logprobs_content:
|
||||
return None
|
||||
probs = []
|
||||
for tok in logprobs_content:
|
||||
lp = (
|
||||
tok.get("logprob")
|
||||
if isinstance(tok, dict)
|
||||
else getattr(tok, "logprob", None)
|
||||
)
|
||||
if lp is None:
|
||||
continue
|
||||
probs.append(math.exp(lp))
|
||||
if not probs:
|
||||
return None
|
||||
return sum(probs) / len(probs)
|
||||
|
||||
|
||||
def _generate_one(
|
||||
item: BatchInputItem,
|
||||
client,
|
||||
model_name: str,
|
||||
max_tokens_default: int,
|
||||
temperature: float,
|
||||
top_p: float,
|
||||
timeout: float,
|
||||
request_logprobs_default: bool,
|
||||
) -> GenerationResult:
|
||||
prompt = item.prompt or PROMPT_MAPPING[item.prompt_type]
|
||||
with timing_span("openai_scale_image", prompt_type=item.prompt_type, image_size=item.image.size):
|
||||
image = scale_to_fit(item.image)
|
||||
with timing_span("openai_build_messages", prompt_type=item.prompt_type):
|
||||
messages = _build_messages(image, prompt)
|
||||
|
||||
max_tokens = item.max_tokens or max_tokens_default
|
||||
request_logprobs = item.request_logprobs or request_logprobs_default
|
||||
|
||||
kwargs = dict(
|
||||
model=model_name,
|
||||
messages=messages,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
timeout=timeout,
|
||||
)
|
||||
if request_logprobs:
|
||||
kwargs["logprobs"] = True
|
||||
|
||||
# Structured output: prefer OpenAI-standard response_format (works on both
|
||||
# vllm and llama.cpp). Fall back to vllm's extra_body for guided_regex.
|
||||
if item.guided_json is not None:
|
||||
kwargs["response_format"] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "structured_output",
|
||||
"schema": item.guided_json,
|
||||
"strict": True,
|
||||
},
|
||||
}
|
||||
if item.guided_regex is not None:
|
||||
kwargs.setdefault("extra_body", {})["guided_regex"] = item.guided_regex
|
||||
|
||||
try:
|
||||
with timing_span(
|
||||
"openai_chat_completion",
|
||||
prompt_type=item.prompt_type,
|
||||
max_tokens=max_tokens,
|
||||
):
|
||||
completion = client.chat.completions.create(**kwargs)
|
||||
raw = completion.choices[0].message.content or ""
|
||||
token_count = completion.usage.completion_tokens if completion.usage else 0
|
||||
with timing_span(
|
||||
"openai_parse_completion",
|
||||
prompt_type=item.prompt_type,
|
||||
token_count=token_count,
|
||||
):
|
||||
mean_p = None
|
||||
logprobs_content = None
|
||||
if request_logprobs:
|
||||
choice = completion.choices[0]
|
||||
lp = getattr(choice, "logprobs", None)
|
||||
if lp is not None:
|
||||
content = getattr(lp, "content", None)
|
||||
if content is not None:
|
||||
logprobs_content = [
|
||||
c.model_dump() if hasattr(c, "model_dump") else c
|
||||
for c in content
|
||||
]
|
||||
mean_p = _mean_token_prob(content)
|
||||
return GenerationResult(
|
||||
raw=raw,
|
||||
token_count=token_count,
|
||||
error=False,
|
||||
mean_token_prob=mean_p,
|
||||
logprobs=logprobs_content,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Inference error: {e}")
|
||||
return GenerationResult(raw="", token_count=0, error=True)
|
||||
|
||||
|
||||
def _should_retry(
|
||||
result: GenerationResult,
|
||||
retries: int,
|
||||
max_retries: int,
|
||||
) -> bool:
|
||||
if retries >= max_retries:
|
||||
return False
|
||||
if result.error:
|
||||
return True
|
||||
has_repeat = detect_repeat_token(result.raw) or (
|
||||
len(result.raw) > 50 and detect_repeat_token(result.raw, cut_from_end=50)
|
||||
)
|
||||
return has_repeat
|
||||
|
||||
|
||||
def chat_completions_batch(
|
||||
batch: List[BatchInputItem],
|
||||
client,
|
||||
model_name: str,
|
||||
max_tokens_default: int = 2048,
|
||||
temperature: float = 0.0,
|
||||
top_p: float = 0.1,
|
||||
timeout: float = 600.0,
|
||||
max_workers: Optional[int] = None,
|
||||
max_retries: int = 3,
|
||||
request_logprobs_default: bool = True,
|
||||
) -> List[BatchOutputItem]:
|
||||
"""Run a batch of items through the chat completions endpoint with concurrent workers."""
|
||||
if not batch:
|
||||
return []
|
||||
if max_workers is None:
|
||||
max_workers = min(64, len(batch))
|
||||
collector = get_current_timing()
|
||||
|
||||
def _process(item: BatchInputItem) -> BatchOutputItem:
|
||||
if collector is not None:
|
||||
set_current_timing(collector)
|
||||
result = _generate_one(
|
||||
item,
|
||||
client=client,
|
||||
model_name=model_name,
|
||||
max_tokens_default=max_tokens_default,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
timeout=timeout,
|
||||
request_logprobs_default=request_logprobs_default,
|
||||
)
|
||||
retries = 0
|
||||
while _should_retry(result, retries, max_retries):
|
||||
backoff = 1.5 * (retries + 1) if result.error else 0
|
||||
if backoff:
|
||||
time.sleep(backoff)
|
||||
retry_temp = min(temperature + 0.2 * (retries + 1), 0.8)
|
||||
retry_top_p = 0.95 if not result.error else top_p
|
||||
result = _generate_one(
|
||||
item,
|
||||
client=client,
|
||||
model_name=model_name,
|
||||
max_tokens_default=max_tokens_default,
|
||||
temperature=retry_temp,
|
||||
top_p=retry_top_p,
|
||||
timeout=timeout,
|
||||
request_logprobs_default=request_logprobs_default,
|
||||
)
|
||||
retries += 1
|
||||
return BatchOutputItem(
|
||||
raw=result.raw,
|
||||
token_count=result.token_count,
|
||||
error=result.error,
|
||||
mean_token_prob=result.mean_token_prob,
|
||||
logprobs=result.logprobs,
|
||||
metadata=item.metadata,
|
||||
)
|
||||
|
||||
with timing_span(
|
||||
"openai_batch_threadpool",
|
||||
item_count=len(batch),
|
||||
max_workers=max_workers,
|
||||
max_retries=max_retries,
|
||||
):
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
return list(executor.map(_process, batch))
|
||||
@@ -0,0 +1,351 @@
|
||||
"""Server lifecycle: probe, filelock, sentinel, atexit cleanup.
|
||||
|
||||
Pattern: probe `/health` → if alive return handle → else acquire lock, re-probe,
|
||||
spawn detached, write sentinel, register atexit kill (only the spawner cleans up).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
import subprocess
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Callable, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
from surya.logging import get_logger
|
||||
from surya.settings import settings
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def _cache_dir() -> Path:
|
||||
base = Path(os.path.expanduser("~/.cache/datalab/surya"))
|
||||
base.mkdir(parents=True, exist_ok=True)
|
||||
return base
|
||||
|
||||
|
||||
def _sentinel_path(backend: str) -> Path:
|
||||
return _cache_dir() / f"{backend}_server.json"
|
||||
|
||||
|
||||
def _lock_path(backend: str) -> Path:
|
||||
return _cache_dir() / f"{backend}_server.lock"
|
||||
|
||||
|
||||
def find_free_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def probe_health(base_url: str, timeout: float = 1.0) -> bool:
|
||||
"""Returns True if the server reports healthy at /health."""
|
||||
try:
|
||||
# llama.cpp returns 200 on /health when ready; vllm returns 200 on /health too.
|
||||
with httpx.Client(timeout=timeout) as client:
|
||||
r = client.get(f"{base_url}/health")
|
||||
return r.status_code == 200
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def wait_for_health(
|
||||
base_url: str, total_timeout: float = 300.0, interval: float = 1.0
|
||||
) -> bool:
|
||||
deadline = time.time() + total_timeout
|
||||
while time.time() < deadline:
|
||||
if probe_health(base_url):
|
||||
return True
|
||||
time.sleep(interval)
|
||||
return False
|
||||
|
||||
|
||||
def probe_model_id(openai_base: str, timeout: float = 5.0) -> Optional[str]:
|
||||
"""Returns the model id reported by the running server, or None on failure."""
|
||||
try:
|
||||
with httpx.Client(timeout=timeout) as client:
|
||||
r = client.get(f"{openai_base}/models")
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
models = data.get("data") or []
|
||||
if models:
|
||||
return models[0].get("id")
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
@dataclass
|
||||
class SpawnedServer:
|
||||
base_url: str # full openai base, e.g. "http://127.0.0.1:8765/v1"
|
||||
health_url: str # base for /health, e.g. "http://127.0.0.1:8765"
|
||||
model_name: str # what to pass as `model`
|
||||
pid: Optional[int]
|
||||
backend: str
|
||||
spawned_by_us: bool
|
||||
|
||||
|
||||
class SpawnError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def _read_sentinel(backend: str) -> Optional[dict]:
|
||||
p = _sentinel_path(backend)
|
||||
if not p.exists():
|
||||
return None
|
||||
try:
|
||||
return json.loads(p.read_text())
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _write_sentinel(backend: str, data: dict) -> None:
|
||||
_sentinel_path(backend).write_text(json.dumps(data))
|
||||
|
||||
|
||||
def _delete_sentinel(backend: str) -> None:
|
||||
p = _sentinel_path(backend)
|
||||
if p.exists():
|
||||
try:
|
||||
p.unlink()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _stop_process(pid: int, name: str) -> None:
|
||||
try:
|
||||
# Graceful first
|
||||
os.kill(pid, 15) # SIGTERM
|
||||
for _ in range(20):
|
||||
try:
|
||||
os.kill(pid, 0) # still alive?
|
||||
except ProcessLookupError:
|
||||
logger.info(f"Stopped {name} (pid {pid})")
|
||||
return
|
||||
time.sleep(0.5)
|
||||
# Hard
|
||||
os.kill(pid, 9)
|
||||
logger.warning(f"Force-killed {name} (pid {pid})")
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to stop {name} (pid {pid}): {e}")
|
||||
|
||||
|
||||
def _capture_server_logs(handle: "SpawnHandle", tail: int = 100) -> str:
|
||||
"""Best-effort tail of a server's logs, for surfacing startup failures."""
|
||||
try:
|
||||
if handle.cleanup_kind == "docker":
|
||||
r = subprocess.run(
|
||||
["docker", "logs", "--tail", str(tail), handle.cleanup_id],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=15,
|
||||
)
|
||||
return (r.stdout or "") + (r.stderr or "") or "(no docker logs)"
|
||||
# llama.cpp process backend logs to this file (see llamacpp.py)
|
||||
log_path = Path("~/.cache/datalab/surya/llamacpp_server.log").expanduser()
|
||||
if log_path.exists():
|
||||
lines = log_path.read_text(errors="replace").splitlines()
|
||||
return "\n".join(lines[-tail:]) or "(empty log)"
|
||||
except Exception as e:
|
||||
return f"(could not capture logs: {e})"
|
||||
return "(no logs available)"
|
||||
|
||||
|
||||
def _stop_docker_container(name: str) -> None:
|
||||
try:
|
||||
subprocess.run(
|
||||
["docker", "stop", name], check=False, capture_output=True, timeout=30
|
||||
)
|
||||
logger.info(f"Stopped docker container {name}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to stop docker container {name}: {e}")
|
||||
|
||||
|
||||
def attach_or_spawn(
|
||||
backend: str,
|
||||
expected_model_name: str,
|
||||
spawn_fn: Callable[[int], "SpawnHandle"],
|
||||
health_url_for: Callable[[int], str],
|
||||
openai_url_for: Callable[[int], str],
|
||||
startup_timeout: float = 600.0,
|
||||
) -> SpawnedServer:
|
||||
"""Generic attach-or-spawn with file lock and sentinel.
|
||||
|
||||
`spawn_fn(port)` must launch the server detached and return a SpawnHandle
|
||||
with `pid` (int or None for docker) and a `cleanup_id` (e.g. container name).
|
||||
"""
|
||||
# 0. If user pinned an external URL, attach without lock
|
||||
if settings.SURYA_INFERENCE_URL:
|
||||
base_url = settings.SURYA_INFERENCE_URL.rstrip("/")
|
||||
health_url = base_url[: -len("/v1")] if base_url.endswith("/v1") else base_url
|
||||
if not probe_health(health_url):
|
||||
raise SpawnError(
|
||||
f"SURYA_INFERENCE_URL={base_url} is not reachable at /health. "
|
||||
"Start the server or unset the variable."
|
||||
)
|
||||
model_name = probe_model_id(base_url) or expected_model_name
|
||||
if model_name != expected_model_name:
|
||||
raise SpawnError(
|
||||
f"Model mismatch at {base_url}: expected {expected_model_name!r}, got {model_name!r}. "
|
||||
"Stop the running server or unset SURYA_INFERENCE_URL."
|
||||
)
|
||||
return SpawnedServer(
|
||||
base_url=base_url,
|
||||
health_url=health_url,
|
||||
model_name=model_name,
|
||||
pid=None,
|
||||
backend=backend,
|
||||
spawned_by_us=False,
|
||||
)
|
||||
|
||||
# 1. Probe sentinel without lock
|
||||
existing = _read_sentinel(backend)
|
||||
if existing:
|
||||
port = existing.get("port")
|
||||
pid = existing.get("pid")
|
||||
if port and probe_health(health_url_for(port)):
|
||||
running_model = probe_model_id(openai_url_for(port)) or expected_model_name
|
||||
if running_model != expected_model_name:
|
||||
raise SpawnError(
|
||||
f"Existing {backend} server on port {port} serves {running_model!r}, "
|
||||
f"expected {expected_model_name!r}. Stop it before continuing."
|
||||
)
|
||||
logger.info(f"Attaching to existing {backend} server on port {port}")
|
||||
return SpawnedServer(
|
||||
base_url=openai_url_for(port),
|
||||
health_url=health_url_for(port),
|
||||
model_name=running_model,
|
||||
pid=pid,
|
||||
backend=backend,
|
||||
spawned_by_us=False,
|
||||
)
|
||||
else:
|
||||
_delete_sentinel(backend)
|
||||
|
||||
if not settings.SURYA_INFERENCE_AUTOSTART:
|
||||
raise SpawnError(
|
||||
f"No running {backend} server and SURYA_INFERENCE_AUTOSTART is False. "
|
||||
"Set the variable to True or start the server manually."
|
||||
)
|
||||
|
||||
# 2. Acquire filelock to prevent races
|
||||
try:
|
||||
from filelock import FileLock
|
||||
except ImportError as e:
|
||||
raise SpawnError(
|
||||
"filelock is required for server spawn. pip install filelock"
|
||||
) from e
|
||||
|
||||
lock = FileLock(str(_lock_path(backend)), timeout=120)
|
||||
with lock:
|
||||
# Re-check sentinel inside the lock
|
||||
existing = _read_sentinel(backend)
|
||||
if existing:
|
||||
port = existing.get("port")
|
||||
if port and probe_health(health_url_for(port)):
|
||||
running_model = (
|
||||
probe_model_id(openai_url_for(port)) or expected_model_name
|
||||
)
|
||||
if running_model != expected_model_name:
|
||||
raise SpawnError(
|
||||
f"Existing {backend} server on port {port} serves {running_model!r}, "
|
||||
f"expected {expected_model_name!r}."
|
||||
)
|
||||
return SpawnedServer(
|
||||
base_url=openai_url_for(port),
|
||||
health_url=health_url_for(port),
|
||||
model_name=running_model,
|
||||
pid=existing.get("pid"),
|
||||
backend=backend,
|
||||
spawned_by_us=False,
|
||||
)
|
||||
|
||||
# 3. Spawn fresh
|
||||
port = settings.SURYA_INFERENCE_PORT or find_free_port()
|
||||
logger.info(f"Spawning {backend} server on port {port}")
|
||||
spawn_handle = spawn_fn(port)
|
||||
|
||||
# 4. Write sentinel
|
||||
_write_sentinel(
|
||||
backend,
|
||||
{
|
||||
"port": port,
|
||||
"pid": spawn_handle.pid,
|
||||
"model": expected_model_name,
|
||||
"backend": backend,
|
||||
"cleanup_id": spawn_handle.cleanup_id,
|
||||
"cleanup_kind": spawn_handle.cleanup_kind,
|
||||
},
|
||||
)
|
||||
|
||||
# 5. Register atexit cleanup (only spawner). Skipped when keep-alive is
|
||||
# set so the server outlives this process and later commands attach to
|
||||
# it via the sentinel. (_cleanup is still callable below on startup
|
||||
# failure, where we always tear a half-started server down.)
|
||||
def _cleanup():
|
||||
try:
|
||||
if spawn_handle.cleanup_kind == "docker":
|
||||
_stop_docker_container(spawn_handle.cleanup_id)
|
||||
elif spawn_handle.cleanup_kind == "process":
|
||||
if spawn_handle.pid:
|
||||
_stop_process(spawn_handle.pid, backend)
|
||||
finally:
|
||||
_delete_sentinel(backend)
|
||||
|
||||
if settings.SURYA_INFERENCE_KEEP_ALIVE:
|
||||
logger.info(
|
||||
f"keep-alive: {backend} server on port {port} will stay up "
|
||||
f"after exit (cleanup_id={spawn_handle.cleanup_id!r})"
|
||||
)
|
||||
else:
|
||||
atexit.register(_cleanup)
|
||||
|
||||
# 6. Wait for health
|
||||
health_url = health_url_for(port)
|
||||
if not wait_for_health(health_url, total_timeout=startup_timeout):
|
||||
# Grab the server's own logs *before* cleanup tears the (--rm)
|
||||
# container down, otherwise the actual failure reason is lost and
|
||||
# all the caller sees is this timeout.
|
||||
logs = _capture_server_logs(spawn_handle)
|
||||
_cleanup()
|
||||
raise SpawnError(
|
||||
f"{backend} server failed to become healthy at {health_url} "
|
||||
f"within {startup_timeout}s.\n"
|
||||
f"--- last {backend} server logs ---\n{logs}"
|
||||
)
|
||||
|
||||
# 7. Verify model name
|
||||
running_model = probe_model_id(openai_url_for(port))
|
||||
if running_model and running_model != expected_model_name:
|
||||
logger.warning(
|
||||
f"{backend} server reports model={running_model!r} "
|
||||
f"but expected {expected_model_name!r}; using reported name."
|
||||
)
|
||||
expected_model_name = running_model
|
||||
|
||||
logger.info(
|
||||
f"{backend} server ready on port {port} (model={expected_model_name})"
|
||||
)
|
||||
return SpawnedServer(
|
||||
base_url=openai_url_for(port),
|
||||
health_url=health_url,
|
||||
model_name=expected_model_name,
|
||||
pid=spawn_handle.pid,
|
||||
backend=backend,
|
||||
spawned_by_us=True,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SpawnHandle:
|
||||
pid: Optional[int]
|
||||
cleanup_id: str # container name for docker, str(pid) for process
|
||||
cleanup_kind: str # "docker" | "process"
|
||||
@@ -0,0 +1,224 @@
|
||||
"""vllm backend: spawns the vllm/vllm-openai docker image with MTP=2."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from typing import List, Optional
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
from surya.inference.backends.base import Backend, ServerHandle
|
||||
from surya.inference.backends.openai_client import chat_completions_batch, resolve_max_workers
|
||||
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
|
||||
from surya.timing import timing_span
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
# 24GB baseline (re-tune for surya-2 once benchmarks land)
|
||||
BASELINE_VRAM_GB = 24
|
||||
BASELINE_MAX_BATCHED_TOKENS = 8192
|
||||
BASELINE_MAX_NUM_SEQS = 32
|
||||
|
||||
GPU_VRAM_GB = {
|
||||
"b300": 270,
|
||||
"b200": 180,
|
||||
"h200": 141,
|
||||
"h100": 80,
|
||||
"a100-80": 80,
|
||||
"a100": 40,
|
||||
"a100-40": 40,
|
||||
"l40s": 48,
|
||||
"a10": 24,
|
||||
"l4": 24,
|
||||
"5090": 32,
|
||||
"4090": 24,
|
||||
"3090": 24,
|
||||
"t4": 16,
|
||||
}
|
||||
|
||||
|
||||
def _gpu_settings(gpu: str) -> tuple[int, int]:
|
||||
vram = GPU_VRAM_GB.get(gpu)
|
||||
if vram is None:
|
||||
available = ", ".join(sorted(GPU_VRAM_GB.keys()))
|
||||
raise SpawnError(f"Unknown VLLM_GPU_TYPE {gpu!r}. Available: {available}")
|
||||
ratio = vram / BASELINE_VRAM_GB
|
||||
raw_tokens = BASELINE_MAX_BATCHED_TOKENS * ratio
|
||||
max_batched_tokens = max(1024, 2 ** math.floor(math.log2(raw_tokens)))
|
||||
max_num_seqs = max(8, (int(BASELINE_MAX_NUM_SEQS * ratio) // 8) * 8)
|
||||
return max_batched_tokens, max_num_seqs
|
||||
|
||||
|
||||
def _resolve_docker_binary() -> str:
|
||||
found = shutil.which("docker")
|
||||
if found:
|
||||
return found
|
||||
raise SpawnError(
|
||||
"docker binary not found. Install Docker (https://docs.docker.com/get-docker/) "
|
||||
"and ensure the daemon is running."
|
||||
)
|
||||
|
||||
|
||||
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 VllmBackend(Backend):
|
||||
name = "vllm"
|
||||
|
||||
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 docker.
|
||||
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="docker"
|
||||
),
|
||||
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=settings.VLLM_API_KEY, base_url=self.handle.base_url
|
||||
)
|
||||
return self.handle
|
||||
|
||||
if os.getenv("SUYA_ALLOW_NESTED_DOCKER", "false").lower() not in {"1", "true", "yes"}:
|
||||
raise SpawnError(
|
||||
"Nested Docker vLLM startup is disabled. Start vLLM in this "
|
||||
"container and set SURYA_INFERENCE_URL, for example "
|
||||
"http://127.0.0.1:8000/v1."
|
||||
)
|
||||
|
||||
docker = _resolve_docker_binary()
|
||||
max_batched_tokens, max_num_seqs = _gpu_settings(settings.VLLM_GPU_TYPE)
|
||||
|
||||
def spawn_fn(port: int) -> SpawnHandle:
|
||||
container_name = f"surya-vllm-{port}"
|
||||
hf_cache = os.path.expanduser(settings.DOCKER_HF_CACHE_PATH)
|
||||
cmd = [
|
||||
docker,
|
||||
"run",
|
||||
"--rm",
|
||||
"-d",
|
||||
"--name",
|
||||
container_name,
|
||||
"--runtime",
|
||||
"nvidia",
|
||||
"--gpus",
|
||||
f"device={settings.VLLM_GPUS}",
|
||||
"-v",
|
||||
f"{hf_cache}:/root/.cache/huggingface",
|
||||
"-p",
|
||||
f"{port}:8000",
|
||||
"--ipc=host",
|
||||
settings.VLLM_DOCKER_IMAGE,
|
||||
"--model",
|
||||
settings.SURYA_MODEL_CHECKPOINT,
|
||||
"--no-enforce-eager",
|
||||
"--max-num-seqs",
|
||||
str(max_num_seqs),
|
||||
"--dtype",
|
||||
settings.VLLM_DTYPE,
|
||||
"--max-model-len",
|
||||
str(settings.VLLM_MAX_MODEL_LEN),
|
||||
"--max-num-batched-tokens",
|
||||
str(max_batched_tokens),
|
||||
"--gpu-memory-utilization",
|
||||
str(settings.VLLM_GPU_MEMORY_UTILIZATION),
|
||||
"--enable-prefix-caching",
|
||||
"--mm-processor-kwargs",
|
||||
json.dumps({"min_pixels": 3136, "max_pixels": 6291456}),
|
||||
"--served-model-name",
|
||||
settings.SURYA_MODEL_CHECKPOINT,
|
||||
]
|
||||
if settings.VLLM_ENABLE_MTP:
|
||||
spec_config = json.dumps(
|
||||
{
|
||||
"method": "mtp",
|
||||
"num_speculative_tokens": settings.VLLM_MTP_TOKENS,
|
||||
}
|
||||
)
|
||||
cmd.extend(["--speculative-config", spec_config])
|
||||
for extra in (settings.VLLM_EXTRA_ARGS or "").split():
|
||||
cmd.append(extra)
|
||||
logger.info(f"Spawning: {' '.join(cmd)}")
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, check=False)
|
||||
if result.returncode != 0:
|
||||
raise SpawnError(f"docker run failed: {result.stderr or result.stdout}")
|
||||
return SpawnHandle(
|
||||
pid=None, cleanup_id=container_name, cleanup_kind="docker"
|
||||
)
|
||||
|
||||
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=settings.VLLM_API_KEY,
|
||||
base_url=self.handle.base_url,
|
||||
)
|
||||
return self.handle
|
||||
|
||||
def stop(self) -> None:
|
||||
self.handle = None
|
||||
self._client = None
|
||||
|
||||
def generate(self, batch: List[BatchInputItem]) -> List[BatchOutputItem]:
|
||||
if self.handle is None or self._client is None:
|
||||
with timing_span("vllm_backend_start"):
|
||||
self.start()
|
||||
with timing_span(
|
||||
"vllm_backend_generate",
|
||||
item_count=len(batch),
|
||||
parallel=settings.SURYA_INFERENCE_PARALLEL,
|
||||
):
|
||||
return chat_completions_batch(
|
||||
batch,
|
||||
client=self._client,
|
||||
model_name=self.handle.model_name,
|
||||
timeout=settings.SURYA_INFERENCE_TIMEOUT_SECONDS,
|
||||
max_workers=resolve_max_workers(
|
||||
len(batch), settings.SURYA_INFERENCE_MAX_INFLIGHT
|
||||
),
|
||||
max_retries=settings.SURYA_INFERENCE_MAX_RETRIES,
|
||||
request_logprobs_default=settings.SURYA_INFERENCE_LOGPROBS,
|
||||
)
|
||||
@@ -0,0 +1,182 @@
|
||||
"""Parsers for the three task outputs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Tuple
|
||||
|
||||
|
||||
from surya.logging import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
# ---- Layout (LAYOUT_PROMPT) -------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParsedLayoutBlock:
|
||||
label: str
|
||||
bbox: Tuple[float, float, float, float] # 0-1000 normalized
|
||||
count: int # multiple of 50, model's token estimate
|
||||
|
||||
|
||||
_JSON_ARRAY_RE = re.compile(r"\[.*\]", re.DOTALL)
|
||||
|
||||
|
||||
def _strip_fences(text: str) -> str:
|
||||
cleaned = text.strip()
|
||||
if cleaned.startswith("```"):
|
||||
cleaned = re.sub(r"^```[a-zA-Z]*\n", "", cleaned)
|
||||
cleaned = re.sub(r"\n```\s*$", "", cleaned)
|
||||
return cleaned
|
||||
|
||||
|
||||
def _coerce_bbox(bbox) -> Tuple[float, float, float, float]:
|
||||
if isinstance(bbox, str):
|
||||
parts = [float(x) for x in bbox.replace(",", " ").split()]
|
||||
else:
|
||||
parts = [float(x) for x in bbox]
|
||||
if len(parts) != 4:
|
||||
raise ValueError(f"Bad bbox: {bbox!r}")
|
||||
return (parts[0], parts[1], parts[2], parts[3])
|
||||
|
||||
|
||||
def _coerce_count(value) -> int:
|
||||
if value is None:
|
||||
return 0
|
||||
try:
|
||||
return max(0, int(value))
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
def parse_layout(text: str) -> List[ParsedLayoutBlock]:
|
||||
"""Pull the JSON array out of LAYOUT_PROMPT output and convert to typed blocks.
|
||||
|
||||
Tolerates code fences, missing fields, and stringified bboxes.
|
||||
"""
|
||||
cleaned = _strip_fences(text)
|
||||
m = _JSON_ARRAY_RE.search(cleaned)
|
||||
if not m:
|
||||
raise ValueError(f"No JSON array found in layout output: {text[:500]!r}")
|
||||
raw = json.loads(m.group(0))
|
||||
out: List[ParsedLayoutBlock] = []
|
||||
for item in raw:
|
||||
try:
|
||||
bbox = _coerce_bbox(item["bbox"])
|
||||
except (KeyError, ValueError) as e:
|
||||
logger.warning(f"Skipping layout block with bad bbox: {e}")
|
||||
continue
|
||||
label = str(item.get("label", "block"))
|
||||
count = _coerce_count(item.get("count"))
|
||||
out.append(ParsedLayoutBlock(label=label, bbox=bbox, count=count))
|
||||
return out
|
||||
|
||||
|
||||
# ---- Table rec (TABLE_REC_PROMPT) ------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParsedTableElement:
|
||||
label: str # "Row" or "Col"
|
||||
bbox: Tuple[float, float, float, float]
|
||||
|
||||
|
||||
def parse_table_rec(text: str) -> List[ParsedTableElement]:
|
||||
"""Parse JSON array of {label: "Row"|"Col", bbox: "x0 y0 x1 y1"} from
|
||||
TABLE_REC_PROMPT output. Returns a flat list of Row + Col elements;
|
||||
cell derivation is the caller's job."""
|
||||
cleaned = _strip_fences(text)
|
||||
m = _JSON_ARRAY_RE.search(cleaned)
|
||||
if not m:
|
||||
raise ValueError(f"No JSON array found in table_rec output: {text[:500]!r}")
|
||||
raw = json.loads(m.group(0))
|
||||
out: List[ParsedTableElement] = []
|
||||
for item in raw:
|
||||
label = str(item.get("label", "")).strip()
|
||||
if label not in ("Row", "Col"):
|
||||
continue
|
||||
try:
|
||||
bbox = _coerce_bbox(item["bbox"])
|
||||
except (KeyError, ValueError):
|
||||
continue
|
||||
out.append(ParsedTableElement(label=label, bbox=bbox))
|
||||
return out
|
||||
|
||||
|
||||
# ---- Block HTML (BLOCK_PROMPT for full table path / general block path) ---
|
||||
|
||||
|
||||
def clean_block_html(html: str) -> str:
|
||||
"""Light cleanup of model-emitted HTML for a single block.
|
||||
|
||||
Strips code fences, leading/trailing whitespace. Does NOT validate against
|
||||
ALLOWED_TAGS — the model is expected to comply, and downstream consumers
|
||||
can sanitize further if needed.
|
||||
"""
|
||||
cleaned = _strip_fences(html).strip()
|
||||
return cleaned
|
||||
|
||||
|
||||
# ---- Full-page fallback (HIGH_ACCURACY_BBOX_PROMPT) -----------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParsedFullPageBlock:
|
||||
label: str
|
||||
bbox: Tuple[float, float, float, float] # 0-1000 normalized
|
||||
html: str # inner HTML of the wrapping div
|
||||
|
||||
|
||||
def parse_full_page_html(text: str) -> List[ParsedFullPageBlock]:
|
||||
"""Parse output of HIGH_ACCURACY_BBOX_PROMPT — top-level <div data-bbox=...
|
||||
data-label=...>inner HTML</div> blocks. Returns one entry per top-level div."""
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
cleaned = _strip_fences(text).strip()
|
||||
if not cleaned:
|
||||
return []
|
||||
# The model outputs a sequence of top-level divs (no surrounding root).
|
||||
# BeautifulSoup parses fine without one.
|
||||
soup = BeautifulSoup(cleaned, "html.parser")
|
||||
divs = soup.find_all("div", recursive=False)
|
||||
out: List[ParsedFullPageBlock] = []
|
||||
for div in divs:
|
||||
label = div.get("data-label")
|
||||
bbox_str = div.get("data-bbox")
|
||||
if not label or not bbox_str:
|
||||
continue
|
||||
try:
|
||||
parts = [float(x) for x in bbox_str.split()]
|
||||
except ValueError:
|
||||
continue
|
||||
if len(parts) != 4:
|
||||
continue
|
||||
# Strip nested data-bbox attrs from the inner HTML so downstream
|
||||
# consumers don't see model debug info on every child element.
|
||||
for tag in div.find_all(attrs={"data-bbox": True}):
|
||||
del tag["data-bbox"]
|
||||
for tag in div.find_all(attrs={"data-label": True}):
|
||||
del tag["data-label"]
|
||||
inner = "".join(str(c) for c in div.contents).strip()
|
||||
out.append(
|
||||
ParsedFullPageBlock(
|
||||
label=str(label),
|
||||
bbox=(parts[0], parts[1], parts[2], parts[3]),
|
||||
html=inner,
|
||||
)
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def denorm_bbox(bbox, img_w: int, img_h: int, scale: int = 1000):
|
||||
x0, y0, x1, y1 = bbox
|
||||
return (
|
||||
x0 / scale * img_w,
|
||||
y0 / scale * img_h,
|
||||
x1 / scale * img_w,
|
||||
y1 / scale * img_h,
|
||||
)
|
||||
@@ -0,0 +1,158 @@
|
||||
"""Prompt strings for surya2. The exact wording is the model's training-time
|
||||
contract — do not paraphrase without retraining."""
|
||||
|
||||
from surya.inference.schema import PROMPT_TYPE_BLOCK as PROMPT_TYPE_BLOCK
|
||||
from surya.inference.schema import (
|
||||
PROMPT_TYPE_HIGH_ACCURACY_BBOX as PROMPT_TYPE_HIGH_ACCURACY_BBOX,
|
||||
)
|
||||
from surya.inference.schema import PROMPT_TYPE_LAYOUT as PROMPT_TYPE_LAYOUT
|
||||
from surya.inference.schema import PROMPT_TYPE_TABLE_REC as PROMPT_TYPE_TABLE_REC
|
||||
|
||||
ALLOWED_TAGS = [
|
||||
"math",
|
||||
"br",
|
||||
"i",
|
||||
"b",
|
||||
"u",
|
||||
"del",
|
||||
"sup",
|
||||
"sub",
|
||||
"table",
|
||||
"tr",
|
||||
"td",
|
||||
"p",
|
||||
"th",
|
||||
"div",
|
||||
"pre",
|
||||
"h1",
|
||||
"h2",
|
||||
"h3",
|
||||
"h4",
|
||||
"h5",
|
||||
"ul",
|
||||
"ol",
|
||||
"li",
|
||||
"input",
|
||||
"a",
|
||||
"span",
|
||||
"img",
|
||||
"hr",
|
||||
"tbody",
|
||||
"small",
|
||||
"caption",
|
||||
"strong",
|
||||
"thead",
|
||||
"big",
|
||||
"code",
|
||||
"chem",
|
||||
]
|
||||
|
||||
ALLOWED_ATTRIBUTES = [
|
||||
"class",
|
||||
"colspan",
|
||||
"rowspan",
|
||||
"display",
|
||||
"checked",
|
||||
"type",
|
||||
"border",
|
||||
"value",
|
||||
"style",
|
||||
"href",
|
||||
"alt",
|
||||
"align",
|
||||
"data-bbox",
|
||||
"data-label",
|
||||
]
|
||||
|
||||
# Block labels we don't run OCR on.
|
||||
SKIP_OCR_LABELS = {"Figure", "Image", "Diagram", "Blank-Page"}
|
||||
|
||||
LAYOUT_PROMPT = (
|
||||
"Output the layout of this image as JSON. Each entry is a dict with "
|
||||
'"label", "bbox", and "count" fields. Bbox is x0 y0 x1 y1, normalized 0-1000.'
|
||||
)
|
||||
|
||||
BLOCK_PROMPT = "OCR this block image to HTML."
|
||||
|
||||
TABLE_REC_PROMPT = (
|
||||
"Output the table rows then columns as JSON. Each entry is a dict with "
|
||||
'"label" ("Row" or "Col") and "bbox" (x0 y0 x1 y1, normalized 0-1000).'
|
||||
)
|
||||
|
||||
HIGH_ACCURACY_BBOX_PROMPT = (
|
||||
"OCR this image to HTML. Each block is a div with data-label and data-bbox "
|
||||
"(x0 y0 x1 y1, normalized 0-1000)."
|
||||
)
|
||||
|
||||
|
||||
PROMPT_MAPPING = {
|
||||
"layout": LAYOUT_PROMPT,
|
||||
"block": BLOCK_PROMPT,
|
||||
"table_rec": TABLE_REC_PROMPT,
|
||||
"high_accuracy_bbox": HIGH_ACCURACY_BBOX_PROMPT,
|
||||
}
|
||||
|
||||
|
||||
# JSON schema for LAYOUT_PROMPT — enforced via vllm guided decoding so the
|
||||
# model can't emit malformed JSON. bbox is a "x0 y0 x1 y1" string (model's
|
||||
# training-time format); count is a non-negative integer.
|
||||
LAYOUT_LABEL_SET = [
|
||||
"Caption",
|
||||
"Footnote",
|
||||
"Equation-Block",
|
||||
"List-Group",
|
||||
"Page-Header",
|
||||
"Page-Footer",
|
||||
"Image",
|
||||
"Section-Header",
|
||||
"Table",
|
||||
"Text",
|
||||
"Complex-Block",
|
||||
"Code-Block",
|
||||
"Form",
|
||||
"Table-Of-Contents",
|
||||
"Figure",
|
||||
"Chemical-Block",
|
||||
"Diagram",
|
||||
"Bibliography",
|
||||
"Blank-Page",
|
||||
]
|
||||
|
||||
LAYOUT_JSON_SCHEMA = {
|
||||
"type": "array",
|
||||
"maxItems": 200,
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"label": {"type": "string", "enum": LAYOUT_LABEL_SET},
|
||||
"bbox": {
|
||||
"type": "string",
|
||||
"pattern": r"^\d{1,4} \d{1,4} \d{1,4} \d{1,4}$",
|
||||
},
|
||||
"count": {"type": "integer", "minimum": 0, "maximum": 10000},
|
||||
},
|
||||
"required": ["label", "bbox", "count"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# JSON schema for TABLE_REC_PROMPT — array of {label: Row|Col, bbox: "x0 y0 x1 y1"}.
|
||||
TABLE_REC_LABEL_SET = ["Row", "Col"]
|
||||
|
||||
TABLE_REC_JSON_SCHEMA = {
|
||||
"type": "array",
|
||||
"maxItems": 200,
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"label": {"type": "string", "enum": TABLE_REC_LABEL_SET},
|
||||
"bbox": {
|
||||
"type": "string",
|
||||
"pattern": r"^\d{1,4} \d{1,4} \d{1,4} \d{1,4}$",
|
||||
},
|
||||
},
|
||||
"required": ["label", "bbox"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
PROMPT_TYPE_LAYOUT = "layout"
|
||||
PROMPT_TYPE_BLOCK = "block"
|
||||
PROMPT_TYPE_TABLE_REC = "table_rec"
|
||||
PROMPT_TYPE_HIGH_ACCURACY_BBOX = "high_accuracy_bbox"
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchInputItem:
|
||||
image: Image.Image
|
||||
prompt_type: str
|
||||
prompt: Optional[str] = None # If set, overrides the default prompt for prompt_type
|
||||
max_tokens: Optional[int] = None
|
||||
request_logprobs: bool = False
|
||||
# vllm-native guided decoding — JSON schema, regex, or grammar string.
|
||||
# When set, the server constrains the decode tokens to match the schema.
|
||||
guided_json: Optional[dict] = None
|
||||
guided_regex: Optional[str] = None
|
||||
metadata: dict = field(default_factory=dict) # Free-form, passes through to output
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenerationResult:
|
||||
raw: str
|
||||
token_count: int
|
||||
error: bool = False
|
||||
# Mean of exp(logprob) across response tokens, if logprobs requested
|
||||
mean_token_prob: Optional[float] = None
|
||||
# Per-token logprobs (raw OpenAI-style content list), if requested - phase 2 use
|
||||
logprobs: Optional[List[Any]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchOutputItem:
|
||||
raw: str
|
||||
token_count: int
|
||||
error: bool
|
||||
mean_token_prob: Optional[float] = None
|
||||
logprobs: Optional[List[Any]] = None
|
||||
metadata: dict = field(default_factory=dict)
|
||||
@@ -0,0 +1,96 @@
|
||||
from typing import Tuple
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def scale_to_fit(
|
||||
img: Image.Image,
|
||||
max_size: Tuple[int, int] = (3072, 2048),
|
||||
min_size: Tuple[int, int] = (1792, 28),
|
||||
grid_size: int = 28,
|
||||
) -> Image.Image:
|
||||
resample_method = Image.Resampling.LANCZOS
|
||||
|
||||
width, height = img.size
|
||||
|
||||
if width <= 0 or height <= 0:
|
||||
return img
|
||||
|
||||
original_ar = width / height
|
||||
current_pixels = width * height
|
||||
max_pixels = max_size[0] * max_size[1]
|
||||
min_pixels = min_size[0] * min_size[1]
|
||||
|
||||
scale = 1.0
|
||||
if current_pixels > max_pixels:
|
||||
scale = (max_pixels / current_pixels) ** 0.5
|
||||
elif current_pixels < min_pixels:
|
||||
scale = (min_pixels / current_pixels) ** 0.5
|
||||
|
||||
w_blocks = max(1, round((width * scale) / grid_size))
|
||||
h_blocks = max(1, round((height * scale) / grid_size))
|
||||
|
||||
while (w_blocks * h_blocks * grid_size * grid_size) > max_pixels:
|
||||
if w_blocks == 1 and h_blocks == 1:
|
||||
break
|
||||
|
||||
if w_blocks == 1:
|
||||
h_blocks -= 1
|
||||
continue
|
||||
if h_blocks == 1:
|
||||
w_blocks -= 1
|
||||
continue
|
||||
|
||||
ar_w_loss = abs(((w_blocks - 1) / h_blocks) - original_ar)
|
||||
ar_h_loss = abs((w_blocks / (h_blocks - 1)) - original_ar)
|
||||
|
||||
if ar_w_loss < ar_h_loss:
|
||||
w_blocks -= 1
|
||||
else:
|
||||
h_blocks -= 1
|
||||
|
||||
new_width = w_blocks * grid_size
|
||||
new_height = h_blocks * grid_size
|
||||
|
||||
if (new_width, new_height) == (width, height):
|
||||
return img
|
||||
|
||||
return img.resize((new_width, new_height), resample=resample_method)
|
||||
|
||||
|
||||
def detect_repeat_token(
|
||||
predicted_tokens: str,
|
||||
base_max_repeats: int = 4,
|
||||
window_size: int = 500,
|
||||
cut_from_end: int = 0,
|
||||
scaling_factor: float = 3.0,
|
||||
) -> bool:
|
||||
if cut_from_end > 0:
|
||||
predicted_tokens = predicted_tokens[:-cut_from_end]
|
||||
|
||||
for seq_len in range(1, window_size // 2 + 1):
|
||||
candidate_seq = predicted_tokens[-seq_len:]
|
||||
|
||||
max_repeats = int(base_max_repeats * (1 + scaling_factor / seq_len))
|
||||
|
||||
repeat_count = 0
|
||||
pos = len(predicted_tokens) - seq_len
|
||||
if pos < 0:
|
||||
continue
|
||||
|
||||
while pos >= 0:
|
||||
if predicted_tokens[pos : pos + seq_len] == candidate_seq:
|
||||
repeat_count += 1
|
||||
pos -= seq_len
|
||||
else:
|
||||
break
|
||||
|
||||
if repeat_count > max_repeats:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def image_token_budget(block_count: int, ceiling: int = 4096, floor: int = 64) -> int:
|
||||
"""Per-block max_tokens: count + 100, clamped to [floor, ceiling]."""
|
||||
return min(max(block_count + 100, floor), ceiling)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user