commit 1a585693bea0a0757640507742058c2708cd3705 Author: Fu Dai Date: Wed Jun 17 10:20:02 2026 +0400 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) diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..7458b49 --- /dev/null +++ b/.dockerignore @@ -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 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d6a65a0 --- /dev/null +++ b/.gitignore @@ -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 diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..1ef49ed --- /dev/null +++ b/AGENTS.md @@ -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. diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..ba68135 --- /dev/null +++ b/Dockerfile @@ -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"] diff --git a/IMPLEMENTATION_PLAN_AND_TEST_RESULTS.md b/IMPLEMENTATION_PLAN_AND_TEST_RESULTS.md new file mode 100644 index 0000000..c77f3c0 --- /dev/null +++ b/IMPLEMENTATION_PLAN_AND_TEST_RESULTS.md @@ -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. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..792e81c --- /dev/null +++ b/LICENSE @@ -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. diff --git a/README.md b/README.md new file mode 100644 index 0000000..d34a1d4 --- /dev/null +++ b/README.md @@ -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**
叻報, 1890 | | | 大清光緒十六年
本館新聞除禮拜外日出一張
諸君賜閱本報者無論本埠外埠…
工務局告示 | +| **English**
The Nation, 1846 | | | VOL. IV. No. 181.
DUBLIN, SATURDAY, MARCH 28, 1846.
PRICE 6 D.
DAVIS TESTIMONIAL. | +| **Arabic**
Al-Ahram, 1981 | | | رئيس مجلس الادارة — عبد الله عبد البخاري
المنطقات المنتصرة
السنة ١٠٠٧ — العدد ٣٩٦٣٨ | +| **Russian**
Виттова Пляска, 1905 | | | Виттова Пляска
ОДНОДНЕВНАЯ ГАЗЕТА. ПАЛИ—ТИКО—ФИ—НАНСОВАЯ
Цѣна 5 копѣекъ. | +| **French**
Le Miroir des Sports, 1937 | | | LE MIROIR DES SPORTS
Le plus fort tirage des hebdomadaires sportifs
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://:5002/v1/api/ai/swagger`. + +All OCR endpoints accept the same JSON body and return the same envelope. + +**Request body:** +```json +{ + "file": "", + "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://: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 | + +![Latency vs concurrency](benchmarks/concurrency_sweep_latency.png) + +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//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: +- Surya GitHub: +- Model: diff --git a/SINGLE_CALL_TIMING_RESULT.md b/SINGLE_CALL_TIMING_RESULT.md new file mode 100644 index 0000000..8895209 --- /dev/null +++ b/SINGLE_CALL_TIMING_RESULT.md @@ -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. diff --git a/api.py b/api.py new file mode 100644 index 0000000..98078bf --- /dev/null +++ b/api.py @@ -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) diff --git a/api_test.py b/api_test.py new file mode 100644 index 0000000..2752761 --- /dev/null +++ b/api_test.py @@ -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()) diff --git a/assets/arabic.txt b/assets/arabic.txt new file mode 100644 index 0000000..a103906 --- /dev/null +++ b/assets/arabic.txt @@ -0,0 +1,15 @@ +A colorful illustration of a structural steel brick facade with a complex graphic design of a steel brick facade an +تسس الأطراف سنة ١٨٧٥ : سانيم ويشمارة نقلا +رئيس مجلس الادارة   رئيس الشعب وير   عبد الله عبد البخاري   إبراهيم تلفيح +<img alt= 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() diff --git a/benchmarks/concurrency_sweep.py b/benchmarks/concurrency_sweep.py new file mode 100644 index 0000000..196e44d --- /dev/null +++ b/benchmarks/concurrency_sweep.py @@ -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() diff --git a/benchmarks/concurrency_sweep_latency.png b/benchmarks/concurrency_sweep_latency.png new file mode 100644 index 0000000..d0cc867 Binary files /dev/null and b/benchmarks/concurrency_sweep_latency.png differ diff --git a/benchmarks/concurrency_sweep_results.csv b/benchmarks/concurrency_sweep_results.csv new file mode 100644 index 0000000..3cb470a --- /dev/null +++ b/benchmarks/concurrency_sweep_results.csv @@ -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 diff --git a/benchmarks/concurrency_sweep_results.json b/benchmarks/concurrency_sweep_results.json new file mode 100644 index 0000000..37dfe91 --- /dev/null +++ b/benchmarks/concurrency_sweep_results.json @@ -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": [] + } + ] +} \ No newline at end of file diff --git a/benchmarks/concurrency_test.py b/benchmarks/concurrency_test.py new file mode 100644 index 0000000..7ff1e03 --- /dev/null +++ b/benchmarks/concurrency_test.py @@ -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() diff --git a/benchmarks/concurrency_test_results.json b/benchmarks/concurrency_test_results.json new file mode 100644 index 0000000..842cd91 --- /dev/null +++ b/benchmarks/concurrency_test_results.json @@ -0,0 +1,73 @@ +{ + "started_at": "2026-06-12T10:43:50+0400", + "image": "temp_image_c3c56948-282e-453b-8fac-6c482243d1e5.jpg", + "requests": 16, + "concurrency": 16, + "preflight": { + "docker_runtimes": { + "ok": true, + "stdout": "{\"io.containerd.runc.v2\":{\"path\":\"runc\",\"status\":{\"org.opencontainers.runtime-spec.features\":\"{\\\"ociVersionMin\\\":\\\"1.0.0\\\",\\\"ociVersionMax\\\":\\\"1.2.0\\\",\\\"hooks\\\":[\\\"prestart\\\",\\\"createRuntime\\\",\\\"createContainer\\\",\\\"startContainer\\\",\\\"poststart\\\",\\\"poststop\\\"],\\\"mountOptions\\\":[\\\"async\\\",\\\"atime\\\",\\\"bind\\\",\\\"defaults\\\",\\\"dev\\\",\\\"diratime\\\",\\\"dirsync\\\",\\\"exec\\\",\\\"iversion\\\",\\\"lazytime\\\",\\\"loud\\\",\\\"mand\\\",\\\"noatime\\\",\\\"nodev\\\",\\\"nodiratime\\\",\\\"noexec\\\",\\\"noiversion\\\",\\\"nolazytime\\\",\\\"nomand\\\",\\\"norelatime\\\",\\\"nostrictatime\\\",\\\"nosuid\\\",\\\"nosymfollow\\\",\\\"private\\\",\\\"ratime\\\",\\\"rbind\\\",\\\"rdev\\\",\\\"rdiratime\\\",\\\"relatime\\\",\\\"remount\\\",\\\"rexec\\\",\\\"rnoatime\\\",\\\"rnodev\\\",\\\"rnodiratime\\\",\\\"rnoexec\\\",\\\"rnorelatime\\\",\\\"rnostrictatime\\\",\\\"rnosuid\\\",\\\"rnosymfollow\\\",\\\"ro\\\",\\\"rprivate\\\",\\\"rrelatime\\\",\\\"rro\\\",\\\"rrw\\\",\\\"rshared\\\",\\\"rslave\\\",\\\"rstrictatime\\\",\\\"rsuid\\\",\\\"rsymfollow\\\",\\\"runbindable\\\",\\\"rw\\\",\\\"shared\\\",\\\"silent\\\",\\\"slave\\\",\\\"strictatime\\\",\\\"suid\\\",\\\"symfollow\\\",\\\"sync\\\",\\\"tmpcopyup\\\",\\\"unbindable\\\"],\\\"linux\\\":{\\\"namespaces\\\":[\\\"cgroup\\\",\\\"ipc\\\",\\\"mount\\\",\\\"network\\\",\\\"pid\\\",\\\"time\\\",\\\"user\\\",\\\"uts\\\"],\\\"capabilities\\\":[\\\"CAP_CHOWN\\\",\\\"CAP_DAC_OVERRIDE\\\",\\\"CAP_DAC_READ_SEARCH\\\",\\\"CAP_FOWNER\\\",\\\"CAP_FSETID\\\",\\\"CAP_KILL\\\",\\\"CAP_SETGID\\\",\\\"CAP_SETUID\\\",\\\"CAP_SETPCAP\\\",\\\"CAP_LINUX_IMMUTABLE\\\",\\\"CAP_NET_BIND_SERVICE\\\",\\\"CAP_NET_BROADCAST\\\",\\\"CAP_NET_ADMIN\\\",\\\"CAP_NET_RAW\\\",\\\"CAP_IPC_LOCK\\\",\\\"CAP_IPC_OWNER\\\",\\\"CAP_SYS_MODULE\\\",\\\"CAP_SYS_RAWIO\\\",\\\"CAP_SYS_CHROOT\\\",\\\"CAP_SYS_PTRACE\\\",\\\"CAP_SYS_PACCT\\\",\\\"CAP_SYS_ADMIN\\\",\\\"CAP_SYS_BOOT\\\",\\\"CAP_SYS_NICE\\\",\\\"CAP_SYS_RESOURCE\\\",\\\"CAP_SYS_TIME\\\",\\\"CAP_SYS_TTY_CONFIG\\\",\\\"CAP_MKNOD\\\",\\\"CAP_LEASE\\\",\\\"CAP_AUDIT_WRITE\\\",\\\"CAP_AUDIT_CONTROL\\\",\\\"CAP_SETFCAP\\\",\\\"CAP_MAC_OVERRIDE\\\",\\\"CAP_MAC_ADMIN\\\",\\\"CAP_SYSLOG\\\",\\\"CAP_WAKE_ALARM\\\",\\\"CAP_BLOCK_SUSPEND\\\",\\\"CAP_AUDIT_READ\\\",\\\"CAP_PERFMON\\\",\\\"CAP_BPF\\\",\\\"CAP_CHECKPOINT_RESTORE\\\"],\\\"cgroup\\\":{\\\"v1\\\":true,\\\"v2\\\":true,\\\"systemd\\\":true,\\\"systemdUser\\\":true,\\\"rdma\\\":true},\\\"seccomp\\\":{\\\"enabled\\\":true,\\\"actions\\\":[\\\"SCMP_ACT_ALLOW\\\",\\\"SCMP_ACT_ERRNO\\\",\\\"SCMP_ACT_KILL\\\",\\\"SCMP_ACT_KILL_PROCESS\\\",\\\"SCMP_ACT_KILL_THREAD\\\",\\\"SCMP_ACT_LOG\\\",\\\"SCMP_ACT_NOTIFY\\\",\\\"SCMP_ACT_TRACE\\\",\\\"SCMP_ACT_TRAP\\\"],\\\"operators\\\":[\\\"SCMP_CMP_EQ\\\",\\\"SCMP_CMP_GE\\\",\\\"SCMP_CMP_GT\\\",\\\"SCMP_CMP_LE\\\",\\\"SCMP_CMP_LT\\\",\\\"SCMP_CMP_MASKED_EQ\\\",\\\"SCMP_CMP_NE\\\"],\\\"archs\\\":[\\\"SCMP_ARCH_AARCH64\\\",\\\"SCMP_ARCH_ARM\\\",\\\"SCMP_ARCH_MIPS\\\",\\\"SCMP_ARCH_MIPS64\\\",\\\"SCMP_ARCH_MIPS64N32\\\",\\\"SCMP_ARCH_MIPSEL\\\",\\\"SCMP_ARCH_MIPSEL64\\\",\\\"SCMP_ARCH_MIPSEL64N32\\\",\\\"SCMP_ARCH_PPC\\\",\\\"SCMP_ARCH_PPC64\\\",\\\"SCMP_ARCH_PPC64LE\\\",\\\"SCMP_ARCH_RISCV64\\\",\\\"SCMP_ARCH_S390\\\",\\\"SCMP_ARCH_S390X\\\",\\\"SCMP_ARCH_X32\\\",\\\"SCMP_ARCH_X86\\\",\\\"SCMP_ARCH_X86_64\\\"],\\\"knownFlags\\\":[\\\"SECCOMP_FILTER_FLAG_TSYNC\\\",\\\"SECCOMP_FILTER_FLAG_SPEC_ALLOW\\\",\\\"SECCOMP_FILTER_FLAG_LOG\\\"],\\\"supportedFlags\\\":[\\\"SECCOMP_FILTER_FLAG_TSYNC\\\",\\\"SECCOMP_FILTER_FLAG_SPEC_ALLOW\\\",\\\"SECCOMP_FILTER_FLAG_LOG\\\"]},\\\"apparmor\\\":{\\\"enabled\\\":true},\\\"selinux\\\":{\\\"enabled\\\":true},\\\"intelRdt\\\":{\\\"enabled\\\":true},\\\"mountExtensions\\\":{\\\"idmap\\\":{\\\"enabled\\\":true}}},\\\"annotations\\\":{\\\"io.github.seccomp.libseccomp.version\\\":\\\"2.5.3\\\",\\\"org.opencontainers.runc.checkpoint.enabled\\\":\\\"true\\\",\\\"org.opencontainers.runc.commit\\\":\\\"v1.2.5-0-g59923ef\\\",\\\"org.opencontainers.runc.version\\\":\\\"1.2.5\\\"},\\\"potentiallyUnsafeConfigAnnotations\\\":[\\\"bundle\\\",\\\"org.systemd.property.\\\",\\\"org.criu.config\\\"]}\"}},\"nvidia\":{\"path\":\"nvidia-container-runtime\",\"status\":{\"org.opencontainers.runtime-spec.features\":\"{\\\"ociVersionMin\\\":\\\"1.0.0\\\",\\\"ociVersionMax\\\":\\\"1.2.0\\\",\\\"hooks\\\":[\\\"prestart\\\",\\\"createRuntime\\\",\\\"createContainer\\\",\\\"startContainer\\\",\\\"poststart\\\",\\\"poststop\\\"],\\\"mountOptions\\\":[\\\"async\\\",\\\"atime\\\",\\\"bind\\\",\\\"defaults\\\",\\\"dev\\\",\\\"diratime\\\",\\\"dirsync\\\",\\\"exec\\\",\\\"iversion\\\",\\\"lazytime\\\",\\\"loud\\\",\\\"mand\\\",\\\"noatime\\\",\\\"nodev\\\",\\\"nodiratime\\\",\\\"noexec\\\",\\\"noiversion\\\",\\\"nolazytime\\\",\\\"nomand\\\",\\\"norelatime\\\",\\\"nostrictatime\\\",\\\"nosuid\\\",\\\"nosymfollow\\\",\\\"private\\\",\\\"ratime\\\",\\\"rbind\\\",\\\"rdev\\\",\\\"rdiratime\\\",\\\"relatime\\\",\\\"remount\\\",\\\"rexec\\\",\\\"rnoatime\\\",\\\"rnodev\\\",\\\"rnodiratime\\\",\\\"rnoexec\\\",\\\"rnorelatime\\\",\\\"rnostrictatime\\\",\\\"rnosuid\\\",\\\"rnosymfollow\\\",\\\"ro\\\",\\\"rprivate\\\",\\\"rrelatime\\\",\\\"rro\\\",\\\"rrw\\\",\\\"rshared\\\",\\\"rslave\\\",\\\"rstrictatime\\\",\\\"rsuid\\\",\\\"rsymfollow\\\",\\\"runbindable\\\",\\\"rw\\\",\\\"shared\\\",\\\"silent\\\",\\\"slave\\\",\\\"strictatime\\\",\\\"suid\\\",\\\"symfollow\\\",\\\"sync\\\",\\\"tmpcopyup\\\",\\\"unbindable\\\"],\\\"linux\\\":{\\\"namespaces\\\":[\\\"cgroup\\\",\\\"ipc\\\",\\\"mount\\\",\\\"network\\\",\\\"pid\\\",\\\"time\\\",\\\"user\\\",\\\"uts\\\"],\\\"capabilities\\\":[\\\"CAP_CHOWN\\\",\\\"CAP_DAC_OVERRIDE\\\",\\\"CAP_DAC_READ_SEARCH\\\",\\\"CAP_FOWNER\\\",\\\"CAP_FSETID\\\",\\\"CAP_KILL\\\",\\\"CAP_SETGID\\\",\\\"CAP_SETUID\\\",\\\"CAP_SETPCAP\\\",\\\"CAP_LINUX_IMMUTABLE\\\",\\\"CAP_NET_BIND_SERVICE\\\",\\\"CAP_NET_BROADCAST\\\",\\\"CAP_NET_ADMIN\\\",\\\"CAP_NET_RAW\\\",\\\"CAP_IPC_LOCK\\\",\\\"CAP_IPC_OWNER\\\",\\\"CAP_SYS_MODULE\\\",\\\"CAP_SYS_RAWIO\\\",\\\"CAP_SYS_CHROOT\\\",\\\"CAP_SYS_PTRACE\\\",\\\"CAP_SYS_PACCT\\\",\\\"CAP_SYS_ADMIN\\\",\\\"CAP_SYS_BOOT\\\",\\\"CAP_SYS_NICE\\\",\\\"CAP_SYS_RESOURCE\\\",\\\"CAP_SYS_TIME\\\",\\\"CAP_SYS_TTY_CONFIG\\\",\\\"CAP_MKNOD\\\",\\\"CAP_LEASE\\\",\\\"CAP_AUDIT_WRITE\\\",\\\"CAP_AUDIT_CONTROL\\\",\\\"CAP_SETFCAP\\\",\\\"CAP_MAC_OVERRIDE\\\",\\\"CAP_MAC_ADMIN\\\",\\\"CAP_SYSLOG\\\",\\\"CAP_WAKE_ALARM\\\",\\\"CAP_BLOCK_SUSPEND\\\",\\\"CAP_AUDIT_READ\\\",\\\"CAP_PERFMON\\\",\\\"CAP_BPF\\\",\\\"CAP_CHECKPOINT_RESTORE\\\"],\\\"cgroup\\\":{\\\"v1\\\":true,\\\"v2\\\":true,\\\"systemd\\\":true,\\\"systemdUser\\\":true,\\\"rdma\\\":true},\\\"seccomp\\\":{\\\"enabled\\\":true,\\\"actions\\\":[\\\"SCMP_ACT_ALLOW\\\",\\\"SCMP_ACT_ERRNO\\\",\\\"SCMP_ACT_KILL\\\",\\\"SCMP_ACT_KILL_PROCESS\\\",\\\"SCMP_ACT_KILL_THREAD\\\",\\\"SCMP_ACT_LOG\\\",\\\"SCMP_ACT_NOTIFY\\\",\\\"SCMP_ACT_TRACE\\\",\\\"SCMP_ACT_TRAP\\\"],\\\"operators\\\":[\\\"SCMP_CMP_EQ\\\",\\\"SCMP_CMP_GE\\\",\\\"SCMP_CMP_GT\\\",\\\"SCMP_CMP_LE\\\",\\\"SCMP_CMP_LT\\\",\\\"SCMP_CMP_MASKED_EQ\\\",\\\"SCMP_CMP_NE\\\"],\\\"archs\\\":[\\\"SCMP_ARCH_AARCH64\\\",\\\"SCMP_ARCH_ARM\\\",\\\"SCMP_ARCH_MIPS\\\",\\\"SCMP_ARCH_MIPS64\\\",\\\"SCMP_ARCH_MIPS64N32\\\",\\\"SCMP_ARCH_MIPSEL\\\",\\\"SCMP_ARCH_MIPSEL64\\\",\\\"SCMP_ARCH_MIPSEL64N32\\\",\\\"SCMP_ARCH_PPC\\\",\\\"SCMP_ARCH_PPC64\\\",\\\"SCMP_ARCH_PPC64LE\\\",\\\"SCMP_ARCH_RISCV64\\\",\\\"SCMP_ARCH_S390\\\",\\\"SCMP_ARCH_S390X\\\",\\\"SCMP_ARCH_X32\\\",\\\"SCMP_ARCH_X86\\\",\\\"SCMP_ARCH_X86_64\\\"],\\\"knownFlags\\\":[\\\"SECCOMP_FILTER_FLAG_TSYNC\\\",\\\"SECCOMP_FILTER_FLAG_SPEC_ALLOW\\\",\\\"SECCOMP_FILTER_FLAG_LOG\\\"],\\\"supportedFlags\\\":[\\\"SECCOMP_FILTER_FLAG_TSYNC\\\",\\\"SECCOMP_FILTER_FLAG_SPEC_ALLOW\\\",\\\"SECCOMP_FILTER_FLAG_LOG\\\"]},\\\"apparmor\\\":{\\\"enabled\\\":true},\\\"selinux\\\":{\\\"enabled\\\":true},\\\"intelRdt\\\":{\\\"enabled\\\":true},\\\"mountExtensions\\\":{\\\"idmap\\\":{\\\"enabled\\\":true}}},\\\"annotations\\\":{\\\"io.github.seccomp.libseccomp.version\\\":\\\"2.5.3\\\",\\\"org.opencontainers.runc.checkpoint.enabled\\\":\\\"true\\\",\\\"org.opencontainers.runc.commit\\\":\\\"v1.2.5-0-g59923ef\\\",\\\"org.opencontainers.runc.version\\\":\\\"1.2.5\\\"},\\\"potentiallyUnsafeConfigAnnotations\\\":[\\\"bundle\\\",\\\"org.systemd.property.\\\",\\\"org.criu.config\\\"]}\"}},\"runc\":{\"path\":\"runc\",\"status\":{\"org.opencontainers.runtime-spec.features\":\"{\\\"ociVersionMin\\\":\\\"1.0.0\\\",\\\"ociVersionMax\\\":\\\"1.2.0\\\",\\\"hooks\\\":[\\\"prestart\\\",\\\"createRuntime\\\",\\\"createContainer\\\",\\\"startContainer\\\",\\\"poststart\\\",\\\"poststop\\\"],\\\"mountOptions\\\":[\\\"async\\\",\\\"atime\\\",\\\"bind\\\",\\\"defaults\\\",\\\"dev\\\",\\\"diratime\\\",\\\"dirsync\\\",\\\"exec\\\",\\\"iversion\\\",\\\"lazytime\\\",\\\"loud\\\",\\\"mand\\\",\\\"noatime\\\",\\\"nodev\\\",\\\"nodiratime\\\",\\\"noexec\\\",\\\"noiversion\\\",\\\"nolazytime\\\",\\\"nomand\\\",\\\"norelatime\\\",\\\"nostrictatime\\\",\\\"nosuid\\\",\\\"nosymfollow\\\",\\\"private\\\",\\\"ratime\\\",\\\"rbind\\\",\\\"rdev\\\",\\\"rdiratime\\\",\\\"relatime\\\",\\\"remount\\\",\\\"rexec\\\",\\\"rnoatime\\\",\\\"rnodev\\\",\\\"rnodiratime\\\",\\\"rnoexec\\\",\\\"rnorelatime\\\",\\\"rnostrictatime\\\",\\\"rnosuid\\\",\\\"rnosymfollow\\\",\\\"ro\\\",\\\"rprivate\\\",\\\"rrelatime\\\",\\\"rro\\\",\\\"rrw\\\",\\\"rshared\\\",\\\"rslave\\\",\\\"rstrictatime\\\",\\\"rsuid\\\",\\\"rsymfollow\\\",\\\"runbindable\\\",\\\"rw\\\",\\\"shared\\\",\\\"silent\\\",\\\"slave\\\",\\\"strictatime\\\",\\\"suid\\\",\\\"symfollow\\\",\\\"sync\\\",\\\"tmpcopyup\\\",\\\"unbindable\\\"],\\\"linux\\\":{\\\"namespaces\\\":[\\\"cgroup\\\",\\\"ipc\\\",\\\"mount\\\",\\\"network\\\",\\\"pid\\\",\\\"time\\\",\\\"user\\\",\\\"uts\\\"],\\\"capabilities\\\":[\\\"CAP_CHOWN\\\",\\\"CAP_DAC_OVERRIDE\\\",\\\"CAP_DAC_READ_SEARCH\\\",\\\"CAP_FOWNER\\\",\\\"CAP_FSETID\\\",\\\"CAP_KILL\\\",\\\"CAP_SETGID\\\",\\\"CAP_SETUID\\\",\\\"CAP_SETPCAP\\\",\\\"CAP_LINUX_IMMUTABLE\\\",\\\"CAP_NET_BIND_SERVICE\\\",\\\"CAP_NET_BROADCAST\\\",\\\"CAP_NET_ADMIN\\\",\\\"CAP_NET_RAW\\\",\\\"CAP_IPC_LOCK\\\",\\\"CAP_IPC_OWNER\\\",\\\"CAP_SYS_MODULE\\\",\\\"CAP_SYS_RAWIO\\\",\\\"CAP_SYS_CHROOT\\\",\\\"CAP_SYS_PTRACE\\\",\\\"CAP_SYS_PACCT\\\",\\\"CAP_SYS_ADMIN\\\",\\\"CAP_SYS_BOOT\\\",\\\"CAP_SYS_NICE\\\",\\\"CAP_SYS_RESOURCE\\\",\\\"CAP_SYS_TIME\\\",\\\"CAP_SYS_TTY_CONFIG\\\",\\\"CAP_MKNOD\\\",\\\"CAP_LEASE\\\",\\\"CAP_AUDIT_WRITE\\\",\\\"CAP_AUDIT_CONTROL\\\",\\\"CAP_SETFCAP\\\",\\\"CAP_MAC_OVERRIDE\\\",\\\"CAP_MAC_ADMIN\\\",\\\"CAP_SYSLOG\\\",\\\"CAP_WAKE_ALARM\\\",\\\"CAP_BLOCK_SUSPEND\\\",\\\"CAP_AUDIT_READ\\\",\\\"CAP_PERFMON\\\",\\\"CAP_BPF\\\",\\\"CAP_CHECKPOINT_RESTORE\\\"],\\\"cgroup\\\":{\\\"v1\\\":true,\\\"v2\\\":true,\\\"systemd\\\":true,\\\"systemdUser\\\":true,\\\"rdma\\\":true},\\\"seccomp\\\":{\\\"enabled\\\":true,\\\"actions\\\":[\\\"SCMP_ACT_ALLOW\\\",\\\"SCMP_ACT_ERRNO\\\",\\\"SCMP_ACT_KILL\\\",\\\"SCMP_ACT_KILL_PROCESS\\\",\\\"SCMP_ACT_KILL_THREAD\\\",\\\"SCMP_ACT_LOG\\\",\\\"SCMP_ACT_NOTIFY\\\",\\\"SCMP_ACT_TRACE\\\",\\\"SCMP_ACT_TRAP\\\"],\\\"operators\\\":[\\\"SCMP_CMP_EQ\\\",\\\"SCMP_CMP_GE\\\",\\\"SCMP_CMP_GT\\\",\\\"SCMP_CMP_LE\\\",\\\"SCMP_CMP_LT\\\",\\\"SCMP_CMP_MASKED_EQ\\\",\\\"SCMP_CMP_NE\\\"],\\\"archs\\\":[\\\"SCMP_ARCH_AARCH64\\\",\\\"SCMP_ARCH_ARM\\\",\\\"SCMP_ARCH_MIPS\\\",\\\"SCMP_ARCH_MIPS64\\\",\\\"SCMP_ARCH_MIPS64N32\\\",\\\"SCMP_ARCH_MIPSEL\\\",\\\"SCMP_ARCH_MIPSEL64\\\",\\\"SCMP_ARCH_MIPSEL64N32\\\",\\\"SCMP_ARCH_PPC\\\",\\\"SCMP_ARCH_PPC64\\\",\\\"SCMP_ARCH_PPC64LE\\\",\\\"SCMP_ARCH_RISCV64\\\",\\\"SCMP_ARCH_S390\\\",\\\"SCMP_ARCH_S390X\\\",\\\"SCMP_ARCH_X32\\\",\\\"SCMP_ARCH_X86\\\",\\\"SCMP_ARCH_X86_64\\\"],\\\"knownFlags\\\":[\\\"SECCOMP_FILTER_FLAG_TSYNC\\\",\\\"SECCOMP_FILTER_FLAG_SPEC_ALLOW\\\",\\\"SECCOMP_FILTER_FLAG_LOG\\\"],\\\"supportedFlags\\\":[\\\"SECCOMP_FILTER_FLAG_TSYNC\\\",\\\"SECCOMP_FILTER_FLAG_SPEC_ALLOW\\\",\\\"SECCOMP_FILTER_FLAG_LOG\\\"]},\\\"apparmor\\\":{\\\"enabled\\\":true},\\\"selinux\\\":{\\\"enabled\\\":true},\\\"intelRdt\\\":{\\\"enabled\\\":true},\\\"mountExtensions\\\":{\\\"idmap\\\":{\\\"enabled\\\":true}}},\\\"annotations\\\":{\\\"io.github.seccomp.libseccomp.version\\\":\\\"2.5.3\\\",\\\"org.opencontainers.runc.checkpoint.enabled\\\":\\\"true\\\",\\\"org.opencontainers.runc.commit\\\":\\\"v1.2.5-0-g59923ef\\\",\\\"org.opencontainers.runc.version\\\":\\\"1.2.5\\\"},\\\"potentiallyUnsafeConfigAnnotations\\\":[\\\"bundle\\\",\\\"org.systemd.property.\\\",\\\"org.criu.config\\\"]}\"}}}", + "stderr": "" + }, + "nvidia_smi": { + "ok": true, + "stdout": "GPU 0: NVIDIA A100 80GB PCIe (UUID: GPU-911fab14-997f-e653-67c4-abba6d597ec6)\nGPU 1: NVIDIA A100 80GB PCIe (UUID: GPU-b6584aae-31c5-e793-7d22-98fbebef889f)\nGPU 2: NVIDIA A100 80GB PCIe (UUID: GPU-b72f3ad9-d411-a9da-c4a2-be9f0cb175fe)\nGPU 3: NVIDIA A100 80GB PCIe (UUID: GPU-e1d3ca84-f769-2e6d-4179-d3ec0ebd94f2)", + "stderr": "" + }, + "nvidia_container_runtime": { + "ok": true, + "stdout": "/usr/bin/nvidia-container-runtime", + "stderr": "" + }, + "nvidia_ctk": { + "ok": true, + "stdout": "/usr/bin/nvidia-ctk", + "stderr": "" + }, + "docker_gpus": { + "ok": true, + "stdout": "GPU 0: NVIDIA A100 80GB PCIe (UUID: GPU-911fab14-997f-e653-67c4-abba6d597ec6)\nGPU 1: NVIDIA A100 80GB PCIe (UUID: GPU-b6584aae-31c5-e793-7d22-98fbebef889f)\nGPU 2: NVIDIA A100 80GB PCIe (UUID: GPU-b72f3ad9-d411-a9da-c4a2-be9f0cb175fe)\nGPU 3: NVIDIA A100 80GB PCIe (UUID: GPU-e1d3ca84-f769-2e6d-4179-d3ec0ebd94f2)", + "stderr": "" + }, + "docker_has_nvidia_runtime": true, + "recommended_fix": "Run: sudo nvidia-ctk runtime configure --runtime=docker --config=/etc/docker/daemon.json && sudo systemctl restart docker" + }, + "endpoints": [ + { + "name": "legacy", + "url": "http://127.0.0.1:5002/v1/api/ai/suya_ocr/", + "requests": 16, + "concurrency": 16, + "success": 16, + "failed": 0, + "wall_seconds": 41.21962874662131, + "throughput_rps": 0.3881645829066689, + "latency_seconds": { + "min": 19.009264937601984, + "mean": 38.5007132650353, + "p50": 39.90640107169747, + "p95": 41.13191143143922, + "max": 41.217520573176444 + }, + "errors": [] + }, + { + "name": "vllm", + "url": "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/", + "requests": 16, + "concurrency": 16, + "success": 16, + "failed": 0, + "wall_seconds": 49.16062220837921, + "throughput_rps": 0.3254637407187428, + "latency_seconds": { + "min": 5.31236336287111, + "mean": 36.04962155013345, + "p50": 28.501926544122398, + "p95": 49.097963376902044, + "max": 49.10191335435957 + }, + "errors": [] + } + ] +} \ No newline at end of file diff --git a/benchmarks/concurrency_warmup_result.json b/benchmarks/concurrency_warmup_result.json new file mode 100644 index 0000000..42141c7 --- /dev/null +++ b/benchmarks/concurrency_warmup_result.json @@ -0,0 +1,55 @@ +{ + "started_at": "2026-06-10T12:01:21+0400", + "image": "/path/to/surya/static/images/excerpt.png", + "requests": 1, + "concurrency": 1, + "preflight": { + "docker_runtimes": { + "ok": true, + "stdout": "{\"io.containerd.runc.v2\":{\"path\":\"runc\",\"status\":{\"org.opencontainers.runtime-spec.features\":\"{\\\"ociVersionMin\\\":\\\"1.0.0\\\",\\\"ociVersionMax\\\":\\\"1.2.0\\\",\\\"hooks\\\":[\\\"prestart\\\",\\\"createRuntime\\\",\\\"createContainer\\\",\\\"startContainer\\\",\\\"poststart\\\",\\\"poststop\\\"],\\\"mountOptions\\\":[\\\"async\\\",\\\"atime\\\",\\\"bind\\\",\\\"defaults\\\",\\\"dev\\\",\\\"diratime\\\",\\\"dirsync\\\",\\\"exec\\\",\\\"iversion\\\",\\\"lazytime\\\",\\\"loud\\\",\\\"mand\\\",\\\"noatime\\\",\\\"nodev\\\",\\\"nodiratime\\\",\\\"noexec\\\",\\\"noiversion\\\",\\\"nolazytime\\\",\\\"nomand\\\",\\\"norelatime\\\",\\\"nostrictatime\\\",\\\"nosuid\\\",\\\"nosymfollow\\\",\\\"private\\\",\\\"ratime\\\",\\\"rbind\\\",\\\"rdev\\\",\\\"rdiratime\\\",\\\"relatime\\\",\\\"remount\\\",\\\"rexec\\\",\\\"rnoatime\\\",\\\"rnodev\\\",\\\"rnodiratime\\\",\\\"rnoexec\\\",\\\"rnorelatime\\\",\\\"rnostrictatime\\\",\\\"rnosuid\\\",\\\"rnosymfollow\\\",\\\"ro\\\",\\\"rprivate\\\",\\\"rrelatime\\\",\\\"rro\\\",\\\"rrw\\\",\\\"rshared\\\",\\\"rslave\\\",\\\"rstrictatime\\\",\\\"rsuid\\\",\\\"rsymfollow\\\",\\\"runbindable\\\",\\\"rw\\\",\\\"shared\\\",\\\"silent\\\",\\\"slave\\\",\\\"strictatime\\\",\\\"suid\\\",\\\"symfollow\\\",\\\"sync\\\",\\\"tmpcopyup\\\",\\\"unbindable\\\"],\\\"linux\\\":{\\\"namespaces\\\":[\\\"cgroup\\\",\\\"ipc\\\",\\\"mount\\\",\\\"network\\\",\\\"pid\\\",\\\"time\\\",\\\"user\\\",\\\"uts\\\"],\\\"capabilities\\\":[\\\"CAP_CHOWN\\\",\\\"CAP_DAC_OVERRIDE\\\",\\\"CAP_DAC_READ_SEARCH\\\",\\\"CAP_FOWNER\\\",\\\"CAP_FSETID\\\",\\\"CAP_KILL\\\",\\\"CAP_SETGID\\\",\\\"CAP_SETUID\\\",\\\"CAP_SETPCAP\\\",\\\"CAP_LINUX_IMMUTABLE\\\",\\\"CAP_NET_BIND_SERVICE\\\",\\\"CAP_NET_BROADCAST\\\",\\\"CAP_NET_ADMIN\\\",\\\"CAP_NET_RAW\\\",\\\"CAP_IPC_LOCK\\\",\\\"CAP_IPC_OWNER\\\",\\\"CAP_SYS_MODULE\\\",\\\"CAP_SYS_RAWIO\\\",\\\"CAP_SYS_CHROOT\\\",\\\"CAP_SYS_PTRACE\\\",\\\"CAP_SYS_PACCT\\\",\\\"CAP_SYS_ADMIN\\\",\\\"CAP_SYS_BOOT\\\",\\\"CAP_SYS_NICE\\\",\\\"CAP_SYS_RESOURCE\\\",\\\"CAP_SYS_TIME\\\",\\\"CAP_SYS_TTY_CONFIG\\\",\\\"CAP_MKNOD\\\",\\\"CAP_LEASE\\\",\\\"CAP_AUDIT_WRITE\\\",\\\"CAP_AUDIT_CONTROL\\\",\\\"CAP_SETFCAP\\\",\\\"CAP_MAC_OVERRIDE\\\",\\\"CAP_MAC_ADMIN\\\",\\\"CAP_SYSLOG\\\",\\\"CAP_WAKE_ALARM\\\",\\\"CAP_BLOCK_SUSPEND\\\",\\\"CAP_AUDIT_READ\\\",\\\"CAP_PERFMON\\\",\\\"CAP_BPF\\\",\\\"CAP_CHECKPOINT_RESTORE\\\"],\\\"cgroup\\\":{\\\"v1\\\":true,\\\"v2\\\":true,\\\"systemd\\\":true,\\\"systemdUser\\\":true,\\\"rdma\\\":true},\\\"seccomp\\\":{\\\"enabled\\\":true,\\\"actions\\\":[\\\"SCMP_ACT_ALLOW\\\",\\\"SCMP_ACT_ERRNO\\\",\\\"SCMP_ACT_KILL\\\",\\\"SCMP_ACT_KILL_PROCESS\\\",\\\"SCMP_ACT_KILL_THREAD\\\",\\\"SCMP_ACT_LOG\\\",\\\"SCMP_ACT_NOTIFY\\\",\\\"SCMP_ACT_TRACE\\\",\\\"SCMP_ACT_TRAP\\\"],\\\"operators\\\":[\\\"SCMP_CMP_EQ\\\",\\\"SCMP_CMP_GE\\\",\\\"SCMP_CMP_GT\\\",\\\"SCMP_CMP_LE\\\",\\\"SCMP_CMP_LT\\\",\\\"SCMP_CMP_MASKED_EQ\\\",\\\"SCMP_CMP_NE\\\"],\\\"archs\\\":[\\\"SCMP_ARCH_AARCH64\\\",\\\"SCMP_ARCH_ARM\\\",\\\"SCMP_ARCH_MIPS\\\",\\\"SCMP_ARCH_MIPS64\\\",\\\"SCMP_ARCH_MIPS64N32\\\",\\\"SCMP_ARCH_MIPSEL\\\",\\\"SCMP_ARCH_MIPSEL64\\\",\\\"SCMP_ARCH_MIPSEL64N32\\\",\\\"SCMP_ARCH_PPC\\\",\\\"SCMP_ARCH_PPC64\\\",\\\"SCMP_ARCH_PPC64LE\\\",\\\"SCMP_ARCH_RISCV64\\\",\\\"SCMP_ARCH_S390\\\",\\\"SCMP_ARCH_S390X\\\",\\\"SCMP_ARCH_X32\\\",\\\"SCMP_ARCH_X86\\\",\\\"SCMP_ARCH_X86_64\\\"],\\\"knownFlags\\\":[\\\"SECCOMP_FILTER_FLAG_TSYNC\\\",\\\"SECCOMP_FILTER_FLAG_SPEC_ALLOW\\\",\\\"SECCOMP_FILTER_FLAG_LOG\\\"],\\\"supportedFlags\\\":[\\\"SECCOMP_FILTER_FLAG_TSYNC\\\",\\\"SECCOMP_FILTER_FLAG_SPEC_ALLOW\\\",\\\"SECCOMP_FILTER_FLAG_LOG\\\"]},\\\"apparmor\\\":{\\\"enabled\\\":true},\\\"selinux\\\":{\\\"enabled\\\":true},\\\"intelRdt\\\":{\\\"enabled\\\":true},\\\"mountExtensions\\\":{\\\"idmap\\\":{\\\"enabled\\\":true}}},\\\"annotations\\\":{\\\"io.github.seccomp.libseccomp.version\\\":\\\"2.5.3\\\",\\\"org.opencontainers.runc.checkpoint.enabled\\\":\\\"true\\\",\\\"org.opencontainers.runc.commit\\\":\\\"v1.2.5-0-g59923ef\\\",\\\"org.opencontainers.runc.version\\\":\\\"1.2.5\\\"},\\\"potentiallyUnsafeConfigAnnotations\\\":[\\\"bundle\\\",\\\"org.systemd.property.\\\",\\\"org.criu.config\\\"]}\"}},\"nvidia\":{\"path\":\"nvidia-container-runtime\",\"status\":{\"org.opencontainers.runtime-spec.features\":\"{\\\"ociVersionMin\\\":\\\"1.0.0\\\",\\\"ociVersionMax\\\":\\\"1.2.0\\\",\\\"hooks\\\":[\\\"prestart\\\",\\\"createRuntime\\\",\\\"createContainer\\\",\\\"startContainer\\\",\\\"poststart\\\",\\\"poststop\\\"],\\\"mountOptions\\\":[\\\"async\\\",\\\"atime\\\",\\\"bind\\\",\\\"defaults\\\",\\\"dev\\\",\\\"diratime\\\",\\\"dirsync\\\",\\\"exec\\\",\\\"iversion\\\",\\\"lazytime\\\",\\\"loud\\\",\\\"mand\\\",\\\"noatime\\\",\\\"nodev\\\",\\\"nodiratime\\\",\\\"noexec\\\",\\\"noiversion\\\",\\\"nolazytime\\\",\\\"nomand\\\",\\\"norelatime\\\",\\\"nostrictatime\\\",\\\"nosuid\\\",\\\"nosymfollow\\\",\\\"private\\\",\\\"ratime\\\",\\\"rbind\\\",\\\"rdev\\\",\\\"rdiratime\\\",\\\"relatime\\\",\\\"remount\\\",\\\"rexec\\\",\\\"rnoatime\\\",\\\"rnodev\\\",\\\"rnodiratime\\\",\\\"rnoexec\\\",\\\"rnorelatime\\\",\\\"rnostrictatime\\\",\\\"rnosuid\\\",\\\"rnosymfollow\\\",\\\"ro\\\",\\\"rprivate\\\",\\\"rrelatime\\\",\\\"rro\\\",\\\"rrw\\\",\\\"rshared\\\",\\\"rslave\\\",\\\"rstrictatime\\\",\\\"rsuid\\\",\\\"rsymfollow\\\",\\\"runbindable\\\",\\\"rw\\\",\\\"shared\\\",\\\"silent\\\",\\\"slave\\\",\\\"strictatime\\\",\\\"suid\\\",\\\"symfollow\\\",\\\"sync\\\",\\\"tmpcopyup\\\",\\\"unbindable\\\"],\\\"linux\\\":{\\\"namespaces\\\":[\\\"cgroup\\\",\\\"ipc\\\",\\\"mount\\\",\\\"network\\\",\\\"pid\\\",\\\"time\\\",\\\"user\\\",\\\"uts\\\"],\\\"capabilities\\\":[\\\"CAP_CHOWN\\\",\\\"CAP_DAC_OVERRIDE\\\",\\\"CAP_DAC_READ_SEARCH\\\",\\\"CAP_FOWNER\\\",\\\"CAP_FSETID\\\",\\\"CAP_KILL\\\",\\\"CAP_SETGID\\\",\\\"CAP_SETUID\\\",\\\"CAP_SETPCAP\\\",\\\"CAP_LINUX_IMMUTABLE\\\",\\\"CAP_NET_BIND_SERVICE\\\",\\\"CAP_NET_BROADCAST\\\",\\\"CAP_NET_ADMIN\\\",\\\"CAP_NET_RAW\\\",\\\"CAP_IPC_LOCK\\\",\\\"CAP_IPC_OWNER\\\",\\\"CAP_SYS_MODULE\\\",\\\"CAP_SYS_RAWIO\\\",\\\"CAP_SYS_CHROOT\\\",\\\"CAP_SYS_PTRACE\\\",\\\"CAP_SYS_PACCT\\\",\\\"CAP_SYS_ADMIN\\\",\\\"CAP_SYS_BOOT\\\",\\\"CAP_SYS_NICE\\\",\\\"CAP_SYS_RESOURCE\\\",\\\"CAP_SYS_TIME\\\",\\\"CAP_SYS_TTY_CONFIG\\\",\\\"CAP_MKNOD\\\",\\\"CAP_LEASE\\\",\\\"CAP_AUDIT_WRITE\\\",\\\"CAP_AUDIT_CONTROL\\\",\\\"CAP_SETFCAP\\\",\\\"CAP_MAC_OVERRIDE\\\",\\\"CAP_MAC_ADMIN\\\",\\\"CAP_SYSLOG\\\",\\\"CAP_WAKE_ALARM\\\",\\\"CAP_BLOCK_SUSPEND\\\",\\\"CAP_AUDIT_READ\\\",\\\"CAP_PERFMON\\\",\\\"CAP_BPF\\\",\\\"CAP_CHECKPOINT_RESTORE\\\"],\\\"cgroup\\\":{\\\"v1\\\":true,\\\"v2\\\":true,\\\"systemd\\\":true,\\\"systemdUser\\\":true,\\\"rdma\\\":true},\\\"seccomp\\\":{\\\"enabled\\\":true,\\\"actions\\\":[\\\"SCMP_ACT_ALLOW\\\",\\\"SCMP_ACT_ERRNO\\\",\\\"SCMP_ACT_KILL\\\",\\\"SCMP_ACT_KILL_PROCESS\\\",\\\"SCMP_ACT_KILL_THREAD\\\",\\\"SCMP_ACT_LOG\\\",\\\"SCMP_ACT_NOTIFY\\\",\\\"SCMP_ACT_TRACE\\\",\\\"SCMP_ACT_TRAP\\\"],\\\"operators\\\":[\\\"SCMP_CMP_EQ\\\",\\\"SCMP_CMP_GE\\\",\\\"SCMP_CMP_GT\\\",\\\"SCMP_CMP_LE\\\",\\\"SCMP_CMP_LT\\\",\\\"SCMP_CMP_MASKED_EQ\\\",\\\"SCMP_CMP_NE\\\"],\\\"archs\\\":[\\\"SCMP_ARCH_AARCH64\\\",\\\"SCMP_ARCH_ARM\\\",\\\"SCMP_ARCH_MIPS\\\",\\\"SCMP_ARCH_MIPS64\\\",\\\"SCMP_ARCH_MIPS64N32\\\",\\\"SCMP_ARCH_MIPSEL\\\",\\\"SCMP_ARCH_MIPSEL64\\\",\\\"SCMP_ARCH_MIPSEL64N32\\\",\\\"SCMP_ARCH_PPC\\\",\\\"SCMP_ARCH_PPC64\\\",\\\"SCMP_ARCH_PPC64LE\\\",\\\"SCMP_ARCH_RISCV64\\\",\\\"SCMP_ARCH_S390\\\",\\\"SCMP_ARCH_S390X\\\",\\\"SCMP_ARCH_X32\\\",\\\"SCMP_ARCH_X86\\\",\\\"SCMP_ARCH_X86_64\\\"],\\\"knownFlags\\\":[\\\"SECCOMP_FILTER_FLAG_TSYNC\\\",\\\"SECCOMP_FILTER_FLAG_SPEC_ALLOW\\\",\\\"SECCOMP_FILTER_FLAG_LOG\\\"],\\\"supportedFlags\\\":[\\\"SECCOMP_FILTER_FLAG_TSYNC\\\",\\\"SECCOMP_FILTER_FLAG_SPEC_ALLOW\\\",\\\"SECCOMP_FILTER_FLAG_LOG\\\"]},\\\"apparmor\\\":{\\\"enabled\\\":true},\\\"selinux\\\":{\\\"enabled\\\":true},\\\"intelRdt\\\":{\\\"enabled\\\":true},\\\"mountExtensions\\\":{\\\"idmap\\\":{\\\"enabled\\\":true}}},\\\"annotations\\\":{\\\"io.github.seccomp.libseccomp.version\\\":\\\"2.5.3\\\",\\\"org.opencontainers.runc.checkpoint.enabled\\\":\\\"true\\\",\\\"org.opencontainers.runc.commit\\\":\\\"v1.2.5-0-g59923ef\\\",\\\"org.opencontainers.runc.version\\\":\\\"1.2.5\\\"},\\\"potentiallyUnsafeConfigAnnotations\\\":[\\\"bundle\\\",\\\"org.systemd.property.\\\",\\\"org.criu.config\\\"]}\"}},\"runc\":{\"path\":\"runc\",\"status\":{\"org.opencontainers.runtime-spec.features\":\"{\\\"ociVersionMin\\\":\\\"1.0.0\\\",\\\"ociVersionMax\\\":\\\"1.2.0\\\",\\\"hooks\\\":[\\\"prestart\\\",\\\"createRuntime\\\",\\\"createContainer\\\",\\\"startContainer\\\",\\\"poststart\\\",\\\"poststop\\\"],\\\"mountOptions\\\":[\\\"async\\\",\\\"atime\\\",\\\"bind\\\",\\\"defaults\\\",\\\"dev\\\",\\\"diratime\\\",\\\"dirsync\\\",\\\"exec\\\",\\\"iversion\\\",\\\"lazytime\\\",\\\"loud\\\",\\\"mand\\\",\\\"noatime\\\",\\\"nodev\\\",\\\"nodiratime\\\",\\\"noexec\\\",\\\"noiversion\\\",\\\"nolazytime\\\",\\\"nomand\\\",\\\"norelatime\\\",\\\"nostrictatime\\\",\\\"nosuid\\\",\\\"nosymfollow\\\",\\\"private\\\",\\\"ratime\\\",\\\"rbind\\\",\\\"rdev\\\",\\\"rdiratime\\\",\\\"relatime\\\",\\\"remount\\\",\\\"rexec\\\",\\\"rnoatime\\\",\\\"rnodev\\\",\\\"rnodiratime\\\",\\\"rnoexec\\\",\\\"rnorelatime\\\",\\\"rnostrictatime\\\",\\\"rnosuid\\\",\\\"rnosymfollow\\\",\\\"ro\\\",\\\"rprivate\\\",\\\"rrelatime\\\",\\\"rro\\\",\\\"rrw\\\",\\\"rshared\\\",\\\"rslave\\\",\\\"rstrictatime\\\",\\\"rsuid\\\",\\\"rsymfollow\\\",\\\"runbindable\\\",\\\"rw\\\",\\\"shared\\\",\\\"silent\\\",\\\"slave\\\",\\\"strictatime\\\",\\\"suid\\\",\\\"symfollow\\\",\\\"sync\\\",\\\"tmpcopyup\\\",\\\"unbindable\\\"],\\\"linux\\\":{\\\"namespaces\\\":[\\\"cgroup\\\",\\\"ipc\\\",\\\"mount\\\",\\\"network\\\",\\\"pid\\\",\\\"time\\\",\\\"user\\\",\\\"uts\\\"],\\\"capabilities\\\":[\\\"CAP_CHOWN\\\",\\\"CAP_DAC_OVERRIDE\\\",\\\"CAP_DAC_READ_SEARCH\\\",\\\"CAP_FOWNER\\\",\\\"CAP_FSETID\\\",\\\"CAP_KILL\\\",\\\"CAP_SETGID\\\",\\\"CAP_SETUID\\\",\\\"CAP_SETPCAP\\\",\\\"CAP_LINUX_IMMUTABLE\\\",\\\"CAP_NET_BIND_SERVICE\\\",\\\"CAP_NET_BROADCAST\\\",\\\"CAP_NET_ADMIN\\\",\\\"CAP_NET_RAW\\\",\\\"CAP_IPC_LOCK\\\",\\\"CAP_IPC_OWNER\\\",\\\"CAP_SYS_MODULE\\\",\\\"CAP_SYS_RAWIO\\\",\\\"CAP_SYS_CHROOT\\\",\\\"CAP_SYS_PTRACE\\\",\\\"CAP_SYS_PACCT\\\",\\\"CAP_SYS_ADMIN\\\",\\\"CAP_SYS_BOOT\\\",\\\"CAP_SYS_NICE\\\",\\\"CAP_SYS_RESOURCE\\\",\\\"CAP_SYS_TIME\\\",\\\"CAP_SYS_TTY_CONFIG\\\",\\\"CAP_MKNOD\\\",\\\"CAP_LEASE\\\",\\\"CAP_AUDIT_WRITE\\\",\\\"CAP_AUDIT_CONTROL\\\",\\\"CAP_SETFCAP\\\",\\\"CAP_MAC_OVERRIDE\\\",\\\"CAP_MAC_ADMIN\\\",\\\"CAP_SYSLOG\\\",\\\"CAP_WAKE_ALARM\\\",\\\"CAP_BLOCK_SUSPEND\\\",\\\"CAP_AUDIT_READ\\\",\\\"CAP_PERFMON\\\",\\\"CAP_BPF\\\",\\\"CAP_CHECKPOINT_RESTORE\\\"],\\\"cgroup\\\":{\\\"v1\\\":true,\\\"v2\\\":true,\\\"systemd\\\":true,\\\"systemdUser\\\":true,\\\"rdma\\\":true},\\\"seccomp\\\":{\\\"enabled\\\":true,\\\"actions\\\":[\\\"SCMP_ACT_ALLOW\\\",\\\"SCMP_ACT_ERRNO\\\",\\\"SCMP_ACT_KILL\\\",\\\"SCMP_ACT_KILL_PROCESS\\\",\\\"SCMP_ACT_KILL_THREAD\\\",\\\"SCMP_ACT_LOG\\\",\\\"SCMP_ACT_NOTIFY\\\",\\\"SCMP_ACT_TRACE\\\",\\\"SCMP_ACT_TRAP\\\"],\\\"operators\\\":[\\\"SCMP_CMP_EQ\\\",\\\"SCMP_CMP_GE\\\",\\\"SCMP_CMP_GT\\\",\\\"SCMP_CMP_LE\\\",\\\"SCMP_CMP_LT\\\",\\\"SCMP_CMP_MASKED_EQ\\\",\\\"SCMP_CMP_NE\\\"],\\\"archs\\\":[\\\"SCMP_ARCH_AARCH64\\\",\\\"SCMP_ARCH_ARM\\\",\\\"SCMP_ARCH_MIPS\\\",\\\"SCMP_ARCH_MIPS64\\\",\\\"SCMP_ARCH_MIPS64N32\\\",\\\"SCMP_ARCH_MIPSEL\\\",\\\"SCMP_ARCH_MIPSEL64\\\",\\\"SCMP_ARCH_MIPSEL64N32\\\",\\\"SCMP_ARCH_PPC\\\",\\\"SCMP_ARCH_PPC64\\\",\\\"SCMP_ARCH_PPC64LE\\\",\\\"SCMP_ARCH_RISCV64\\\",\\\"SCMP_ARCH_S390\\\",\\\"SCMP_ARCH_S390X\\\",\\\"SCMP_ARCH_X32\\\",\\\"SCMP_ARCH_X86\\\",\\\"SCMP_ARCH_X86_64\\\"],\\\"knownFlags\\\":[\\\"SECCOMP_FILTER_FLAG_TSYNC\\\",\\\"SECCOMP_FILTER_FLAG_SPEC_ALLOW\\\",\\\"SECCOMP_FILTER_FLAG_LOG\\\"],\\\"supportedFlags\\\":[\\\"SECCOMP_FILTER_FLAG_TSYNC\\\",\\\"SECCOMP_FILTER_FLAG_SPEC_ALLOW\\\",\\\"SECCOMP_FILTER_FLAG_LOG\\\"]},\\\"apparmor\\\":{\\\"enabled\\\":true},\\\"selinux\\\":{\\\"enabled\\\":true},\\\"intelRdt\\\":{\\\"enabled\\\":true},\\\"mountExtensions\\\":{\\\"idmap\\\":{\\\"enabled\\\":true}}},\\\"annotations\\\":{\\\"io.github.seccomp.libseccomp.version\\\":\\\"2.5.3\\\",\\\"org.opencontainers.runc.checkpoint.enabled\\\":\\\"true\\\",\\\"org.opencontainers.runc.commit\\\":\\\"v1.2.5-0-g59923ef\\\",\\\"org.opencontainers.runc.version\\\":\\\"1.2.5\\\"},\\\"potentiallyUnsafeConfigAnnotations\\\":[\\\"bundle\\\",\\\"org.systemd.property.\\\",\\\"org.criu.config\\\"]}\"}}}", + "stderr": "" + }, + "nvidia_smi": { + "ok": true, + "stdout": "GPU 0: NVIDIA A100 80GB PCIe (UUID: GPU-911fab14-997f-e653-67c4-abba6d597ec6)\nGPU 1: NVIDIA A100 80GB PCIe (UUID: GPU-b6584aae-31c5-e793-7d22-98fbebef889f)\nGPU 2: NVIDIA A100 80GB PCIe (UUID: GPU-b72f3ad9-d411-a9da-c4a2-be9f0cb175fe)\nGPU 3: NVIDIA A100 80GB PCIe (UUID: GPU-e1d3ca84-f769-2e6d-4179-d3ec0ebd94f2)", + "stderr": "" + }, + "nvidia_container_runtime": { + "ok": true, + "stdout": "/usr/bin/nvidia-container-runtime", + "stderr": "" + }, + "nvidia_ctk": { + "ok": true, + "stdout": "/usr/bin/nvidia-ctk", + "stderr": "" + }, + "docker_gpus": { + "ok": true, + "stdout": "GPU 0: NVIDIA A100 80GB PCIe (UUID: GPU-911fab14-997f-e653-67c4-abba6d597ec6)\nGPU 1: NVIDIA A100 80GB PCIe (UUID: GPU-b6584aae-31c5-e793-7d22-98fbebef889f)\nGPU 2: NVIDIA A100 80GB PCIe (UUID: GPU-b72f3ad9-d411-a9da-c4a2-be9f0cb175fe)\nGPU 3: NVIDIA A100 80GB PCIe (UUID: GPU-e1d3ca84-f769-2e6d-4179-d3ec0ebd94f2)", + "stderr": "" + }, + "docker_has_nvidia_runtime": true, + "recommended_fix": "Run: sudo nvidia-ctk runtime configure --runtime=docker --config=/etc/docker/daemon.json && sudo systemctl restart docker" + }, + "endpoints": [ + { + "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": 372.96746190078557, + "throughput_rps": 0.002681199037856052, + "latency_seconds": { + "min": 372.96655484382063, + "mean": 372.96655484382063, + "p50": 372.96655484382063, + "p95": 372.96655484382063, + "max": 372.96655484382063 + }, + "errors": [] + } + ] +} \ No newline at end of file diff --git a/logger.yaml b/logger.yaml new file mode 100644 index 0000000..df3929d --- /dev/null +++ b/logger.yaml @@ -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] \ No newline at end of file diff --git a/models.py b/models.py new file mode 100644 index 0000000..e4c46a2 --- /dev/null +++ b/models.py @@ -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), + } diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..63837ff --- /dev/null +++ b/pyproject.toml @@ -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"] diff --git a/requirements-bench.txt b/requirements-bench.txt new file mode 100644 index 0000000..f6936df --- /dev/null +++ b/requirements-bench.txt @@ -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 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..be3e907 --- /dev/null +++ b/requirements.txt @@ -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 diff --git a/scripts/capture_ocr_text.py b/scripts/capture_ocr_text.py new file mode 100644 index 0000000..6062398 --- /dev/null +++ b/scripts/capture_ocr_text.py @@ -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() diff --git a/scripts/cer_divergence.py b/scripts/cer_divergence.py new file mode 100644 index 0000000..91f4468 --- /dev/null +++ b/scripts/cer_divergence.py @@ -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() diff --git a/scripts/configure_nvidia_docker.sh b/scripts/configure_nvidia_docker.sh new file mode 100755 index 0000000..8eff0e9 --- /dev/null +++ b/scripts/configure_nvidia_docker.sh @@ -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 diff --git a/scripts/docker_nvidia_runtime_compat.sh b/scripts/docker_nvidia_runtime_compat.sh new file mode 100755 index 0000000..cd5a9ed --- /dev/null +++ b/scripts/docker_nvidia_runtime_compat.sh @@ -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[@]}" diff --git a/scripts/parse_timing.py b/scripts/parse_timing.py new file mode 100644 index 0000000..62323ae --- /dev/null +++ b/scripts/parse_timing.py @@ -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() diff --git a/scripts/poll_vllm_metrics.py b/scripts/poll_vllm_metrics.py new file mode 100644 index 0000000..8db999b --- /dev/null +++ b/scripts/poll_vllm_metrics.py @@ -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() diff --git a/scripts/quant/__init__.py b/scripts/quant/__init__.py new file mode 100644 index 0000000..8f99522 --- /dev/null +++ b/scripts/quant/__init__.py @@ -0,0 +1 @@ +"""Quantization benchmark harness for the Surya-OCR-2 recognition VLM.""" diff --git a/scripts/quant/aggregate.py b/scripts/quant/aggregate.py new file mode 100644 index 0000000..66f6546 --- /dev/null +++ b/scripts/quant/aggregate.py @@ -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) diff --git a/scripts/quant/bbox_iou.py b/scripts/quant/bbox_iou.py new file mode 100644 index 0000000..12b9f03 --- /dev/null +++ b/scripts/quant/bbox_iou.py @@ -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, + } diff --git a/scripts/quant/build_model.py b/scripts/quant/build_model.py new file mode 100644 index 0000000..4c9b563 --- /dev/null +++ b/scripts/quant/build_model.py @@ -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 diff --git a/scripts/quant/capture.py b/scripts/quant/capture.py new file mode 100644 index 0000000..9485309 --- /dev/null +++ b/scripts/quant/capture.py @@ -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() diff --git a/scripts/quant/finalize.py b/scripts/quant/finalize.py new file mode 100644 index 0000000..dfec336 --- /dev/null +++ b/scripts/quant/finalize.py @@ -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() diff --git a/scripts/quant/manifest.py b/scripts/quant/manifest.py new file mode 100644 index 0000000..87685d8 --- /dev/null +++ b/scripts/quant/manifest.py @@ -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 diff --git a/scripts/quant/plot.py b/scripts/quant/plot.py new file mode 100644 index 0000000..41c1c41 --- /dev/null +++ b/scripts/quant/plot.py @@ -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 diff --git a/scripts/quant/recipes.py b/scripts/quant/recipes.py new file mode 100644 index 0000000..7f0202c --- /dev/null +++ b/scripts/quant/recipes.py @@ -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 diff --git a/scripts/quant/run_all.py b/scripts/quant/run_all.py new file mode 100644 index 0000000..086bb70 --- /dev/null +++ b/scripts/quant/run_all.py @@ -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() diff --git a/scripts/quant/score_methods.sh b/scripts/quant/score_methods.sh new file mode 100644 index 0000000..96c92d2 --- /dev/null +++ b/scripts/quant/score_methods.sh @@ -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" diff --git a/scripts/quant/serve.py b/scripts/quant/serve.py new file mode 100644 index 0000000..186ec26 --- /dev/null +++ b/scripts/quant/serve.py @@ -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)]) diff --git a/scripts/quant/tune_mtp.sh b/scripts/quant/tune_mtp.sh new file mode 100644 index 0000000..cc0c4cc --- /dev/null +++ b/scripts/quant/tune_mtp.sh @@ -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" diff --git a/scripts/quant/tune_vllm.sh b/scripts/quant/tune_vllm.sh new file mode 100644 index 0000000..f5fc487 --- /dev/null +++ b/scripts/quant/tune_vllm.sh @@ -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" diff --git a/scripts/start_single_container.sh b/scripts/start_single_container.sh new file mode 100755 index 0000000..5b2a7ec --- /dev/null +++ b/scripts/start_single_container.sh @@ -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 diff --git a/surya/__init__.py b/surya/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/surya/common/__init__.py b/surya/common/__init__.py new file mode 100644 index 0000000..b28b04f --- /dev/null +++ b/surya/common/__init__.py @@ -0,0 +1,3 @@ + + + diff --git a/surya/common/blank.py b/surya/common/blank.py new file mode 100644 index 0000000..d11c721 --- /dev/null +++ b/surya/common/blank.py @@ -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 diff --git a/surya/common/load.py b/surya/common/load.py new file mode 100644 index 0000000..e1c7643 --- /dev/null +++ b/surya/common/load.py @@ -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() diff --git a/surya/common/polygon.py b/surya/common/polygon.py new file mode 100644 index 0000000..80e7aec --- /dev/null +++ b/surya/common/polygon.py @@ -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)) diff --git a/surya/common/predictor.py b/surya/common/predictor.py new file mode 100644 index 0000000..d7e6c72 --- /dev/null +++ b/surya/common/predictor.py @@ -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() diff --git a/surya/common/pretrained.py b/surya/common/pretrained.py new file mode 100644 index 0000000..6bf5606 --- /dev/null +++ b/surya/common/pretrained.py @@ -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 diff --git a/surya/common/s3.py b/surya/common/s3.py new file mode 100644 index 0000000..98b42de --- /dev/null +++ b/surya/common/s3.py @@ -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) diff --git a/surya/common/util.py b/surya/common/util.py new file mode 100644 index 0000000..bd093f5 --- /dev/null +++ b/surya/common/util.py @@ -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, + ] diff --git a/surya/debug/draw.py b/surya/debug/draw.py new file mode 100644 index 0000000..3d15731 --- /dev/null +++ b/surya/debug/draw.py @@ -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 diff --git a/surya/debug/fonts.py b/surya/debug/fonts.py new file mode 100644 index 0000000..e9e1878 --- /dev/null +++ b/surya/debug/fonts.py @@ -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 \ No newline at end of file diff --git a/surya/debug/katex.js b/surya/debug/katex.js new file mode 100644 index 0000000..ac12e7e --- /dev/null +++ b/surya/debug/katex.js @@ -0,0 +1,64 @@ + + + + \ No newline at end of file diff --git a/surya/debug/render_html.py b/surya/debug/render_html.py new file mode 100644 index 0000000..a0f93e1 --- /dev/null +++ b/surya/debug/render_html.py @@ -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""" + + + + + {katex_script} + + +""") + + 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'{text}') + else: + # Plain text, escape it + escaped_text = htmllib.escape(text) + html_content.append(f'{escaped_text}') + + html_content.append("") + + return "\n".join(html_content), image_size \ No newline at end of file diff --git a/surya/debug/text.py b/surya/debug/text.py new file mode 100644 index 0000000..ce120e2 --- /dev/null +++ b/surya/debug/text.py @@ -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 diff --git a/surya/detection/__init__.py b/surya/detection/__init__.py new file mode 100644 index 0000000..298bceb --- /dev/null +++ b/surya/detection/__init__.py @@ -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() diff --git a/surya/detection/heatmap.py b/surya/detection/heatmap.py new file mode 100644 index 0000000..93ffe04 --- /dev/null +++ b/surya/detection/heatmap.py @@ -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 diff --git a/surya/detection/loader.py b/surya/detection/loader.py new file mode 100644 index 0000000..189ed0e --- /dev/null +++ b/surya/detection/loader.py @@ -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) diff --git a/surya/detection/model/__init__.py b/surya/detection/model/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/surya/detection/model/config.py b/surya/detection/model/config.py new file mode 100644 index 0000000..b205ab4 --- /dev/null +++ b/surya/detection/model/config.py @@ -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 \ No newline at end of file diff --git a/surya/detection/model/encoderdecoder.py b/surya/detection/model/encoderdecoder.py new file mode 100644 index 0000000..af5ea9c --- /dev/null +++ b/surya/detection/model/encoderdecoder.py @@ -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 + ) diff --git a/surya/detection/parallel.py b/surya/detection/parallel.py new file mode 100644 index 0000000..2779c8e --- /dev/null +++ b/surya/detection/parallel.py @@ -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) diff --git a/surya/detection/processor.py b/surya/detection/processor.py new file mode 100644 index 0000000..1cb44a0 --- /dev/null +++ b/surya/detection/processor.py @@ -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) diff --git a/surya/detection/schema.py b/surya/detection/schema.py new file mode 100644 index 0000000..5e32088 --- /dev/null +++ b/surya/detection/schema.py @@ -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] diff --git a/surya/detection/util.py b/surya/detection/util.py new file mode 100644 index 0000000..594cf4d --- /dev/null +++ b/surya/detection/util.py @@ -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] diff --git a/surya/endpoint/__init__.py b/surya/endpoint/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/surya/endpoint/app.py b/surya/endpoint/app.py new file mode 100644 index 0000000..cb78481 --- /dev/null +++ b/surya/endpoint/app.py @@ -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) diff --git a/surya/endpoint/legacy.py b/surya/endpoint/legacy.py new file mode 100644 index 0000000..b9bfdab --- /dev/null +++ b/surya/endpoint/legacy.py @@ -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 "

Welcome to SURYA OCR API!

" + + +@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 diff --git a/surya/endpoint/openai.py b/surya/endpoint/openai.py new file mode 100644 index 0000000..5c3e391 --- /dev/null +++ b/surya/endpoint/openai.py @@ -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")) diff --git a/surya/endpoint/schemas.py b/surya/endpoint/schemas.py new file mode 100644 index 0000000..715978b --- /dev/null +++ b/surya/endpoint/schemas.py @@ -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 diff --git a/surya/endpoint/service.py b/surya/endpoint/service.py new file mode 100644 index 0000000..1d56fb1 --- /dev/null +++ b/surya/endpoint/service.py @@ -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} diff --git a/surya/inference/__init__.py b/surya/inference/__init__.py new file mode 100644 index 0000000..4686c22 --- /dev/null +++ b/surya/inference/__init__.py @@ -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 diff --git a/surya/inference/backends/__init__.py b/surya/inference/backends/__init__.py new file mode 100644 index 0000000..f81cf02 --- /dev/null +++ b/surya/inference/backends/__init__.py @@ -0,0 +1,2 @@ +from surya.inference.backends.base import Backend as Backend +from surya.inference.backends.base import ServerHandle as ServerHandle diff --git a/surya/inference/backends/base.py b/surya/inference/backends/base.py new file mode 100644 index 0000000..d97dad2 --- /dev/null +++ b/surya/inference/backends/base.py @@ -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 diff --git a/surya/inference/backends/llamacpp.py b/surya/inference/backends/llamacpp.py new file mode 100644 index 0000000..c97c6fe --- /dev/null +++ b/surya/inference/backends/llamacpp.py @@ -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, + ) diff --git a/surya/inference/backends/openai_client.py b/surya/inference/backends/openai_client.py new file mode 100644 index 0000000..5ae167d --- /dev/null +++ b/surya/inference/backends/openai_client.py @@ -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)) diff --git a/surya/inference/backends/spawn.py b/surya/inference/backends/spawn.py new file mode 100644 index 0000000..e1d6ff1 --- /dev/null +++ b/surya/inference/backends/spawn.py @@ -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" diff --git a/surya/inference/backends/vllm.py b/surya/inference/backends/vllm.py new file mode 100644 index 0000000..a51420c --- /dev/null +++ b/surya/inference/backends/vllm.py @@ -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, + ) diff --git a/surya/inference/parsers.py b/surya/inference/parsers.py new file mode 100644 index 0000000..728d3a6 --- /dev/null +++ b/surya/inference/parsers.py @@ -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
inner HTML
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, + ) diff --git a/surya/inference/prompts.py b/surya/inference/prompts.py new file mode 100644 index 0000000..c9a1425 --- /dev/null +++ b/surya/inference/prompts.py @@ -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, + }, +} diff --git a/surya/inference/schema.py b/surya/inference/schema.py new file mode 100644 index 0000000..f171da5 --- /dev/null +++ b/surya/inference/schema.py @@ -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) diff --git a/surya/inference/util.py b/surya/inference/util.py new file mode 100644 index 0000000..1691170 --- /dev/null +++ b/surya/inference/util.py @@ -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) diff --git a/surya/input/load.py b/surya/input/load.py new file mode 100644 index 0000000..4b5fd27 --- /dev/null +++ b/surya/input/load.py @@ -0,0 +1,77 @@ +from typing import List +import PIL + +from surya.input.processing import open_pdf, get_page_images +from surya.logging import get_logger +from surya.settings import settings +import os +import filetype +from PIL import Image + +logger = get_logger() + + +def get_name_from_path(path): + return os.path.basename(path).split(".")[0] + + +def load_pdf(pdf_path, page_range: List[int] | None = None, dpi=settings.IMAGE_DPI): + doc = open_pdf(pdf_path) + last_page = len(doc) + + if page_range: + assert all([0 <= page < last_page for page in page_range]), ( + f"Invalid page range: {page_range}" + ) + else: + page_range = list(range(last_page)) + + images = get_page_images(doc, page_range, dpi=dpi) + doc.close() + names = [get_name_from_path(pdf_path) for _ in page_range] + return images, names + + +def load_image(image_path): + image = Image.open(image_path).convert("RGB") + name = get_name_from_path(image_path) + return [image], [name] + + +def load_from_file( + input_path, page_range: List[int] | None = None, dpi=settings.IMAGE_DPI +): + input_type = filetype.guess(input_path) + if input_type and input_type.extension == "pdf": + return load_pdf(input_path, page_range, dpi=dpi) + else: + return load_image(input_path) + + +def load_from_folder( + folder_path, page_range: List[int] | None = None, dpi=settings.IMAGE_DPI +): + image_paths = [ + os.path.join(folder_path, image_name) + for image_name in os.listdir(folder_path) + if not image_name.startswith(".") + ] + image_paths = [ip for ip in image_paths if not os.path.isdir(ip)] + + images = [] + names = [] + for path in image_paths: + extension = filetype.guess(path) + if extension and extension.extension == "pdf": + image, name = load_pdf(path, page_range, dpi=dpi) + images.extend(image) + names.extend(name) + else: + try: + image, name = load_image(path) + images.extend(image) + names.extend(name) + except PIL.UnidentifiedImageError: + logger.warning(f"Could not load image {path}") + continue + return images, names diff --git a/surya/input/processing.py b/surya/input/processing.py new file mode 100644 index 0000000..03a6003 --- /dev/null +++ b/surya/input/processing.py @@ -0,0 +1,17 @@ +from typing import List + +import pypdfium2 + +from surya.settings import settings + + +def open_pdf(pdf_filepath): + return pypdfium2.PdfDocument(pdf_filepath) + + +def get_page_images(doc, indices: List, dpi=settings.IMAGE_DPI): + images = [ + doc[i].render(scale=dpi / 72, draw_annots=False).to_pil() for i in indices + ] + images = [image.convert("RGB") for image in images] + return images diff --git a/surya/layout/__init__.py b/surya/layout/__init__.py new file mode 100644 index 0000000..f5b1213 --- /dev/null +++ b/surya/layout/__init__.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +from typing import List, Optional + +from PIL import Image + +from surya.common.blank import is_blank_region +from surya.inference import SuryaInferenceManager, get_default_manager +from surya.inference.parsers import denorm_bbox, parse_layout +from surya.inference.prompts import LAYOUT_JSON_SCHEMA, PROMPT_TYPE_LAYOUT +from surya.inference.schema import BatchInputItem +from surya.layout.label import LAYOUT_PRED_RELABEL, TEXT_LABELS +from surya.layout.schema import LayoutBox, LayoutResult +from surya.logging import get_logger +from surya.settings import settings +from surya.timing import timing_span + +logger = get_logger() + + +class LayoutPredictor: + """Run LAYOUT_PROMPT on full pages, parse JSON, return LayoutResult per image.""" + + def __init__(self, manager: Optional[SuryaInferenceManager] = None): + self.manager = manager # If None, get_default_manager() is used at call time + self._disable_tqdm = settings.DISABLE_TQDM + + @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 to(self, *args, **kwargs): + # Manager-backed; .to() is a no-op for compatibility with BasePredictor callers. + return + + def __call__( + self, + images: List[Image.Image], + target_image_sizes: Optional[List[tuple]] = None, + max_tokens: Optional[int] = None, + ) -> List[LayoutResult]: + """Run layout on a batch of images. + + target_image_sizes: optional list of (width, height) tuples — if + provided, bboxes are denormalized to these sizes instead of each + input image's size. Useful when layout runs on a low-DPI render but + you want bboxes in the OCR image's coordinate space. + """ + if not images: + return [] + manager = self.manager or get_default_manager() + + max_tokens = max_tokens or settings.SURYA_MAX_TOKENS_LAYOUT + guided = LAYOUT_JSON_SCHEMA if settings.SURYA_GUIDED_LAYOUT else None + with timing_span("layout_build_batch", image_count=len(images), max_tokens=max_tokens): + batch = [ + BatchInputItem( + image=img, + prompt_type=PROMPT_TYPE_LAYOUT, + max_tokens=max_tokens, + guided_json=guided, + ) + for img in images + ] + with timing_span("layout_manager_generate", item_count=len(batch)): + outputs = manager.generate(batch) + + if target_image_sizes is not None and len(target_image_sizes) != len(images): + raise ValueError("target_image_sizes must match images length") + + with timing_span("layout_parse_outputs", item_count=len(outputs)): + results: List[LayoutResult] = [] + for idx, (img, out) in enumerate(zip(images, outputs)): + if target_image_sizes is not None: + w, h = target_image_sizes[idx] + else: + w, h = img.size + page_bbox = [0, 0, float(w), float(h)] + if out.error or not out.raw: + results.append( + LayoutResult( + bboxes=[], image_bbox=page_bbox, raw=out.raw, error=True + ) + ) + continue + try: + parsed = parse_layout(out.raw) + except Exception as e: + logger.warning(f"Layout parse failed: {e}; raw[:300]={out.raw[:300]!r}") + results.append( + LayoutResult( + bboxes=[], image_bbox=page_bbox, raw=out.raw, error=True + ) + ) + continue + + confidence = out.mean_token_prob if out.mean_token_prob is not None else 1.0 + img_w, img_h = img.size + boxes: List[LayoutBox] = [] + dropped_blank = 0 + for blk in parsed: + canon = LAYOUT_PRED_RELABEL.get(blk.label, blk.label) + # Drop text-labeled blocks the model hallucinated over an + # essentially-blank region (mostly white OR near-uniform + # color). Visual blocks (Picture / Figure / Table / etc.) + # are allowed to be uniform — that's normal content. + if canon in TEXT_LABELS: + img_bbox = denorm_bbox( + blk.bbox, img_w, img_h, scale=settings.BBOX_SCALE + ) + x0, y0, x1, y1 = (max(0, int(v)) for v in img_bbox) + if x1 > x0 and y1 > y0: + if is_blank_region(img.crop((x0, y0, x1, y1))): + dropped_blank += 1 + continue + pixel_bbox = denorm_bbox(blk.bbox, w, h, scale=settings.BBOX_SCALE) + boxes.append( + LayoutBox( + polygon=list(pixel_bbox), + label=canon, + raw_label=blk.label, + position=len(boxes), + count=blk.count, + confidence=confidence, + ) + ) + if dropped_blank: + logger.info( + f"dropped {dropped_blank} text-labeled layout block(s) over " + f"blank/uniform regions" + ) + results.append( + LayoutResult( + bboxes=boxes, image_bbox=page_bbox, raw=out.raw, error=False + ) + ) + return results diff --git a/surya/layout/label.py b/surya/layout/label.py new file mode 100644 index 0000000..bb4c6e3 --- /dev/null +++ b/surya/layout/label.py @@ -0,0 +1,44 @@ +"""Surya2 layout labels emitted by the model + canonicalization to surya's +public label vocabulary.""" + +# Canonical text-bearing labels — used by blank-region filters to decide +# which blocks may be dropped when their underlying image region is empty. +# Excludes Picture/Figure/Diagram/Table/Form/Equation/etc., which can legitimately +# contain whitespace or solid fills. +TEXT_LABELS = frozenset( + { + "Text", + "SectionHeader", + "PageHeader", + "PageFooter", + "Caption", + "Footnote", + "Code", + "Bibliography", + } +) + + +# Canonicalize raw model labels to public surya label names. Marker and other +# downstream consumers depend on these names. +LAYOUT_PRED_RELABEL = { + "Caption": "Caption", + "Footnote": "Footnote", + "Equation-Block": "Equation", + "List-Group": "ListGroup", + "Page-Header": "PageHeader", + "Page-Footer": "PageFooter", + "Image": "Picture", + "Section-Header": "SectionHeader", + "Table": "Table", + "Text": "Text", + "Complex-Block": "Figure", + "Code-Block": "Code", + "Form": "Form", + "Table-Of-Contents": "TableOfContents", + "Figure": "Figure", + "Chemical-Block": "ChemicalBlock", + "Diagram": "Diagram", + "Bibliography": "Bibliography", + "Blank-Page": "BlankPage", +} diff --git a/surya/layout/schema.py b/surya/layout/schema.py new file mode 100644 index 0000000..d58acdd --- /dev/null +++ b/surya/layout/schema.py @@ -0,0 +1,19 @@ +from typing import List, Optional + +from pydantic import BaseModel + +from surya.common.polygon import PolygonBox + + +class LayoutBox(PolygonBox): + label: str # canonicalized via LAYOUT_PRED_RELABEL + raw_label: str # original model label, before canonicalization + position: int # reading order index + count: int = 0 # model's token estimate for OCR output (multiple of 50) + + +class LayoutResult(BaseModel): + bboxes: List[LayoutBox] + image_bbox: List[float] + raw: Optional[str] = None # raw model output, useful for debugging + error: bool = False diff --git a/surya/logging.py b/surya/logging.py new file mode 100644 index 0000000..ee51dbd --- /dev/null +++ b/surya/logging.py @@ -0,0 +1,27 @@ +import logging +import warnings +from surya.settings import settings + + +def configure_logging(): + logger = get_logger() + + # Remove any existing handlers to prevent duplicates + for handler in logger.handlers[:]: + logger.removeHandler(handler) + + # Add our handler + handler = logging.StreamHandler() + formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s") + handler.setFormatter(formatter) + logger.addHandler(handler) + + # Prevent propagation to parent loggers to avoid double logging + logger.propagate = False + + logger.setLevel(settings.LOGLEVEL) + warnings.simplefilter(action="ignore", category=FutureWarning) + + +def get_logger(): + return logging.getLogger("surya") diff --git a/surya/ocr_error/__init__.py b/surya/ocr_error/__init__.py new file mode 100644 index 0000000..dd87657 --- /dev/null +++ b/surya/ocr_error/__init__.py @@ -0,0 +1,52 @@ +import math +from typing import List, Optional + +from tqdm import tqdm + +from surya.common.predictor import BasePredictor +from surya.ocr_error.loader import OCRErrorModelLoader +from surya.ocr_error.model.config import ID2LABEL +from surya.ocr_error.schema import OCRErrorDetectionResult +from surya.settings import settings + + +class OCRErrorPredictor(BasePredictor): + model_loader_cls = OCRErrorModelLoader + batch_size = settings.OCR_ERROR_BATCH_SIZE + default_batch_sizes = {"cpu": 8, "mps": 8, "cuda": 64} + + def __call__(self, texts: List[str], batch_size: Optional[int] = None): + return self.batch_ocr_error_detection(texts, batch_size) + + def batch_ocr_error_detection( + self, texts: List[str], batch_size: Optional[int] = None + ): + if batch_size is None: + batch_size = self.get_batch_size() + + num_batches = math.ceil(len(texts) / batch_size) + texts_processed = self.processor( + texts, padding="longest", truncation=True, return_tensors="pt" + ) + predictions = [] + for batch_idx in tqdm( + range(num_batches), + desc="Running OCR Error Detection", + disable=self.disable_tqdm, + ): + start_idx, end_idx = batch_idx * batch_size, (batch_idx + 1) * batch_size + batch_input_ids = texts_processed.input_ids[start_idx:end_idx].to( + self.model.device + ) + batch_attention_mask = texts_processed.attention_mask[start_idx:end_idx].to( + self.model.device + ) + + with settings.INFERENCE_MODE(): + pred = self.model(batch_input_ids, attention_mask=batch_attention_mask) + logits = pred.logits.argmax(dim=1).cpu().tolist() + predictions.extend(logits) + + return OCRErrorDetectionResult( + texts=texts, labels=[ID2LABEL[p] for p in predictions] + ) diff --git a/surya/ocr_error/loader.py b/surya/ocr_error/loader.py new file mode 100644 index 0000000..d1519ab --- /dev/null +++ b/surya/ocr_error/loader.py @@ -0,0 +1,48 @@ +from typing import Optional + + +from surya.common.load import ModelLoader +from surya.logging import get_logger +from surya.ocr_error.model.config import DistilBertConfig +from surya.ocr_error.model.encoder import DistilBertForSequenceClassification +from surya.ocr_error.tokenizer import DistilBertTokenizer +from surya.settings import settings + +logger = get_logger() + + +class OCRErrorModelLoader(ModelLoader): + def __init__(self, checkpoint: Optional[str] = None): + super().__init__(checkpoint) + + if self.checkpoint is None: + self.checkpoint = settings.OCR_ERROR_MODEL_CHECKPOINT + + def model( + self, + device=settings.TORCH_DEVICE_MODEL, + dtype=settings.MODEL_DTYPE, + attention_implementation: Optional[str] = None, + ) -> DistilBertForSequenceClassification: + if device is None: + device = settings.TORCH_DEVICE_MODEL + if dtype is None: + dtype = settings.MODEL_DTYPE + + config = DistilBertConfig.from_pretrained(self.checkpoint) + model = ( + DistilBertForSequenceClassification.from_pretrained( + self.checkpoint, + dtype=dtype, + config=config, + ) + .to(device) + .eval() + ) + + return model + + def processor( + self, device=settings.TORCH_DEVICE_MODEL, dtype=settings.MODEL_DTYPE + ) -> DistilBertTokenizer: + return DistilBertTokenizer.from_pretrained(self.checkpoint) diff --git a/surya/ocr_error/model/__init__.py b/surya/ocr_error/model/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/surya/ocr_error/model/config.py b/surya/ocr_error/model/config.py new file mode 100644 index 0000000..9d46621 --- /dev/null +++ b/surya/ocr_error/model/config.py @@ -0,0 +1,68 @@ +from collections import OrderedDict +from typing import Mapping + +from transformers.configuration_utils import PretrainedConfig +from transformers.onnx import OnnxConfig + +from surya.common.s3 import S3DownloaderMixin + +ID2LABEL = { + 0: 'good', + 1: 'bad' +} + +class DistilBertConfig(S3DownloaderMixin, PretrainedConfig): + model_type = "distilbert" + attribute_map = { + "hidden_size": "dim", + "num_attention_heads": "n_heads", + "num_hidden_layers": "n_layers", + } + + def __init__( + self, + vocab_size=30522, + max_position_embeddings=512, + sinusoidal_pos_embds=False, + n_layers=6, + n_heads=12, + dim=768, + hidden_dim=4 * 768, + dropout=0.1, + attention_dropout=0.1, + activation="gelu", + initializer_range=0.02, + qa_dropout=0.1, + seq_classif_dropout=0.2, + pad_token_id=0, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.sinusoidal_pos_embds = sinusoidal_pos_embds + self.n_layers = n_layers + self.n_heads = n_heads + self.dim = dim + self.hidden_dim = hidden_dim + self.dropout = dropout + self.attention_dropout = attention_dropout + self.activation = activation + self.initializer_range = initializer_range + self.qa_dropout = qa_dropout + self.seq_classif_dropout = seq_classif_dropout + super().__init__(**kwargs, pad_token_id=pad_token_id) + + +class DistilBertOnnxConfig(OnnxConfig): + @property + def inputs(self) -> Mapping[str, Mapping[int, str]]: + if self.task == "multiple-choice": + dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"} + else: + dynamic_axis = {0: "batch", 1: "sequence"} + return OrderedDict( + [ + ("input_ids", dynamic_axis), + ("attention_mask", dynamic_axis), + ] + ) \ No newline at end of file diff --git a/surya/ocr_error/model/encoder.py b/surya/ocr_error/model/encoder.py new file mode 100644 index 0000000..cac1c50 --- /dev/null +++ b/surya/ocr_error/model/encoder.py @@ -0,0 +1,910 @@ +from __future__ import annotations + +import math +from typing import Optional, Set, List, Tuple, Union, Dict + +import numpy as np +import torch +from torch import nn +from torch.nn import functional as F, MSELoss, CrossEntropyLoss, BCEWithLogitsLoss +from transformers import apply_chunking_to_forward +from transformers.activations import get_activation +from transformers.modeling_outputs import BaseModelOutput, SequenceClassifierOutput +from transformers.pytorch_utils import ( + find_pruneable_heads_and_indices, + prune_linear_layer, +) + +from transformers.utils import ( + is_flash_attn_greater_or_equal_2_10, +) + +from surya.common.pretrained import SuryaPreTrainedModel + +from surya.common.s3 import S3DownloaderMixin +from surya.ocr_error.model.config import DistilBertConfig + + +def _get_unpad_data(attention_mask): + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +def create_sinusoidal_embeddings(n_pos: int, dim: int, out: torch.Tensor): + position_enc = np.array( + [ + [pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] + for pos in range(n_pos) + ] + ) + out.requires_grad = False + out[:, 0::2] = torch.FloatTensor(np.sin(position_enc[:, 0::2])) + out[:, 1::2] = torch.FloatTensor(np.cos(position_enc[:, 1::2])) + out.detach_() + + +class Embeddings(nn.Module): + def __init__(self, config: DistilBertConfig): + super().__init__() + self.word_embeddings = nn.Embedding( + config.vocab_size, config.dim, padding_idx=config.pad_token_id + ) + self.position_embeddings = nn.Embedding( + config.max_position_embeddings, config.dim + ) + + self.LayerNorm = nn.LayerNorm(config.dim, eps=1e-12) + self.dropout = nn.Dropout(config.dropout) + self.register_buffer( + "position_ids", + torch.arange(config.max_position_embeddings).expand((1, -1)), + persistent=False, + ) + + def forward( + self, input_ids: torch.Tensor, input_embeds: Optional[torch.Tensor] = None + ) -> torch.Tensor: + """ + Parameters: + input_ids (torch.Tensor): + torch.tensor(bs, max_seq_length) The token ids to embed. + input_embeds (*optional*, torch.Tensor): + The pre-computed word embeddings. Can only be passed if the input ids are `None`. + + + Returns: torch.tensor(bs, max_seq_length, dim) The embedded tokens (plus position embeddings, no token_type + embeddings) + """ + if input_ids is not None: + input_embeds = self.word_embeddings(input_ids) # (bs, max_seq_length, dim) + + seq_length = input_embeds.size(1) + + # Setting the position-ids to the registered buffer in constructor, it helps + # when tracing the model without passing position-ids, solves + # isues similar to issue #5664 + if hasattr(self, "position_ids"): + position_ids = self.position_ids[:, :seq_length] + else: + position_ids = torch.arange( + seq_length, dtype=torch.long, device=input_ids.device + ) # (max_seq_length) + position_ids = position_ids.unsqueeze(0).expand_as( + input_ids + ) # (bs, max_seq_length) + + position_embeddings = self.position_embeddings( + position_ids + ) # (bs, max_seq_length, dim) + + embeddings = input_embeds + position_embeddings # (bs, max_seq_length, dim) + embeddings = self.LayerNorm(embeddings) # (bs, max_seq_length, dim) + embeddings = self.dropout(embeddings) # (bs, max_seq_length, dim) + return embeddings + + +class MultiHeadSelfAttention(nn.Module): + def __init__(self, config: DistilBertConfig): + super().__init__() + self.config = config + + self.n_heads = config.n_heads + self.dim = config.dim + self.dropout = nn.Dropout(p=config.attention_dropout) + self.is_causal = False + + # Have an even number of multi heads that divide the dimensions + if self.dim % self.n_heads != 0: + # Raise value errors for even multi-head attention nodes + raise ValueError( + f"self.n_heads: {self.n_heads} must divide self.dim: {self.dim} evenly" + ) + + self.q_lin = nn.Linear(in_features=config.dim, out_features=config.dim) + self.k_lin = nn.Linear(in_features=config.dim, out_features=config.dim) + self.v_lin = nn.Linear(in_features=config.dim, out_features=config.dim) + self.out_lin = nn.Linear(in_features=config.dim, out_features=config.dim) + + self.pruned_heads: Set[int] = set() + self.attention_head_size = self.dim // self.n_heads + + def prune_heads(self, heads: List[int]): + if len(heads) == 0: + return + heads, index = find_pruneable_heads_and_indices( + heads, self.n_heads, self.attention_head_size, self.pruned_heads + ) + # Prune linear layers + self.q_lin = prune_linear_layer(self.q_lin, index) + self.k_lin = prune_linear_layer(self.k_lin, index) + self.v_lin = prune_linear_layer(self.v_lin, index) + self.out_lin = prune_linear_layer(self.out_lin, index, dim=1) + # Update hyper params + self.n_heads = self.n_heads - len(heads) + self.dim = self.attention_head_size * self.n_heads + self.pruned_heads = self.pruned_heads.union(heads) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: torch.Tensor, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + ) -> Tuple[torch.Tensor, ...]: + """ + Parameters: + query: torch.tensor(bs, seq_length, dim) + key: torch.tensor(bs, seq_length, dim) + value: torch.tensor(bs, seq_length, dim) + mask: torch.tensor(bs, seq_length) + + Returns: + weights: torch.tensor(bs, n_heads, seq_length, seq_length) Attention weights context: torch.tensor(bs, + seq_length, dim) Contextualized layer. Optional: only if `output_attentions=True` + """ + bs, q_length, dim = query.size() + k_length = key.size(1) + # assert dim == self.dim, f'Dimensions do not match: {dim} input vs {self.dim} configured' + # assert key.size() == value.size() + + dim_per_head = self.dim // self.n_heads + + mask_reshp = (bs, 1, 1, k_length) + + def shape(x: torch.Tensor) -> torch.Tensor: + """separate heads""" + return x.view(bs, -1, self.n_heads, dim_per_head).transpose(1, 2) + + def unshape(x: torch.Tensor) -> torch.Tensor: + """group heads""" + return ( + x.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * dim_per_head) + ) + + q = shape(self.q_lin(query)) # (bs, n_heads, q_length, dim_per_head) + k = shape(self.k_lin(key)) # (bs, n_heads, k_length, dim_per_head) + v = shape(self.v_lin(value)) # (bs, n_heads, k_length, dim_per_head) + + q = q / math.sqrt(dim_per_head) # (bs, n_heads, q_length, dim_per_head) + scores = torch.matmul(q, k.transpose(2, 3)) # (bs, n_heads, q_length, k_length) + mask = ( + (mask == 0).view(mask_reshp).expand_as(scores) + ) # (bs, n_heads, q_length, k_length) + scores = scores.masked_fill( + mask, torch.tensor(torch.finfo(scores.dtype).min) + ) # (bs, n_heads, q_length, k_length) + + weights = nn.functional.softmax( + scores, dim=-1 + ) # (bs, n_heads, q_length, k_length) + weights = self.dropout(weights) # (bs, n_heads, q_length, k_length) + + # Mask heads if we want to + if head_mask is not None: + weights = weights * head_mask + + context = torch.matmul(weights, v) # (bs, n_heads, q_length, dim_per_head) + context = unshape(context) # (bs, q_length, dim) + context = self.out_lin(context) # (bs, q_length, dim) + + if output_attentions: + return (context, weights) + else: + return (context,) + + +class DistilBertFlashAttention2(MultiHeadSelfAttention): + """ + DistilBert flash attention module. This module inherits from `MultiHeadSelfAttention` as the weights of the module + stays untouched. The only required change would be on the forward pass where it needs to correctly call the public + API of flash attention and deal with padding tokens in case the input contains any of them. + """ + + # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__ + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: torch.Tensor, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + ) -> Tuple[torch.Tensor, ...]: + """ + Parameters: + query: torch.tensor(bs, seq_length, dim) + key: torch.tensor(bs, seq_length, dim) + value: torch.tensor(bs, seq_length, dim) + mask: torch.tensor(bs, seq_length) + + Returns: + weights: torch.tensor(bs, n_heads, seq_length, seq_length) Attention weights context: torch.tensor(bs, + seq_length, dim) Contextualized layer. Optional: only if `output_attentions=True` + """ + batch_size, q_length, dim = query.size() + + dim_per_head = self.dim // self.n_heads + + def reshape(x: torch.Tensor) -> torch.Tensor: + """separate heads""" + return x.view(batch_size, -1, self.n_heads, dim_per_head) + + # Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + query_states = reshape(self.q_lin(query)) + key_states = reshape(self.k_lin(key)) + value_states = reshape(self.v_lin(value)) + + attn_dropout = self.config.attention_dropout if self.training else 0.0 + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slowdown training & inference so it is recommended to not cast the LayerNorms + # in fp32. (LlamaRMSNorm handles it correctly) + + if query_states.dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_lin.weight.dtype + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + attn_weights = self._flash_attention_forward( + query_states, key_states, value_states, mask, q_length, dropout=attn_dropout + ) + + attn_weights_reshaped = attn_weights.reshape( + batch_size, q_length, self.n_heads * dim_per_head + ) + attn_output = self.out_lin(attn_weights_reshaped) + + if output_attentions: + return (attn_output, attn_weights) + else: + return (attn_output,) + + # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward with causal=True->causal=False + def _flash_attention_forward( + self, + query_states, + key_states, + value_states, + attention_mask, + query_length, + dropout=0.0, + softmax_scale=None, + ): + """ + Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token + first unpad the input, then computes the attention scores and pad the final attention scores. + + Args: + query_states (`torch.Tensor`): + Input query states to be passed to Flash Attention API + key_states (`torch.Tensor`): + Input key states to be passed to Flash Attention API + value_states (`torch.Tensor`): + Input value states to be passed to Flash Attention API + attention_mask (`torch.Tensor`): + The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the + position of padding tokens and 1 for the position of non-padding tokens. + dropout (`float`): + Attention dropout + softmax_scale (`float`, *optional*): + The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim) + """ + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.bert_padding import pad_input + + if not self._flash_attn_uses_top_left_mask: + causal = self.is_causal + else: + # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__. + causal = self.is_causal and query_length != 1 + + # Contains at least one padding token in the sequence + if attention_mask is not None: + batch_size = query_states.shape[0] + ( + query_states, + key_states, + value_states, + indices_q, + cu_seq_lens, + max_seq_lens, + ) = self._upad_input( + query_states, key_states, value_states, attention_mask, query_length + ) + + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens + + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=dropout, + softmax_scale=softmax_scale, + causal=causal, + ) + + attn_output = pad_input( + attn_output_unpad, indices_q, batch_size, query_length + ) + else: + attn_output = flash_attn_func( + query_states, + key_states, + value_states, + dropout, + softmax_scale=softmax_scale, + causal=causal, + ) + + return attn_output + + # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input with num_heads->n_heads + def _upad_input( + self, query_layer, key_layer, value_layer, attention_mask, query_length + ): + from flash_attn.bert_padding import index_first_axis, unpad_input + + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape + + key_layer = index_first_axis( + key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), + indices_k, + ) + value_layer = index_first_axis( + value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), + indices_k, + ) + if query_length == kv_seq_len: + query_layer = index_first_axis( + query_layer.reshape(batch_size * kv_seq_len, self.n_heads, head_dim), + indices_k, + ) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query_layer.device + ) # There is a memcpy here, that is very bad. + indices_q = cu_seqlens_q[:-1] + query_layer = query_layer.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -query_length:] + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input( + query_layer, attention_mask + ) + + return ( + query_layer, + key_layer, + value_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +class FFN(nn.Module): + def __init__(self, config: DistilBertConfig): + super().__init__() + self.dropout = nn.Dropout(p=config.dropout) + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.lin1 = nn.Linear(in_features=config.dim, out_features=config.hidden_dim) + self.lin2 = nn.Linear(in_features=config.hidden_dim, out_features=config.dim) + self.activation = get_activation(config.activation) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + return apply_chunking_to_forward( + self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, input + ) + + def ff_chunk(self, input: torch.Tensor) -> torch.Tensor: + x = self.lin1(input) + x = self.activation(x) + x = self.lin2(x) + x = self.dropout(x) + return x + + +DISTILBERT_ATTENTION_CLASSES = { + "eager": MultiHeadSelfAttention, + "flash_attention_2": DistilBertFlashAttention2, +} + + +class TransformerBlock(nn.Module): + def __init__(self, config: DistilBertConfig): + super().__init__() + + # Have an even number of Configure multi-heads + if config.dim % config.n_heads != 0: + raise ValueError( + f"config.n_heads {config.n_heads} must divide config.dim {config.dim} evenly" + ) + + self.attention = DISTILBERT_ATTENTION_CLASSES[config._attn_implementation]( + config + ) + self.sa_layer_norm = nn.LayerNorm(normalized_shape=config.dim, eps=1e-12) + + self.ffn = FFN(config) + self.output_layer_norm = nn.LayerNorm(normalized_shape=config.dim, eps=1e-12) + + def forward( + self, + x: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + ) -> Tuple[torch.Tensor, ...]: + """ + Parameters: + x: torch.tensor(bs, seq_length, dim) + attn_mask: torch.tensor(bs, seq_length) + + Returns: + sa_weights: torch.tensor(bs, n_heads, seq_length, seq_length) The attention weights ffn_output: + torch.tensor(bs, seq_length, dim) The output of the transformer block contextualization. + """ + # Self-Attention + sa_output = self.attention( + query=x, + key=x, + value=x, + mask=attn_mask, + head_mask=head_mask, + output_attentions=output_attentions, + ) + if output_attentions: + sa_output, sa_weights = ( + sa_output # (bs, seq_length, dim), (bs, n_heads, seq_length, seq_length) + ) + else: # To handle these `output_attentions` or `output_hidden_states` cases returning tuples + sa_output = sa_output[0] + + sa_output = self.sa_layer_norm(sa_output + x) # (bs, seq_length, dim) + + # Feed Forward Network + ffn_output = self.ffn(sa_output) # (bs, seq_length, dim) + ffn_output: torch.Tensor = self.output_layer_norm( + ffn_output + sa_output + ) # (bs, seq_length, dim) + + output = (ffn_output,) + if output_attentions: + output = (sa_weights,) + output + return output + + +class Transformer(nn.Module): + def __init__(self, config: DistilBertConfig): + super().__init__() + self.n_layers = config.n_layers + self.layer = nn.ModuleList( + [TransformerBlock(config) for _ in range(config.n_layers)] + ) + self.gradient_checkpointing = False + + def forward( + self, + x: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: Optional[bool] = None, + ) -> Union[BaseModelOutput, Tuple[torch.Tensor, ...]]: # docstyle-ignore + """ + Parameters: + x: torch.tensor(bs, seq_length, dim) Input sequence embedded. + attn_mask: torch.tensor(bs, seq_length) Attention mask on the sequence. + + Returns: + hidden_state: torch.tensor(bs, seq_length, dim) Sequence of hidden states in the last (top) + layer all_hidden_states: Tuple[torch.tensor(bs, seq_length, dim)] + Tuple of length n_layers with the hidden states from each layer. + Optional: only if output_hidden_states=True + all_attentions: Tuple[torch.tensor(bs, n_heads, seq_length, seq_length)] + Tuple of length n_layers with the attention weights from each layer + Optional: only if output_attentions=True + """ + all_hidden_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + hidden_state = x + for i, layer_module in enumerate(self.layer): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_state,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + layer_module.__call__, + hidden_state, + attn_mask, + head_mask[i], + output_attentions, + ) + else: + layer_outputs = layer_module( + hidden_state, + attn_mask, + head_mask[i], + output_attentions, + ) + + hidden_state = layer_outputs[-1] + + if output_attentions: + if len(layer_outputs) != 2: + raise ValueError( + f"The length of the layer_outputs should be 2, but it is {len(layer_outputs)}" + ) + + attentions = layer_outputs[0] + all_attentions = all_attentions + (attentions,) + else: + if len(layer_outputs) != 1: + raise ValueError( + f"The length of the layer_outputs should be 1, but it is {len(layer_outputs)}" + ) + + # Add last layer + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_state,) + + if not return_dict: + return tuple( + v + for v in [hidden_state, all_hidden_states, all_attentions] + if v is not None + ) + return BaseModelOutput( + last_hidden_state=hidden_state, + hidden_states=all_hidden_states, + attentions=all_attentions, + ) + + +class DistilBertPreTrainedModel(SuryaPreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = DistilBertConfig + load_tf_weights = None + base_model_prefix = "distilbert" + supports_gradient_checkpointing = True + _supports_flash_attn_2 = True + + def _init_weights(self, module: nn.Module): + """Initialize the weights.""" + if isinstance(module, nn.Linear): + # 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) + elif isinstance(module, Embeddings) and self.config.sinusoidal_pos_embds: + create_sinusoidal_embeddings( + self.config.max_position_embeddings, + self.config.dim, + module.position_embeddings.weight, + ) + + +class DistilBertModel(DistilBertPreTrainedModel): + def __init__(self, config: DistilBertConfig): + super().__init__(config) + + self.embeddings = Embeddings(config) # Embeddings + self.transformer = Transformer(config) # Encoder + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + + # Initialize weights and apply final processing + self.post_init() + + def get_position_embeddings(self) -> nn.Embedding: + """ + Returns the position embeddings + """ + return self.embeddings.position_embeddings + + def resize_position_embeddings(self, new_num_position_embeddings: int): + """ + Resizes position embeddings of the model if `new_num_position_embeddings != config.max_position_embeddings`. + + Arguments: + new_num_position_embeddings (`int`): + The number of new position embedding matrix. If position embeddings are learned, increasing the size + will add newly initialized vectors at the end, whereas reducing the size will remove vectors from the + end. If position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the + size will add correct vectors at the end following the position encoding algorithm, whereas reducing + the size will remove vectors from the end. + """ + num_position_embeds_diff = ( + new_num_position_embeddings - self.config.max_position_embeddings + ) + + # no resizing needs to be done if the length stays the same + if num_position_embeds_diff == 0: + return + + self.config.max_position_embeddings = new_num_position_embeddings + + old_position_embeddings_weight = ( + self.embeddings.position_embeddings.weight.clone() + ) + + self.embeddings.position_embeddings = nn.Embedding( + self.config.max_position_embeddings, self.config.dim + ) + + if self.config.sinusoidal_pos_embds: + create_sinusoidal_embeddings( + n_pos=self.config.max_position_embeddings, + dim=self.config.dim, + out=self.position_embeddings.weight, + ) + else: + with torch.no_grad(): + if num_position_embeds_diff > 0: + self.embeddings.position_embeddings.weight[ + :-num_position_embeds_diff + ] = nn.Parameter(old_position_embeddings_weight) + else: + self.embeddings.position_embeddings.weight = nn.Parameter( + old_position_embeddings_weight[:num_position_embeds_diff] + ) + # move position_embeddings to correct device + self.embeddings.position_embeddings.to(self.device) + + def get_input_embeddings(self) -> nn.Embedding: + return self.embeddings.word_embeddings + + def set_input_embeddings(self, new_embeddings: nn.Embedding): + self.embeddings.word_embeddings = new_embeddings + + def _prune_heads(self, heads_to_prune: Dict[int, List[List[int]]]): + """ + Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base + class PreTrainedModel + """ + for layer, heads in heads_to_prune.items(): + self.transformer.layer[layer].attention.prune_heads(heads) + + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[BaseModelOutput, Tuple[torch.Tensor, ...]]: + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if input_ids is not None and inputs_embeds is not None: + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time" + ) + elif input_ids is not None: + self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask) + input_shape = input_ids.size() + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + device = input_ids.device if input_ids is not None else inputs_embeds.device + + # Prepare head mask if needed + head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) + + embeddings = self.embeddings(input_ids, inputs_embeds) # (bs, seq_length, dim) + + if self._use_flash_attention_2: + attention_mask = ( + attention_mask + if (attention_mask is not None and 0 in attention_mask) + else None + ) + else: + if attention_mask is None: + attention_mask = torch.ones( + input_shape, device=device + ) # (bs, seq_length) + + return self.transformer( + x=embeddings, + attn_mask=attention_mask, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + +class DistilBertForSequenceClassification(S3DownloaderMixin, DistilBertPreTrainedModel): + def __init__(self, config: DistilBertConfig, **kwargs): + super().__init__(config, **kwargs) + self.num_labels = config.num_labels + self.config = config + + self.distilbert = DistilBertModel(config) + self.pre_classifier = nn.Linear(config.dim, config.dim) + self.classifier = nn.Linear(config.dim, config.num_labels) + self.dropout = nn.Dropout(config.seq_classif_dropout) + + # Initialize weights and apply final processing + self.post_init() + + def get_position_embeddings(self) -> nn.Embedding: + """ + Returns the position embeddings + """ + return self.distilbert.get_position_embeddings() + + def resize_position_embeddings(self, new_num_position_embeddings: int): + """ + Resizes position embeddings of the model if `new_num_position_embeddings != config.max_position_embeddings`. + + Arguments: + new_num_position_embeddings (`int`): + The number of new position embedding matrix. If position embeddings are learned, increasing the size + will add newly initialized vectors at the end, whereas reducing the size will remove vectors from the + end. If position embeddings are not learned (*e.g.* sinusoidal position embeddings), increasing the + size will add correct vectors at the end following the position encoding algorithm, whereas reducing + the size will remove vectors from the end. + """ + self.distilbert.resize_position_embeddings(new_num_position_embeddings) + + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[SequenceClassifierOutput, Tuple[torch.Tensor, ...]]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + distilbert_output = self.distilbert( + input_ids=input_ids, + attention_mask=attention_mask, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_state = distilbert_output[0] # (bs, seq_len, dim) + pooled_output = hidden_state[:, 0] # (bs, dim) + pooled_output = self.pre_classifier(pooled_output) # (bs, dim) + pooled_output = nn.ReLU()(pooled_output) # (bs, dim) + pooled_output = self.dropout(pooled_output) # (bs, dim) + logits = self.classifier(pooled_output) # (bs, num_labels) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + if not return_dict: + output = (logits,) + distilbert_output[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=distilbert_output.hidden_states, + attentions=distilbert_output.attentions, + ) diff --git a/surya/ocr_error/schema.py b/surya/ocr_error/schema.py new file mode 100644 index 0000000..564063d --- /dev/null +++ b/surya/ocr_error/schema.py @@ -0,0 +1,8 @@ +from typing import List + +from pydantic import BaseModel + + +class OCRErrorDetectionResult(BaseModel): + texts: List[str] + labels: List[str] diff --git a/surya/ocr_error/tokenizer.py b/surya/ocr_error/tokenizer.py new file mode 100644 index 0000000..7212bcf --- /dev/null +++ b/surya/ocr_error/tokenizer.py @@ -0,0 +1,525 @@ +import collections +import os +import unicodedata +from typing import List, Optional, Tuple + +from transformers.tokenization_utils import ( + PreTrainedTokenizer, + _is_control, + _is_punctuation, + _is_whitespace, +) + +from surya.common.s3 import S3DownloaderMixin + +VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"} + + +# Copied from transformers.models.bert.tokenization_bert.load_vocab +def load_vocab(vocab_file): + """Loads a vocabulary file into a dictionary.""" + vocab = collections.OrderedDict() + with open(vocab_file, "r", encoding="utf-8") as reader: + tokens = reader.readlines() + for index, token in enumerate(tokens): + token = token.rstrip("\n") + vocab[token] = index + return vocab + + +# Copied from transformers.models.bert.tokenization_bert.whitespace_tokenize +def whitespace_tokenize(text): + """Runs basic whitespace cleaning and splitting on a piece of text.""" + text = text.strip() + if not text: + return [] + tokens = text.split() + return tokens + + +class DistilBertTokenizer(S3DownloaderMixin, PreTrainedTokenizer): + r""" + Construct a DistilBERT tokenizer. Based on WordPiece. + + This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to + this superclass for more information regarding those methods. + + Args: + vocab_file (`str`): + File containing the vocabulary. + do_lower_case (`bool`, *optional*, defaults to `True`): + Whether or not to lowercase the input when tokenizing. + do_basic_tokenize (`bool`, *optional*, defaults to `True`): + Whether or not to do basic tokenization before WordPiece. + never_split (`Iterable`, *optional*): + Collection of tokens which will never be split during tokenization. Only has an effect when + `do_basic_tokenize=True` + unk_token (`str`, *optional*, defaults to `"[UNK]"`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + sep_token (`str`, *optional*, defaults to `"[SEP]"`): + The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for + sequence classification or for a text and a question for question answering. It is also used as the last + token of a sequence built with special tokens. + pad_token (`str`, *optional*, defaults to `"[PAD]"`): + The token used for padding, for example when batching sequences of different lengths. + cls_token (`str`, *optional*, defaults to `"[CLS]"`): + The classifier token which is used when doing sequence classification (classification of the whole sequence + instead of per-token classification). It is the first token of the sequence when built with special tokens. + mask_token (`str`, *optional*, defaults to `"[MASK]"`): + The token used for masking values. This is the token used when training this model with masked language + modeling. This is the token which the model will try to predict. + tokenize_chinese_chars (`bool`, *optional*, defaults to `True`): + Whether or not to tokenize Chinese characters. + + This should likely be deactivated for Japanese (see this + [issue](https://github.com/huggingface/transformers/issues/328)). + strip_accents (`bool`, *optional*): + Whether or not to strip all accents. If this option is not specified, then it will be determined by the + value for `lowercase` (as in the original BERT). + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + + def __init__( + self, + vocab_file, + do_lower_case=True, + do_basic_tokenize=True, + never_split=None, + unk_token="[UNK]", + sep_token="[SEP]", + pad_token="[PAD]", + cls_token="[CLS]", + mask_token="[MASK]", + tokenize_chinese_chars=True, + strip_accents=None, + **kwargs, + ): + if not os.path.isfile(vocab_file): + raise ValueError( + f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained" + " model use `tokenizer = DistilBertTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`" + ) + self.vocab = load_vocab(vocab_file) + self.ids_to_tokens = collections.OrderedDict( + [(ids, tok) for tok, ids in self.vocab.items()] + ) + self.do_basic_tokenize = do_basic_tokenize + if do_basic_tokenize: + self.basic_tokenizer = BasicTokenizer( + do_lower_case=do_lower_case, + never_split=never_split, + tokenize_chinese_chars=tokenize_chinese_chars, + strip_accents=strip_accents, + ) + self.wordpiece_tokenizer = WordpieceTokenizer( + vocab=self.vocab, unk_token=str(unk_token) + ) + + super().__init__( + do_lower_case=do_lower_case, + do_basic_tokenize=do_basic_tokenize, + never_split=never_split, + unk_token=unk_token, + sep_token=sep_token, + pad_token=pad_token, + cls_token=cls_token, + mask_token=mask_token, + tokenize_chinese_chars=tokenize_chinese_chars, + strip_accents=strip_accents, + **kwargs, + ) + + @property + # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.do_lower_case + def do_lower_case(self): + return self.basic_tokenizer.do_lower_case + + @property + # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.vocab_size + def vocab_size(self): + return len(self.vocab) + + # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.get_vocab + def get_vocab(self): + return dict(self.vocab, **self.added_tokens_encoder) + + # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._tokenize + def _tokenize(self, text, split_special_tokens=False): + split_tokens = [] + if self.do_basic_tokenize: + for token in self.basic_tokenizer.tokenize( + text, + never_split=self.all_special_tokens + if not split_special_tokens + else None, + ): + # If the token is part of the never_split set + if token in self.basic_tokenizer.never_split: + split_tokens.append(token) + else: + split_tokens += self.wordpiece_tokenizer.tokenize(token) + else: + split_tokens = self.wordpiece_tokenizer.tokenize(text) + return split_tokens + + # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._convert_token_to_id + def _convert_token_to_id(self, token): + """Converts a token (str) in an id using the vocab.""" + return self.vocab.get(token, self.vocab.get(self.unk_token)) + + # Copied from transformers.models.bert.tokenization_bert.BertTokenizer._convert_id_to_token + def _convert_id_to_token(self, index): + """Converts an index (integer) in a token (str) using the vocab.""" + return self.ids_to_tokens.get(index, self.unk_token) + + # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.convert_tokens_to_string + def convert_tokens_to_string(self, tokens): + """Converts a sequence of tokens (string) in a single string.""" + out_string = " ".join(tokens).replace(" ##", "").strip() + return out_string + + # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.build_inputs_with_special_tokens + def build_inputs_with_special_tokens( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None + ) -> List[int]: + """ + Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and + adding special tokens. A BERT sequence has the following format: + + - single sequence: `[CLS] X [SEP]` + - pair of sequences: `[CLS] A [SEP] B [SEP]` + + Args: + token_ids_0 (`List[int]`): + List of IDs to which the special tokens will be added. + token_ids_1 (`List[int]`, *optional*): + Optional second list of IDs for sequence pairs. + + Returns: + `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens. + """ + if token_ids_1 is None: + return [self.cls_token_id] + token_ids_0 + [self.sep_token_id] + cls = [self.cls_token_id] + sep = [self.sep_token_id] + return cls + token_ids_0 + sep + token_ids_1 + sep + + # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.get_special_tokens_mask + def get_special_tokens_mask( + self, + token_ids_0: List[int], + token_ids_1: Optional[List[int]] = None, + already_has_special_tokens: bool = False, + ) -> List[int]: + """ + Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding + special tokens using the tokenizer `prepare_for_model` method. + + Args: + token_ids_0 (`List[int]`): + List of IDs. + token_ids_1 (`List[int]`, *optional*): + Optional second list of IDs for sequence pairs. + already_has_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the token list is already formatted with special tokens for the model. + + Returns: + `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. + """ + + if already_has_special_tokens: + return super().get_special_tokens_mask( + token_ids_0=token_ids_0, + token_ids_1=token_ids_1, + already_has_special_tokens=True, + ) + + if token_ids_1 is not None: + return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1] + return [1] + ([0] * len(token_ids_0)) + [1] + + # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.create_token_type_ids_from_sequences + def create_token_type_ids_from_sequences( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None + ) -> List[int]: + """ + Create a mask from the two sequences passed to be used in a sequence-pair classification task. A BERT sequence + pair mask has the following format: + + ``` + 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 + | first sequence | second sequence | + ``` + + If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s). + + Args: + token_ids_0 (`List[int]`): + List of IDs. + token_ids_1 (`List[int]`, *optional*): + Optional second list of IDs for sequence pairs. + + Returns: + `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s). + """ + sep = [self.sep_token_id] + cls = [self.cls_token_id] + if token_ids_1 is None: + return len(cls + token_ids_0 + sep) * [0] + return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1] + + # Copied from transformers.models.bert.tokenization_bert.BertTokenizer.save_vocabulary + def save_vocabulary( + self, save_directory: str, filename_prefix: Optional[str] = None + ) -> Tuple[str]: + index = 0 + if os.path.isdir(save_directory): + vocab_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + VOCAB_FILES_NAMES["vocab_file"], + ) + else: + vocab_file = ( + filename_prefix + "-" if filename_prefix else "" + ) + save_directory + with open(vocab_file, "w", encoding="utf-8") as writer: + for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]): + if index != token_index: + # logger.warning( + # f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive." + # " Please check that the vocabulary is not corrupted!" + # ) + index = token_index + writer.write(token + "\n") + index += 1 + return (vocab_file,) + + +# Copied from transformers.models.bert.tokenization_bert.BasicTokenizer +class BasicTokenizer(object): + """ + Constructs a BasicTokenizer that will run basic tokenization (punctuation splitting, lower casing, etc.). + + Args: + do_lower_case (`bool`, *optional*, defaults to `True`): + Whether or not to lowercase the input when tokenizing. + never_split (`Iterable`, *optional*): + Collection of tokens which will never be split during tokenization. Only has an effect when + `do_basic_tokenize=True` + tokenize_chinese_chars (`bool`, *optional*, defaults to `True`): + Whether or not to tokenize Chinese characters. + + This should likely be deactivated for Japanese (see this + [issue](https://github.com/huggingface/transformers/issues/328)). + strip_accents (`bool`, *optional*): + Whether or not to strip all accents. If this option is not specified, then it will be determined by the + value for `lowercase` (as in the original BERT). + do_split_on_punc (`bool`, *optional*, defaults to `True`): + In some instances we want to skip the basic punctuation splitting so that later tokenization can capture + the full context of the words, such as contractions. + """ + + def __init__( + self, + do_lower_case=True, + never_split=None, + tokenize_chinese_chars=True, + strip_accents=None, + do_split_on_punc=True, + ): + if never_split is None: + never_split = [] + self.do_lower_case = do_lower_case + self.never_split = set(never_split) + self.tokenize_chinese_chars = tokenize_chinese_chars + self.strip_accents = strip_accents + self.do_split_on_punc = do_split_on_punc + + def tokenize(self, text, never_split=None): + """ + Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer. + + Args: + never_split (`List[str]`, *optional*) + Kept for backward compatibility purposes. Now implemented directly at the base class level (see + [`PreTrainedTokenizer.tokenize`]) List of token not to split. + """ + # union() returns a new set by concatenating the two sets. + never_split = ( + self.never_split.union(set(never_split)) + if never_split + else self.never_split + ) + text = self._clean_text(text) + + # This was added on November 1st, 2018 for the multilingual and Chinese + # models. This is also applied to the English models now, but it doesn't + # matter since the English models were not trained on any Chinese data + # and generally don't have any Chinese data in them (there are Chinese + # characters in the vocabulary because Wikipedia does have some Chinese + # words in the English Wikipedia.). + if self.tokenize_chinese_chars: + text = self._tokenize_chinese_chars(text) + # prevents treating the same character with different unicode codepoints as different characters + unicode_normalized_text = unicodedata.normalize("NFC", text) + orig_tokens = whitespace_tokenize(unicode_normalized_text) + split_tokens = [] + for token in orig_tokens: + if token not in never_split: + if self.do_lower_case: + token = token.lower() + if self.strip_accents is not False: + token = self._run_strip_accents(token) + elif self.strip_accents: + token = self._run_strip_accents(token) + split_tokens.extend(self._run_split_on_punc(token, never_split)) + + output_tokens = whitespace_tokenize(" ".join(split_tokens)) + return output_tokens + + def _run_strip_accents(self, text): + """Strips accents from a piece of text.""" + text = unicodedata.normalize("NFD", text) + output = [] + for char in text: + cat = unicodedata.category(char) + if cat == "Mn": + continue + output.append(char) + return "".join(output) + + def _run_split_on_punc(self, text, never_split=None): + """Splits punctuation on a piece of text.""" + if not self.do_split_on_punc or ( + never_split is not None and text in never_split + ): + return [text] + chars = list(text) + i = 0 + start_new_word = True + output = [] + while i < len(chars): + char = chars[i] + if _is_punctuation(char): + output.append([char]) + start_new_word = True + else: + if start_new_word: + output.append([]) + start_new_word = False + output[-1].append(char) + i += 1 + + return ["".join(x) for x in output] + + def _tokenize_chinese_chars(self, text): + """Adds whitespace around any CJK character.""" + output = [] + for char in text: + cp = ord(char) + if self._is_chinese_char(cp): + output.append(" ") + output.append(char) + output.append(" ") + else: + output.append(char) + return "".join(output) + + def _is_chinese_char(self, cp): + """Checks whether CP is the codepoint of a CJK character.""" + # This defines a "chinese character" as anything in the CJK Unicode block: + # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block) + # + # Note that the CJK Unicode block is NOT all Japanese and Korean characters, + # despite its name. The modern Korean Hangul alphabet is a different block, + # as is Japanese Hiragana and Katakana. Those alphabets are used to write + # space-separated words, so they are not treated specially and handled + # like the all of the other languages. + if ( + (cp >= 0x4E00 and cp <= 0x9FFF) + or (cp >= 0x3400 and cp <= 0x4DBF) # + or (cp >= 0x20000 and cp <= 0x2A6DF) # + or (cp >= 0x2A700 and cp <= 0x2B73F) # + or (cp >= 0x2B740 and cp <= 0x2B81F) # + or (cp >= 0x2B820 and cp <= 0x2CEAF) # + or (cp >= 0xF900 and cp <= 0xFAFF) + or (cp >= 0x2F800 and cp <= 0x2FA1F) # + ): # + return True + + return False + + def _clean_text(self, text): + """Performs invalid character removal and whitespace cleanup on text.""" + output = [] + for char in text: + cp = ord(char) + if cp == 0 or cp == 0xFFFD or _is_control(char): + continue + if _is_whitespace(char): + output.append(" ") + else: + output.append(char) + return "".join(output) + + +# Copied from transformers.models.bert.tokenization_bert.WordpieceTokenizer +class WordpieceTokenizer(object): + """Runs WordPiece tokenization.""" + + def __init__(self, vocab, unk_token, max_input_chars_per_word=100): + self.vocab = vocab + self.unk_token = unk_token + self.max_input_chars_per_word = max_input_chars_per_word + + def tokenize(self, text): + """ + Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform + tokenization using the given vocabulary. + + For example, `input = "unaffable"` wil return as output `["un", "##aff", "##able"]`. + + Args: + text: A single token or whitespace separated tokens. This should have + already been passed through *BasicTokenizer*. + + Returns: + A list of wordpiece tokens. + """ + + output_tokens = [] + for token in whitespace_tokenize(text): + chars = list(token) + if len(chars) > self.max_input_chars_per_word: + output_tokens.append(self.unk_token) + continue + + is_bad = False + start = 0 + sub_tokens = [] + while start < len(chars): + end = len(chars) + cur_substr = None + while start < end: + substr = "".join(chars[start:end]) + if start > 0: + substr = "##" + substr + if substr in self.vocab: + cur_substr = substr + break + end -= 1 + if cur_substr is None: + is_bad = True + break + sub_tokens.append(cur_substr) + start = end + + if is_bad: + output_tokens.append(self.unk_token) + else: + output_tokens.extend(sub_tokens) + return output_tokens diff --git a/surya/recognition/__init__.py b/surya/recognition/__init__.py new file mode 100644 index 0000000..90aff61 --- /dev/null +++ b/surya/recognition/__init__.py @@ -0,0 +1,404 @@ +"""RecognitionPredictor: per-block OCR via BLOCK_PROMPT. + +Given page images and corresponding LayoutResult (or any list of LayoutBox), +crops each block, runs BLOCK_PROMPT, returns PageOCRResult per page. +""" + +from __future__ import annotations + +from typing import List, Optional + +from PIL import Image + +from surya.common.blank import is_blank_region +from surya.inference import SuryaInferenceManager, get_default_manager +from surya.inference.parsers import clean_block_html, parse_full_page_html +from surya.inference.prompts import ( + PROMPT_TYPE_BLOCK, + PROMPT_TYPE_HIGH_ACCURACY_BBOX, + SKIP_OCR_LABELS, +) +from surya.inference.schema import BatchInputItem +from surya.inference.util import image_token_budget +from surya.layout.label import LAYOUT_PRED_RELABEL, TEXT_LABELS +from surya.layout.schema import LayoutResult +from surya.logging import get_logger +from surya.recognition.schema import ( + BlockOCRResult, + PageOCRResult, +) +from surya.settings import settings +from surya.timing import timing_span + +logger = get_logger() + + +# Surya's canonical labels we shouldn't OCR (mirrors model-emitted SKIP_OCR_LABELS +# after canonicalization). +SKIP_CANON_LABELS = {LAYOUT_PRED_RELABEL.get(lbl, lbl) for lbl in SKIP_OCR_LABELS} + + +def _crop_block(image: Image.Image, polygon, pad: int = 4) -> Image.Image: + xs = [p[0] for p in polygon] + ys = [p[1] for p in polygon] + x0 = max(0, int(min(xs)) - pad) + y0 = max(0, int(min(ys)) - pad) + x1 = min(image.size[0], int(max(xs)) + pad) + y1 = min(image.size[1], int(max(ys)) + pad) + if x1 <= x0 or y1 <= y0: + return image.crop((0, 0, 1, 1)) + return image.crop((x0, y0, x1, y1)) + + +def _drop_blank_text_blocks( + image: Image.Image, + blocks: List[BlockOCRResult], +) -> List[BlockOCRResult]: + """Drop text-labeled blocks whose source page region is essentially blank. + + Full-page OCR can emit text divs for regions that are visually empty + (margins, gutter space) — the model hallucinates a paragraph where there + is none. We crop the region, count near-white pixels, and drop the block + when the fraction exceeds ``blank_pixel_fraction``. Only text-like labels + (see ``TEXT_LABELS``) are eligible: tables, forms, equations, and visual + blocks may legitimately contain large whitespace and are left untouched. + """ + kept: List[BlockOCRResult] = [] + dropped = 0 + for blk in blocks: + if blk.label not in TEXT_LABELS or blk.skipped or blk.error: + kept.append(blk) + continue + crop = _crop_block(image, blk.polygon) + if not is_blank_region(crop): + kept.append(blk) + continue + dropped += 1 + if dropped: + logger.info(f"dropped {dropped} blank text block(s) from full-page OCR") + return kept + + +def _detect_repeat_loop( + text: str, + base_max_repeats: int = 4, + window_size: int = 500, + scaling_factor: float = 3.0, +) -> bool: + """True iff the tail of ``text`` ends in a repeating sequence. + + Ported from chandra's detect_repeat_token. For each candidate length + 1..window_size/2, takes that many trailing chars and counts consecutive + identical preceding blocks. Shorter loops need many repeats to count; + longer ones only need a few. Catches the typical decoder failure mode + where a page output gets stuck emitting the same div / phrase until it + hits max_tokens. + """ + if not text: + return False + for seq_len in range(1, window_size // 2 + 1): + candidate = text[-seq_len:] + max_repeats = int(base_max_repeats * (1 + scaling_factor / seq_len)) + repeats = 0 + pos = len(text) - seq_len + while pos >= 0 and text[pos : pos + seq_len] == candidate: + repeats += 1 + pos -= seq_len + if repeats > max_repeats: + return True + return False + + +class RecognitionPredictor: + """Per-block OCR. Construct with a SuryaInferenceManager (or rely on default).""" + + def __init__(self, manager: Optional[SuryaInferenceManager] = None): + self.manager = manager + self._disable_tqdm = settings.DISABLE_TQDM + + @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 to(self, *args, **kwargs): + return + + def __call__( + self, + images: List[Image.Image], + layout_results: Optional[List[LayoutResult]] = None, + *, + full_page: Optional[bool] = None, + ) -> List[PageOCRResult]: + """Run OCR on each page. + + Mode resolution: + - ``full_page=None`` (default): block mode if ``layout_results`` is + given, else full-page mode. This is the most-do-what-I-mean form. + - ``full_page=True``: full-page OCR (single HIGH_ACCURACY_BBOX_PROMPT + request per page). ``layout_results`` is ignored — a warning is + logged if it was supplied. + - ``full_page=False``: block mode (per-layout-block OCR request). + ``layout_results`` is required. + + Full-page is the more accurate path; block mode is for callers that + specifically need per-block crops (e.g. for downstream merging with + text-line detection). + """ + if not images: + return [] + if full_page is None: + full_page = layout_results is None + if full_page: + if layout_results is not None: + logger.info( + "RecognitionPredictor called with full_page=True and " + "layout_results; layout will be used as fallback if the " + "full-page output devolves into a repetition loop." + ) + return self._full_page_ocr(images, fallback_layout=layout_results) + if layout_results is None: + raise ValueError("layout_results required when full_page=False") + if len(images) != len(layout_results): + raise ValueError( + f"images and layout_results must be same length " + f"({len(images)} vs {len(layout_results)})" + ) + manager = self.manager or get_default_manager() + + # Build a flat batch across all pages for max concurrency + batch: List[BatchInputItem] = [] + block_index_map: List[tuple[int, int]] = [] # (page_idx, block_idx) + skipped_flags: List[bool] = [] + + with timing_span("recognition_build_block_batch", page_count=len(images)): + for page_idx, (img, layout) in enumerate(zip(images, layout_results)): + page_boxes = sorted(layout.bboxes, key=lambda b: (b.position, b.bbox[1], b.bbox[0])) + page_boxes = page_boxes[: settings.SURYA_MAX_BLOCKS_PER_PAGE] + if len(layout.bboxes) > len(page_boxes): + logger.info( + f"capped OCR blocks for page {page_idx}: " + f"{len(layout.bboxes)} -> {len(page_boxes)}" + ) + for block_idx, box in enumerate(page_boxes): + skip = box.label in SKIP_CANON_LABELS + skipped_flags.append(skip) + if skip: + continue + crop = _crop_block(img, box.polygon) + max_tokens = image_token_budget( + box.count, ceiling=settings.SURYA_MAX_TOKENS_BLOCK_CEILING + ) + batch.append( + BatchInputItem( + image=crop, + prompt_type=PROMPT_TYPE_BLOCK, + max_tokens=max_tokens, + metadata={"page_idx": page_idx, "block_idx": block_idx}, + ) + ) + block_index_map.append((page_idx, block_idx)) + + with timing_span( + "recognition_manager_generate", + item_count=len(batch), + skipped_count=sum(1 for flag in skipped_flags if flag), + max_tokens_sum=sum(item.max_tokens or 0 for item in batch), + ): + outputs = manager.generate(batch) if batch else [] + + # Index outputs by (page_idx, block_idx) + with timing_span( + "recognition_assemble_pages", + output_count=len(outputs), + token_count=sum(out.token_count or 0 for out in outputs), + ): + out_by_key = {} + for out in outputs: + key = (out.metadata["page_idx"], out.metadata["block_idx"]) + out_by_key[key] = out + + # Assemble PageOCRResult per page + results: List[PageOCRResult] = [] + for page_idx, (img, layout) in enumerate(zip(images, layout_results)): + w, h = img.size + blocks: List[BlockOCRResult] = [] + page_boxes = sorted(layout.bboxes, key=lambda b: (b.position, b.bbox[1], b.bbox[0])) + page_boxes = page_boxes[: settings.SURYA_MAX_BLOCKS_PER_PAGE] + for block_idx, box in enumerate(page_boxes): + skip = box.label in SKIP_CANON_LABELS + if skip: + blocks.append( + BlockOCRResult( + polygon=box.polygon, + label=box.label, + raw_label=box.raw_label, + reading_order=box.position, + html="", + skipped=True, + confidence=1.0, + ) + ) + continue + out = out_by_key.get((page_idx, block_idx)) + if out is None or out.error: + blocks.append( + BlockOCRResult( + polygon=box.polygon, + label=box.label, + raw_label=box.raw_label, + reading_order=box.position, + html="", + skipped=False, + error=True, + confidence=0.0, + ) + ) + continue + html = clean_block_html(out.raw) + conf = out.mean_token_prob if out.mean_token_prob is not None else 1.0 + blocks.append( + BlockOCRResult( + polygon=box.polygon, + label=box.label, + raw_label=box.raw_label, + reading_order=box.position, + html=html, + skipped=False, + error=False, + confidence=conf, + raw_logprobs=out.logprobs, + ) + ) + results.append( + PageOCRResult(blocks=blocks, image_bbox=[0, 0, float(w), float(h)]) + ) + return results + + def _full_page_ocr( + self, + images: List[Image.Image], + fallback_layout: Optional[List[LayoutResult]] = None, + ) -> List[PageOCRResult]: + """One HIGH_ACCURACY_BBOX_PROMPT request per page; parses divs into blocks. + + On per-page failure (parse error, empty output, or a detected + repetition loop in the decoder output), falls back to layout + + block-mode OCR for that page only. ``fallback_layout``, if given, + provides per-page LayoutResults to use on fallback; otherwise the + LayoutPredictor is invoked lazily for just the affected pages. + """ + manager = self.manager or get_default_manager() + with timing_span("recognition_build_full_page_batch", page_count=len(images)): + batch = [ + BatchInputItem( + image=img, + prompt_type=PROMPT_TYPE_HIGH_ACCURACY_BBOX, + max_tokens=settings.SURYA_MAX_TOKENS_FULL_PAGE, + metadata={"page_idx": i}, + ) + for i, img in enumerate(images) + ] + with timing_span( + "recognition_full_page_generate", + item_count=len(batch), + max_tokens=settings.SURYA_MAX_TOKENS_FULL_PAGE, + ): + outputs = manager.generate(batch) + out_by_page = {o.metadata["page_idx"]: o for o in outputs} + + results: List[Optional[PageOCRResult]] = [None] * len(images) + needs_fallback: List[int] = [] + for page_idx, img in enumerate(images): + w, h = img.size + page_bbox = [0, 0, float(w), float(h)] + out = out_by_page.get(page_idx) + if out is None or out.error: + # Hard failure (request lost / server error). Always fallback. + needs_fallback.append(page_idx) + continue + if not out.raw: + # Empty model output. If the page is genuinely blank, the + # model is correct — return an empty result. Only fall back + # when the page has content the model failed to emit. + if is_blank_region(img): + results[page_idx] = PageOCRResult(blocks=[], image_bbox=page_bbox) + else: + logger.info( + f"empty full-page output for non-blank page {page_idx}; " + f"falling back to layout + block OCR" + ) + needs_fallback.append(page_idx) + continue + if _detect_repeat_loop(out.raw): + logger.info( + f"full-page output for page {page_idx} appears to loop; " + f"falling back to layout + block OCR" + ) + needs_fallback.append(page_idx) + continue + try: + parsed = parse_full_page_html(out.raw) + except Exception as e: + logger.warning( + f"Full-page parse failed for page {page_idx}: {e}; " + f"falling back to layout + block OCR" + ) + needs_fallback.append(page_idx) + continue + confidence = out.mean_token_prob if out.mean_token_prob is not None else 1.0 + blocks: List[BlockOCRResult] = [] + for idx, item in enumerate(parsed): + x0 = item.bbox[0] / settings.BBOX_SCALE * w + y0 = item.bbox[1] / settings.BBOX_SCALE * h + x1 = item.bbox[2] / settings.BBOX_SCALE * w + y1 = item.bbox[3] / settings.BBOX_SCALE * h + polygon = [[x0, y0], [x1, y0], [x1, y1], [x0, y1]] + canon = LAYOUT_PRED_RELABEL.get(item.label, item.label) + skipped = canon in SKIP_CANON_LABELS + blocks.append( + BlockOCRResult( + polygon=polygon, + label=canon, + raw_label=item.label, + reading_order=idx, + html="" if skipped else item.html, + skipped=skipped, + error=False, + confidence=confidence, + ) + ) + blocks = _drop_blank_text_blocks(img, blocks) + results[page_idx] = PageOCRResult(blocks=blocks, image_bbox=page_bbox) + + # Block-mode fallback for any pages whose full-page output failed or looped. + if needs_fallback: + fb_images = [images[i] for i in needs_fallback] + if fallback_layout is not None: + fb_layouts = [fallback_layout[i] for i in needs_fallback] + else: + # Lazy import to avoid the surya.layout ↔ surya.recognition cycle. + from surya.layout import LayoutPredictor + + logger.info( + f"running layout for {len(fb_images)} page(s) requiring " + f"block-mode fallback" + ) + fb_layouts = LayoutPredictor(self.manager)(fb_images) + fb_results = self.__call__(fb_images, fb_layouts, full_page=False) + for fb_idx, page_idx in enumerate(needs_fallback): + results[page_idx] = fb_results[fb_idx] + + # Backfill any still-None pages with empty results (defensive — shouldn't happen). + out_results: List[PageOCRResult] = [] + for page_idx, img in enumerate(images): + r = results[page_idx] + if r is None: + w, h = img.size + r = PageOCRResult(blocks=[], image_bbox=[0, 0, float(w), float(h)]) + out_results.append(r) + return out_results diff --git a/surya/recognition/schema.py b/surya/recognition/schema.py new file mode 100644 index 0000000..e914b5b --- /dev/null +++ b/surya/recognition/schema.py @@ -0,0 +1,19 @@ +from typing import List + +from pydantic import BaseModel + +from surya.common.polygon import PolygonBox + + +class BlockOCRResult(PolygonBox): + label: str # canonicalized layout label (Picture, Text, ...) + raw_label: str = "" # original model label + reading_order: int # 0-indexed position in layout output + html: str = "" # block HTML (BLOCK_PROMPT output, "" if skipped) + skipped: bool = False # True if label was in SKIP_OCR_LABELS + error: bool = False + + +class PageOCRResult(BaseModel): + blocks: List[BlockOCRResult] + image_bbox: List[float] diff --git a/surya/scripts/__init__.py b/surya/scripts/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/surya/scripts/build_gguf.py b/surya/scripts/build_gguf.py new file mode 100644 index 0000000..5f5fe4e --- /dev/null +++ b/surya/scripts/build_gguf.py @@ -0,0 +1,311 @@ +"""Build a GGUF artifact from a surya VLM checkpoint, suitable for stock llama.cpp. + +Three patches are applied to llama.cpp's convert_hf_to_gguf.py before +running the HF→GGUF conversion. All fixes are baked into the output +artifact, so the resulting GGUF runs on unpatched llama.cpp. + +Two checkpoint-side patches are also applied to a working copy: + - config.json: architectures → ["Qwen3_5ForConditionalGeneration"] + - tokenizer_config.json: tokenizer_class → "PreTrainedTokenizerFast", + strip backend / extra_special_tokens / is_local + +The original checkpoint is never modified — patches land in a sibling +working dir under --out-dir. + +Usage: + python -m surya.scripts.build_gguf \\ + --checkpoint datalab-to/surya-2.1.6 \\ + --out-dir ./gguf-build +""" + +import argparse +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path + +LLAMA_CPP_REPO = "https://github.com/ggerganov/llama.cpp.git" +# Pinned commit the converter patches were authored against. Bump when +# the upstream converter drifts; re-validate the anchor strings below. +LLAMA_CPP_REV = "bbeb89d76c41bc250f16e4a6fefcc9b530d6e3f3" + + +# ---- llama.cpp converter patches ----------------------------------------- +# Each patch is (sentinel, anchor, replacement). Idempotent: if `sentinel` +# is already in the file, the patch is skipped. + +_PATCH_REGISTER_HASH_ANCHOR = """ if chkhsh == "862f827721df956049dff5ca81a57f29e575280bc622e290d3bf4e35eca29015": + # ref: https://huggingface.co/codefuse-ai/F2LLM-v2-4B + res = "f2llmv2" +""" + +_PATCH_REGISTER_HASH_REPLACEMENT = ( + _PATCH_REGISTER_HASH_ANCHOR + + """ if chkhsh == "11865354be60ff9206694aed04242190f1807029bbd750bfbced09b4d26f1ad2": + # surya-2.1.0 char-level WordLevel tokenizer (Split regex=".") + res = "default" +""" +) + +_PATCH_VOCAB_ANCHOR = """ def _set_vocab_gpt2(self) -> None: + tokens, toktypes, tokpre = self.get_vocab_base() + self.gguf_writer.add_tokenizer_model("gpt2") + self.gguf_writer.add_tokenizer_pre(tokpre) + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) + special_vocab.add_to_gguf(self.gguf_writer) +""" + +_PATCH_VOCAB_REPLACEMENT = ''' def _set_vocab_gpt2(self) -> None: + tokens, toktypes, tokpre = self.get_vocab_base() + # surya: char-level / WordLevel vocabs contain raw bytes (e.g. " ", + # "\\n"). llama.cpp's gpt2 vocab decoder applies bytes_to_unicode + # when emitting NORMAL tokens, so encode bytes here for round-trip. + # Idempotent on already-encoded vocabs. + tokens = self._maybe_encode_gpt2_bytes(tokens, toktypes) + self.gguf_writer.add_tokenizer_model("gpt2") + self.gguf_writer.add_tokenizer_pre(tokpre) + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) + special_vocab.add_to_gguf(self.gguf_writer) + # surya: char-level / WordLevel vocabs have no merges, but the gpt2 + # vocab loader requires the field. Write a single dummy entry. + if not special_vocab.merges: + self.gguf_writer.add_token_merges(["a a"]) + + @staticmethod + def _maybe_encode_gpt2_bytes(tokens: list[str], toktypes: list[int]) -> list[str]: + """Apply GPT-2 bytes_to_unicode to NORMAL tokens iff any contain raw + whitespace/control bytes. Idempotent on already-encoded vocabs.""" + bs = (list(range(ord("!"), ord("~") + 1)) + + list(range(ord("¡"), ord("¬") + 1)) + + list(range(ord("®"), ord("ÿ") + 1))) + cs = bs[:] + n = 0 + for b in range(256): + if b not in bs: + bs.append(b) + cs.append(2 ** 8 + n) + n += 1 + byte_to_unicode = {b: chr(c) for b, c in zip(bs, cs)} + printable = set(range(ord("!"), ord("~") + 1)) + needs = False + for tok, ttype in zip(tokens, toktypes): + if ttype != gguf.TokenType.NORMAL: + continue + for ch in tok: + cb = ord(ch) + if cb < 0x80 and cb not in printable: + needs = True + break + if needs: + break + if not needs: + return tokens + out: list[str] = [] + for tok, ttype in zip(tokens, toktypes): + if ttype == gguf.TokenType.NORMAL: + out.append("".join(byte_to_unicode[b] for b in tok.encode("utf-8"))) + else: + out.append(tok) + return out +''' + +CONVERTER_PATCHES = [ + { + "name": "register surya-2.1.0 pre-tokenizer hash", + "sentinel": '"11865354be60ff9206694aed04242190f1807029bbd750bfbced09b4d26f1ad2"', + "anchor": _PATCH_REGISTER_HASH_ANCHOR, + "replacement": _PATCH_REGISTER_HASH_REPLACEMENT, + }, + { + "name": "byte-encode NORMAL tokens + dummy merges in _set_vocab_gpt2", + "sentinel": "_maybe_encode_gpt2_bytes", + "anchor": _PATCH_VOCAB_ANCHOR, + "replacement": _PATCH_VOCAB_REPLACEMENT, + }, +] + + +def patch_converter(convert_py: Path) -> None: + text = convert_py.read_text() + changed = False + for p in CONVERTER_PATCHES: + if p["sentinel"] in text: + print(f" [skip] {p['name']} (already applied)") + continue + if p["anchor"] not in text: + raise RuntimeError( + f"converter patch {p['name']!r} could not find its anchor. " + f"llama.cpp upstream may have drifted; re-pin LLAMA_CPP_REV " + f"and re-validate anchors." + ) + text = text.replace(p["anchor"], p["replacement"], 1) + print(f" [apply] {p['name']}") + changed = True + if changed: + convert_py.write_text(text) + + +# ---- checkpoint-side patches ---------------------------------------------- + + +def patch_checkpoint(src: Path, dst: Path) -> None: + """Symlink-clone src into dst, with config.json + tokenizer_config.json + rewritten for stock transformers/llama.cpp compatibility.""" + if dst.exists(): + shutil.rmtree(dst) + dst.mkdir(parents=True) + overrides = {"config.json", "tokenizer_config.json"} + for entry in src.iterdir(): + if entry.name in overrides: + continue + os.symlink(entry.resolve(), dst / entry.name) + + cfg = json.loads((src / "config.json").read_text()) + cfg["architectures"] = ["Qwen3_5ForConditionalGeneration"] + (dst / "config.json").write_text(json.dumps(cfg, indent=2)) + + tk = json.loads((src / "tokenizer_config.json").read_text()) + tk["tokenizer_class"] = "PreTrainedTokenizerFast" + for k in ("backend", "extra_special_tokens", "is_local"): + tk.pop(k, None) + (dst / "tokenizer_config.json").write_text(json.dumps(tk, indent=2)) + + +# ---- llama.cpp resolution ------------------------------------------------- + + +def ensure_llama_cpp(repo_dir: Path, rev: str) -> Path: + if not repo_dir.exists(): + repo_dir.parent.mkdir(parents=True, exist_ok=True) + print(f"[clone] {LLAMA_CPP_REPO} → {repo_dir}") + subprocess.check_call(["git", "clone", LLAMA_CPP_REPO, str(repo_dir)]) + head = subprocess.check_output( + ["git", "-C", str(repo_dir), "rev-parse", "HEAD"], text=True + ).strip() + if head != rev: + # Discard any prior patches so the checkout is clean. + subprocess.check_call( + ["git", "-C", str(repo_dir), "reset", "--hard", "--quiet", "HEAD"] + ) + subprocess.check_call( + ["git", "-C", str(repo_dir), "fetch", "--quiet", "origin"] + ) + print(f"[checkout] llama.cpp @ {rev}") + subprocess.check_call(["git", "-C", str(repo_dir), "checkout", "--quiet", rev]) + return repo_dir + + +# ---- checkpoint resolution ------------------------------------------------ + + +def resolve_checkpoint(checkpoint: str) -> Path: + p = Path(checkpoint) + if p.exists(): + return p.resolve() + from huggingface_hub import snapshot_download + + print(f"[download] {checkpoint}") + return Path(snapshot_download(checkpoint)) + + +# ---- main ----------------------------------------------------------------- + + +def main() -> int: + from surya.settings import settings + + ap = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + ap.add_argument( + "--checkpoint", + default=settings.SURYA_MODEL_CHECKPOINT, + help="HF repo id or local checkpoint dir", + ) + ap.add_argument("--out-dir", type=Path, default=Path("./gguf-build")) + ap.add_argument( + "--name", + default="surya-2", + help="Output basename. Produces .gguf and -mmproj.gguf", + ) + ap.add_argument( + "--llama-cpp-dir", + type=Path, + default=Path.home() / ".cache" / "datalab" / "llama.cpp", + ) + ap.add_argument("--llama-cpp-rev", default=LLAMA_CPP_REV) + ap.add_argument( + "--outtype", + default="f16", + help="convert_hf_to_gguf --outtype (f16, bf16, q8_0, ...)", + ) + ap.add_argument( + "--keep-work", + action="store_true", + help="Keep the patched-checkpoint working dir on success", + ) + args = ap.parse_args() + + args.out_dir.mkdir(parents=True, exist_ok=True) + work = args.out_dir / "_patched_ckpt" + + src = resolve_checkpoint(args.checkpoint) + print(f"[checkpoint] {src}") + + print("[patch] checkpoint config + tokenizer_config") + patch_checkpoint(src, work) + + repo = ensure_llama_cpp(args.llama_cpp_dir, args.llama_cpp_rev) + convert_py = repo / "convert_hf_to_gguf.py" + print("[patch] llama.cpp convert_hf_to_gguf.py") + patch_converter(convert_py) + + out_llm = (args.out_dir / f"{args.name}.gguf").resolve() + out_mmproj = (args.out_dir / f"{args.name}-mmproj.gguf").resolve() + + print(f"[convert] LLM → {out_llm}") + subprocess.check_call( + [ + sys.executable, + str(convert_py), + str(work), + "--outfile", + str(out_llm), + "--outtype", + args.outtype, + ] + ) + print(f"[convert] mmproj → {out_mmproj}") + subprocess.check_call( + [ + sys.executable, + str(convert_py), + str(work), + "--mmproj", + "--outfile", + str(out_mmproj), + "--outtype", + args.outtype, + ] + ) + + if not args.keep_work: + shutil.rmtree(work) + + print() + print(f" LLM: {out_llm}") + print(f" mmproj: {out_mmproj}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/surya/scripts/config.py b/surya/scripts/config.py new file mode 100644 index 0000000..921627a --- /dev/null +++ b/surya/scripts/config.py @@ -0,0 +1,98 @@ +from typing import List + +import click +import os +from surya.input.load import load_from_folder, load_from_file +from surya.settings import settings + + +class CLILoader: + def __init__(self, filepath: str, cli_options: dict, highres: bool = False): + self.page_range = cli_options.get("page_range") + if self.page_range: + self.page_range = self.parse_range_str(self.page_range) + self.filepath = filepath + self.config = cli_options + self.save_images = cli_options.get("images", False) + self.debug = cli_options.get("debug", False) + self.output_dir = cli_options.get("output_dir") + + # Opt in to leaving the inference server up so later commands reuse it. + if cli_options.get("keep_server"): + settings.SURYA_INFERENCE_KEEP_ALIVE = True + + self.load(highres) + + @staticmethod + def common_options(fn): + fn = click.argument("input_path", type=click.Path(exists=True), required=True)( + fn + ) + fn = click.option( + "--output_dir", + type=click.Path(exists=False), + required=False, + default=os.path.join(settings.RESULT_DIR, "surya"), + help="Directory to save output.", + )(fn) + fn = click.option( + "--page_range", + type=str, + default=None, + help="Page range to convert, specify comma separated page numbers or ranges. Example: 0,5-10,20", + )(fn) + fn = click.option( + "--images", + is_flag=True, + help="Save images of detected bboxes.", + default=False, + )(fn) + fn = click.option( + "--debug", "-d", is_flag=True, help="Enable debug mode.", default=False + )(fn) + fn = click.option( + "--keep_server", + is_flag=True, + default=False, + help="Keep the inference server (vllm/llama.cpp) running after this command exits so later commands reuse it instead of re-spawning.", + )(fn) + return fn + + def load(self, highres: bool = False): + highres_images = None + if os.path.isdir(self.filepath): + images, names = load_from_folder(self.filepath, self.page_range) + folder_name = os.path.basename(self.filepath) + if highres: + highres_images, _ = load_from_folder( + self.filepath, self.page_range, settings.IMAGE_DPI_HIGHRES + ) + else: + images, names = load_from_file(self.filepath, self.page_range) + folder_name = os.path.basename(self.filepath).split(".")[0] + if highres: + highres_images, _ = load_from_file( + self.filepath, self.page_range, settings.IMAGE_DPI_HIGHRES + ) + + self.images = images + self.highres_images = highres_images + self.names = names + + self.result_path = os.path.abspath(os.path.join(self.output_dir, folder_name)) + os.makedirs(self.result_path, exist_ok=True) + + @staticmethod + def parse_range_str(range_str: str) -> List[int]: + range_lst = range_str.split(",") + page_lst = [] + for i in range_lst: + if "-" in i: + start, end = i.split("-") + page_lst += list(range(int(start), int(end) + 1)) + else: + page_lst.append(int(i)) + page_lst = sorted( + list(set(page_lst)) + ) # Deduplicate page numbers and sort in order + return page_lst diff --git a/surya/scripts/detect_layout.py b/surya/scripts/detect_layout.py new file mode 100644 index 0000000..ffe3745 --- /dev/null +++ b/surya/scripts/detect_layout.py @@ -0,0 +1,62 @@ +import time +import click +import copy +import json +from collections import defaultdict + +from surya.inference import SuryaInferenceManager +from surya.layout import LayoutPredictor +from surya.debug.draw import draw_polys_on_image +from surya.logging import configure_logging, get_logger +from surya.scripts.config import CLILoader +import os + +configure_logging() +logger = get_logger() + + +@click.command(help="Detect layout of an input file or folder (PDFs or image).") +@CLILoader.common_options +def detect_layout_cli(input_path: str, **kwargs): + loader = CLILoader(input_path, kwargs) + + manager = SuryaInferenceManager() + layout_predictor = LayoutPredictor(manager) + + start = time.time() + layout_predictions = layout_predictor(loader.images) + + if loader.debug: + logger.debug(f"Layout took {time.time() - start} seconds") + + if loader.save_images: + for idx, (image, layout_pred, name) in enumerate( + zip(loader.images, layout_predictions, loader.names) + ): + polygons = [p.polygon for p in layout_pred.bboxes] + labels = [f"{p.label}-{p.position}" for p in layout_pred.bboxes] + bbox_image = draw_polys_on_image( + polygons, copy.deepcopy(image), labels=labels + ) + bbox_image.save( + os.path.join(loader.result_path, f"{name}_{idx}_layout.png") + ) + + predictions_by_page = defaultdict(list) + for idx, (pred, name, image) in enumerate( + zip(layout_predictions, loader.names, loader.images) + ): + out_pred = pred.model_dump() + out_pred["page"] = len(predictions_by_page[name]) + 1 + predictions_by_page[name].append(out_pred) + + with open( + os.path.join(loader.result_path, "results.json"), "w+", encoding="utf-8" + ) as f: + json.dump(predictions_by_page, f, ensure_ascii=False) + + logger.info(f"Wrote results to {loader.result_path}") + + +if __name__ == "__main__": + detect_layout_cli() diff --git a/surya/scripts/detect_text.py b/surya/scripts/detect_text.py new file mode 100644 index 0000000..efa2d38 --- /dev/null +++ b/surya/scripts/detect_text.py @@ -0,0 +1,59 @@ +import click +import copy +import json +import time +from collections import defaultdict + +from surya.detection import DetectionPredictor +from surya.debug.draw import draw_polys_on_image +from surya.logging import configure_logging, get_logger +from surya.scripts.config import CLILoader +import os + +configure_logging() +logger = get_logger() + + +@click.command(help="Detect bboxes in an input file or folder (PDFs or image).") +@CLILoader.common_options +def detect_text_cli(input_path: str, **kwargs): + loader = CLILoader(input_path, kwargs) + + det_predictor = DetectionPredictor() + + start = time.time() + predictions = det_predictor(loader.images, include_maps=loader.debug) + end = time.time() + if loader.debug: + logger.debug(f"Detection took {end - start} seconds") + + if loader.save_images: + for idx, (image, pred, name) in enumerate( + zip(loader.images, predictions, loader.names) + ): + polygons = [p.polygon for p in pred.bboxes] + bbox_image = draw_polys_on_image(polygons, copy.deepcopy(image)) + bbox_image.save(os.path.join(loader.result_path, f"{name}_{idx}_bbox.png")) + + if loader.debug: + heatmap = pred.heatmap + heatmap.save(os.path.join(loader.result_path, f"{name}_{idx}_heat.png")) + + predictions_by_page = defaultdict(list) + for idx, (pred, name, image) in enumerate( + zip(predictions, loader.names, loader.images) + ): + out_pred = pred.model_dump(exclude=["heatmap", "affinity_map"]) + out_pred["page"] = len(predictions_by_page[name]) + 1 + predictions_by_page[name].append(out_pred) + + with open( + os.path.join(loader.result_path, "results.json"), "w+", encoding="utf-8" + ) as f: + json.dump(predictions_by_page, f, ensure_ascii=False) + + logger.info(f"Wrote results to {loader.result_path}") + + +if __name__ == "__main__": + detect_text_cli() diff --git a/surya/scripts/ocr_text.py b/surya/scripts/ocr_text.py new file mode 100644 index 0000000..7ae3f0a --- /dev/null +++ b/surya/scripts/ocr_text.py @@ -0,0 +1,48 @@ +import os +import click +import json +import time +from collections import defaultdict + +from surya.inference import SuryaInferenceManager +from surya.logging import configure_logging, get_logger +from surya.recognition import RecognitionPredictor +from surya.scripts.config import CLILoader + +configure_logging() +logger = get_logger() + + +@click.command(help="OCR text — full-page OCR (one VLM call per page).") +@CLILoader.common_options +def ocr_text_cli(input_path: str, **kwargs): + # Full-page OCR is the default path: one VLM call per page returns layout + # + content together. Pages whose full-page output fails to parse fall + # back to layout + per-block OCR automatically (see RecognitionPredictor). + loader = CLILoader(input_path, kwargs, highres=True) + + manager = SuryaInferenceManager() + rec_predictor = RecognitionPredictor(manager) + + start = time.time() + page_results = rec_predictor(loader.highres_images, full_page=True) + + if loader.debug: + logger.debug(f"OCR took {time.time() - start:.2f} seconds") + + out_preds = defaultdict(list) + for name, page in zip(loader.names, page_results): + out_pred = page.model_dump() + out_pred["page"] = len(out_preds[name]) + 1 + out_preds[name].append(out_pred) + + with open( + os.path.join(loader.result_path, "results.json"), "w+", encoding="utf-8" + ) as f: + json.dump(out_preds, f, ensure_ascii=False) + + logger.info(f"Wrote results to {loader.result_path}") + + +if __name__ == "__main__": + ocr_text_cli() diff --git a/surya/scripts/run_streamlit_app.py b/surya/scripts/run_streamlit_app.py new file mode 100644 index 0000000..658b2ce --- /dev/null +++ b/surya/scripts/run_streamlit_app.py @@ -0,0 +1,9 @@ +import subprocess +import os + + +def streamlit_app_cli(): + cur_dir = os.path.dirname(os.path.abspath(__file__)) + ocr_app_path = os.path.join(cur_dir, "streamlit_app.py") + cmd = ["streamlit", "run", ocr_app_path, "--server.fileWatcherType", "none", "--server.headless", "true"] + subprocess.run(cmd, env={**os.environ, "IN_STREAMLIT": "true"}) \ No newline at end of file diff --git a/surya/scripts/screenshot_app.py b/surya/scripts/screenshot_app.py new file mode 100644 index 0000000..34b88ab --- /dev/null +++ b/surya/scripts/screenshot_app.py @@ -0,0 +1,226 @@ +"""Screenshot-friendly Surya viewer. + +Shows a PDF/image page on the left and full-page OCR output on the right, side +by side, for clean screenshots. You can scroll through pages and preview them +before running OCR, then export the side-by-side view as a PNG. + +Run with `surya_screenshot`, then open http://localhost:8504. +""" + +from __future__ import annotations + +import base64 +import io +import os +import tempfile +import uuid +from typing import List, Optional + +import pypdfium2 +from flask import Flask, jsonify, render_template, request +from PIL import Image +from werkzeug.utils import secure_filename + +from surya.inference import SuryaInferenceManager +from surya.logging import configure_logging, get_logger +from surya.recognition import RecognitionPredictor +from surya.recognition.schema import PageOCRResult +from surya.settings import settings + +configure_logging() +logger = get_logger() + +app = Flask(__name__) + +ALLOWED_EXT = {".pdf", ".png", ".jpg", ".jpeg", ".gif", ".webp"} +UPLOAD_DIR = os.path.join(tempfile.gettempdir(), "surya_screenshot") +os.makedirs(UPLOAD_DIR, exist_ok=True) + +_rec: Optional[RecognitionPredictor] = None + + +def get_rec() -> RecognitionPredictor: + """Lazily build the recognition predictor (shared inference manager).""" + global _rec + if _rec is None: + _rec = RecognitionPredictor(SuryaInferenceManager()) + return _rec + + +# Datalab-flavored palette for layout block overlays, keyed by canonical label. +LABEL_COLORS = { + "Text": "#2563eb", + "SectionHeader": "#0ea5e9", + "PageHeader": "#7c3aed", + "PageFooter": "#7c3aed", + "Caption": "#c026d3", + "Footnote": "#64748b", + "Equation": "#9333ea", + "Table": "#f59e0b", + "TableOfContents": "#f59e0b", + "Form": "#ea580c", + "ListGroup": "#10b981", + "Picture": "#db2777", + "Figure": "#db2777", + "Diagram": "#db2777", + "Code": "#0d9488", + "default": "#ef4444", +} + + +def _logo_data_url() -> str: + path = os.path.join(settings.BASE_DIR, "static", "datalab-logo.png") + try: + with open(path, "rb") as f: + return "data:image/png;base64," + base64.b64encode(f.read()).decode() + except Exception: + return "" + + +def _pil_to_data_url(img: Image.Image, fmt: str = "PNG") -> str: + buf = io.BytesIO() + img.save(buf, format=fmt) + return ( + f"data:image/{fmt.lower()};base64," + base64.b64encode(buf.getvalue()).decode() + ) + + +def _is_pdf(path: str) -> bool: + return path.lower().endswith(".pdf") + + +def _page_count(path: str) -> int: + if _is_pdf(path): + doc = pypdfium2.PdfDocument(path) + n = len(doc) + doc.close() + return n + return 1 + + +def _render_page(path: str, page: int, dpi: int) -> Image.Image: + """Render a 0-indexed page of a PDF (or load an image file) as RGB.""" + if _is_pdf(path): + doc = pypdfium2.PdfDocument(path) + try: + pil = doc[page].render(scale=dpi / 72).to_pil().convert("RGB") + finally: + doc.close() + return pil + return Image.open(path).convert("RGB") + + +def _assemble_page_html(page: PageOCRResult) -> str: + """Whole-page HTML from a PageOCRResult (math stays in tags).""" + parts: List[str] = [] + for blk in page.blocks: + if blk.skipped: + continue + x0, y0, x1, y1 = (int(c) for c in blk.bbox) + parts.append( + f'
{blk.html or ""}
' + ) + return "\n".join(parts) + + +@app.route("/") +def index(): + return render_template("surya_screenshot.html", logo=_logo_data_url()) + + +@app.route("/info", methods=["POST"]) +def info(): + path = (request.json or {}).get("file_path", "").strip() + if not path: + return jsonify({"error": "file_path is required"}), 400 + if not os.path.exists(path): + return jsonify({"error": f"File not found: {path}"}), 400 + try: + return jsonify({"page_count": _page_count(path)}) + except Exception as e: + return jsonify({"error": str(e)}), 500 + + +@app.route("/upload", methods=["POST"]) +def upload(): + """Accept a drag/drop (or browsed) file, save to a temp path, return it.""" + f = request.files.get("file") + if f is None or not f.filename: + return jsonify({"error": "no file uploaded"}), 400 + ext = os.path.splitext(f.filename)[1].lower() + if ext not in ALLOWED_EXT: + return jsonify({"error": f"unsupported file type: {ext or '(none)'}"}), 400 + safe = secure_filename(f.filename) or f"upload{ext}" + dest = os.path.join(UPLOAD_DIR, f"{uuid.uuid4().hex}_{safe}") + f.save(dest) + try: + return jsonify( + {"file_path": dest, "page_count": _page_count(dest), "name": f.filename} + ) + except Exception as e: + return jsonify({"error": str(e)}), 500 + + +@app.route("/page", methods=["POST"]) +def page(): + """Render a single page for preview (no OCR).""" + data = request.json or {} + path = data.get("file_path", "").strip() + page_num = int(data.get("page", 0)) + if not path or not os.path.exists(path): + return jsonify({"error": "valid file_path is required"}), 400 + try: + img = _render_page(path, page_num, settings.IMAGE_DPI_HIGHRES) + return jsonify( + { + "image_base64": _pil_to_data_url(img), + "width": img.size[0], + "height": img.size[1], + } + ) + except Exception as e: + return jsonify({"error": str(e)}), 500 + + +@app.route("/process", methods=["POST"]) +def process(): + """Run full-page OCR on one page; return the page image + OCR HTML + blocks.""" + data = request.json or {} + path = data.get("file_path", "").strip() + page_num = int(data.get("page", 0)) + if not path or not os.path.exists(path): + return jsonify({"error": "valid file_path is required"}), 400 + try: + img = _render_page(path, page_num, settings.IMAGE_DPI_HIGHRES) + page_result = get_rec()([img], full_page=True)[0] + blocks = [ + { + "bbox": [int(c) for c in blk.bbox], + "label": blk.label, + "color": LABEL_COLORS.get(blk.label, LABEL_COLORS["default"]), + } + for blk in page_result.blocks + if not blk.skipped + ] + return jsonify( + { + "image_base64": _pil_to_data_url(img), + "width": img.size[0], + "height": img.size[1], + "html": _assemble_page_html(page_result), + "blocks": blocks, + "n_blocks": len(page_result.blocks), + } + ) + except Exception as e: + logger.exception("Full-page OCR failed") + return jsonify({"error": str(e)}), 500 + + +def main(): + app.run(host="0.0.0.0", port=8504) + + +if __name__ == "__main__": + main() diff --git a/surya/scripts/streamlit_app.py b/surya/scripts/streamlit_app.py new file mode 100644 index 0000000..24d0cc9 --- /dev/null +++ b/surya/scripts/streamlit_app.py @@ -0,0 +1,490 @@ +"""Surya2 streamlit app — exercise layout, recognition, table_rec via the +inference manager. Detection + OCR-error stay in their own torch paths.""" + +from __future__ import annotations + +import io +import re +import tempfile +import time +from typing import List + +import pypdfium2 +import streamlit as st +import streamlit.components.v1 as components +from PIL import Image, ImageDraw + +from surya.debug.draw import draw_polys_on_image, draw_bboxes_on_image +from surya.detection import TextDetectionResult +from surya.inference import SuryaInferenceManager +from surya.layout import LayoutPredictor +from surya.layout.schema import LayoutResult +from surya.recognition import RecognitionPredictor +from surya.recognition.schema import PageOCRResult +from surya.settings import settings +from surya.table_rec import TableRecPredictor +from surya.table_rec.schema import TableResult + + +# KaTeX-enabled HTML wrapper. The OCR HTML wraps math in ... +# (KaTeX-compatible LaTeX inside), which a browser would otherwise show as +# raw text. We convert those tags to \( \) / \[ \] delimiters and let KaTeX +# auto-render typeset them inside an iframe component. +_KATEX_HEAD = r""" + + + + + +""" + +_KATEX_TAIL = r""" + +""" + +_MATH_RE = re.compile(r"]*)>(.*?)", re.DOTALL | re.IGNORECASE) + + +def _math_to_katex(html_str: str) -> str: + """Rewrite ... tags into KaTeX \\( \\) / \\[ \\] delimiters.""" + + def repl(m: "re.Match") -> str: + attrs, inner = m.group(1), m.group(2) + if re.search(r"""display\s*=\s*["']block["']""", attrs): + return "\\[" + inner + "\\]" + return "\\(" + inner + "\\)" + + return _MATH_RE.sub(repl, html_str or "") + + +def render_ocr_html(html_str: str, height: int = 400) -> None: + """Render OCR HTML with math typeset by KaTeX (iframe component).""" + components.html( + _KATEX_HEAD + _math_to_katex(html_str) + _KATEX_TAIL, + height=height, + scrolling=True, + ) + + +def _assemble_page_html(page: PageOCRResult) -> str: + """Reconstruct a div-block whole-page HTML from a PageOCRResult.""" + parts: List[str] = [] + for blk in page.blocks: + if blk.skipped: + continue + x0, y0, x1, y1 = (int(c) for c in blk.bbox) + body = blk.html or "" + parts.append( + f'
{body}
' + ) + return "\n".join(parts) + + +def _show_timing(label: str, elapsed_s: float, extra: str = "") -> None: + """Render a small caption with wall-clock + optional extra detail.""" + detail = f" — {extra}" if extra else "" + st.caption(f"⏱ {label}: {elapsed_s * 1000:.0f} ms ({elapsed_s:.2f}s){detail}") + + +@st.cache_resource() +def load_predictors_cached(): + manager = SuryaInferenceManager() + layout_predictor = LayoutPredictor(manager) + rec_predictor = RecognitionPredictor(manager) + table_rec_predictor = TableRecPredictor(manager) + + # Lazy-import detection / ocr_error to keep startup snappy when the user + # only wants VLM modes + from surya.detection import DetectionPredictor + from surya.ocr_error import OCRErrorPredictor + + return { + "manager": manager, + "layout": layout_predictor, + "recognition": rec_predictor, + "table_rec": table_rec_predictor, + "detection": DetectionPredictor(), + "ocr_error": OCRErrorPredictor(), + } + + +def text_detection(img) -> tuple[Image.Image, TextDetectionResult, float]: + t = time.perf_counter() + text_pred = predictors["detection"]([img])[0] + elapsed = time.perf_counter() - t + text_polygons = [p.polygon for p in text_pred.bboxes] + det_img = draw_polys_on_image(text_polygons, img.copy()) + return det_img, text_pred, elapsed + + +def layout_detection(img) -> tuple[Image.Image, LayoutResult, float]: + t = time.perf_counter() + pred = predictors["layout"]([img])[0] + elapsed = time.perf_counter() - t + polygons = [p.polygon for p in pred.bboxes] + labels = [ + f"{p.label}-{p.position}-c{p.count}-{round(p.confidence or 0, 2)}" + for p in pred.bboxes + ] + annotated = draw_polys_on_image( + polygons, img.copy(), labels=labels, label_font_size=14 + ) + return annotated, pred, elapsed + + +def block_ocr(img) -> tuple[Image.Image, PageOCRResult, LayoutResult, float, float]: + """Layout → block crops → BLOCK_PROMPT. Returns layout + block-OCR timings.""" + t_layout = time.perf_counter() + layout = predictors["layout"]([img])[0] + layout_elapsed = time.perf_counter() - t_layout + + t_blocks = time.perf_counter() + page_results = predictors["recognition"]([img], [layout]) + blocks_elapsed = time.perf_counter() - t_blocks + page = page_results[0] + + annotated = img.copy() + draw = ImageDraw.Draw(annotated) + for blk in page.blocks: + x0, y0, x1, y1 = blk.bbox + color = "red" if blk.error else ("orange" if blk.skipped else "green") + draw.rectangle((x0, y0, x1, y1), outline=color, width=3) + draw.text((x0 + 4, y0 + 4), f"{blk.reading_order} {blk.label}", fill=color) + return annotated, page, layout, layout_elapsed, blocks_elapsed + + +def full_page_ocr(img) -> tuple[Image.Image, PageOCRResult, float]: + """Single HIGH_ACCURACY_BBOX_PROMPT call on the whole page.""" + t = time.perf_counter() + page_results = predictors["recognition"]([img], full_page=True) + elapsed = time.perf_counter() - t + page = page_results[0] + annotated = img.copy() + draw = ImageDraw.Draw(annotated) + for blk in page.blocks: + x0, y0, x1, y1 = blk.bbox + color = "red" if blk.error else ("orange" if blk.skipped else "green") + draw.rectangle((x0, y0, x1, y1), outline=color, width=3) + draw.text((x0 + 4, y0 + 4), f"{blk.reading_order} {blk.label}", fill=color) + return annotated, page, elapsed + + +def table_recognition( + img: Image.Image, + mode: str, + skip_table_detection: bool, +) -> tuple[Image.Image, List[TableResult], float, float]: + """Returns (annotated_img, table_preds, layout_elapsed, table_rec_elapsed).""" + layout_elapsed = 0.0 + if skip_table_detection: + table_imgs = [img] + table_counts = [0] + table_bboxes = [(0, 0, img.size[0], img.size[1])] + else: + t = time.perf_counter() + layout = predictors["layout"]([img])[0] + layout_elapsed = time.perf_counter() - t + tables = [b for b in layout.bboxes if b.label in ("Table", "TableOfContents")] + if not tables: + return img.copy(), [], layout_elapsed, 0.0 + table_bboxes = [tuple(int(c) for c in b.bbox) for b in tables] + table_imgs = [img.crop(b) for b in table_bboxes] + table_counts = [b.count for b in tables] + + t = time.perf_counter() + if mode == "full": + table_preds = predictors["table_rec"].predict_full( + table_imgs, counts=table_counts + ) + else: + table_preds = predictors["table_rec"].predict_simple(table_imgs) + table_rec_elapsed = time.perf_counter() - t + + out_img = img.copy() + for pred, table_img, tbbox in zip(table_preds, table_imgs, table_bboxes): + if pred.error or pred.mode != "simple" or not pred.rows: + continue + row_bboxes = [r.bbox for r in pred.rows] + col_bboxes = [c.bbox for c in pred.cols] + row_labels = [r.label for r in pred.rows] + col_labels = [c.label for c in pred.cols] + annot = table_img.copy() + annot = draw_bboxes_on_image( + row_bboxes, annot, labels=row_labels, label_font_size=14, color="blue" + ) + annot = draw_bboxes_on_image( + col_bboxes, annot, labels=col_labels, label_font_size=14, color="red" + ) + # Paste annotated crop back at the table's position in the page. + out_img.paste(annot, (tbbox[0], tbbox[1])) + return out_img, table_preds, layout_elapsed, table_rec_elapsed + + +def ocr_errors(pdf_file, page_count, sample_len=512, max_samples=10, max_pages=15): + from pdftext.extraction import plain_text_output + + with tempfile.NamedTemporaryFile(suffix=".pdf") as f: + f.write(pdf_file.getvalue()) + f.seek(0) + + page_middle = page_count // 2 + page_range = range( + max(page_middle - max_pages, 0), min(page_middle + max_pages, page_count) + ) + text = plain_text_output(f.name, page_range=page_range) + + sample_gap = len(text) // max_samples + if len(text) == 0 or sample_gap == 0: + return "This PDF has no text or very little text", ["no text"] + + if sample_gap < sample_len: + sample_gap = sample_len + + samples = [] + for i in range(0, len(text), sample_gap): + samples.append(text[i : i + sample_len]) + + results = predictors["ocr_error"](samples) + label = "This PDF has good text." + if results.labels.count("bad") / len(results.labels) > 0.2: + label = "This PDF may have garbled or bad OCR text." + return label, results.labels + + +def open_pdf(pdf_file): + stream = io.BytesIO(pdf_file.getvalue()) + return pypdfium2.PdfDocument(stream) + + +@st.cache_data() +def get_page_image(pdf_file, page_num, dpi=settings.IMAGE_DPI): + doc = open_pdf(pdf_file) + renderer = doc.render( + pypdfium2.PdfBitmap.to_pil, + page_indices=[page_num - 1], + scale=dpi / 72, + ) + png = list(renderer)[0] + png_image = png.convert("RGB") + doc.close() + return png_image + + +@st.cache_data() +def page_counter(pdf_file): + doc = open_pdf(pdf_file) + doc_len = len(doc) + doc.close() + return doc_len + + +st.set_page_config(layout="wide") +col1, col2 = st.columns([0.55, 0.45]) + +predictors = load_predictors_cached() + +st.markdown( + """ +# Surya 2 Demo + +VLM-backed layout, OCR, and table recognition. The model runs in a local +`llama-server` (or vllm) process, started on first use. + +Modes: +- **Layout**: page → list of blocks with label + bbox + token count +- **Block OCR**: layout + per-block HTML +- **Table Rec (simple)**: row + column bboxes only +- **Table Rec (full)**: full HTML for each detected table +""" +) + +in_file = st.sidebar.file_uploader( + "PDF file or image:", type=["pdf", "png", "jpg", "jpeg", "gif", "webp"] +) + +if in_file is None: + st.stop() + +filetype = in_file.type +page_count = None +if "pdf" in filetype: + page_count = page_counter(in_file) + page_number = st.sidebar.number_input( + f"Page number out of {page_count}:", min_value=1, value=1, max_value=page_count + ) + # Render at high DPI so the OCR / table-rec demos see fine glyphs. + # Layout + detection internally downsample (or accept the small perf hit + # at demo scale); we always render and display the high-DPI page here. + pil_image = get_page_image(in_file, page_number, settings.IMAGE_DPI_HIGHRES) +else: + pil_image = Image.open(in_file).convert("RGB") + page_number = None + +run_full_page_ocr = st.sidebar.button("Run Full-Page OCR") +run_text_det = st.sidebar.button("Run Text Detection") +run_layout = st.sidebar.button("Run Layout Analysis") +run_table_rec = st.sidebar.button("Run Table Rec") +run_block_ocr = st.sidebar.button("Run Block OCR") +run_ocr_errors = st.sidebar.button("Run bad-PDF-text detection") + +table_mode = st.sidebar.radio( + "Table mode", + options=["simple", "full"], + index=0, + help="simple: rows+cols only. full: full HTML.", +) +skip_table_detection = st.sidebar.checkbox( + "Skip table detection", + value=False, + help="Treat the entire page/image as a single table.", +) + +if pil_image is None: + st.stop() + + +if run_text_det: + det_img, text_pred, elapsed = text_detection(pil_image) + with col1: + _show_timing("Text detection", elapsed, f"{len(text_pred.bboxes)} polys") + st.image(det_img, caption="Detected Text", use_container_width=True) + st.json( + text_pred.model_dump(exclude=["heatmap", "affinity_map"]), expanded=False + ) + + +if run_layout: + annotated, pred, elapsed = layout_detection(pil_image) + with col1: + _show_timing("Layout", elapsed, f"{len(pred.bboxes)} blocks") + st.image(annotated, caption="Detected Layout", use_container_width=True) + st.json(pred.model_dump(), expanded=False) + + +if run_block_ocr: + annotated, page, layout, t_layout, t_blocks = block_ocr(pil_image) + with col1: + n_blocks = len(page.blocks) + n_ok = sum(1 for b in page.blocks if not b.skipped and not b.error) + _show_timing("Block OCR — layout", t_layout, f"{n_blocks} blocks") + _show_timing("Block OCR — per-block OCR", t_blocks, f"{n_ok} OCR'd") + _show_timing("Block OCR — total", t_layout + t_blocks) + st.image( + annotated, + caption="Block OCR (green=ok, orange=skipped, red=error)", + use_container_width=True, + ) + full_html = _assemble_page_html(page) + with st.expander("Full page HTML (rendered)", expanded=False): + render_ocr_html(full_html, height=600) + with st.expander("Full page HTML (source)", expanded=False): + st.code(full_html, language="html") + for blk in page.blocks: + with st.expander( + f"#{blk.reading_order} {blk.label} (conf {blk.confidence:.2f})" + ): + # Diagnostics: show numeric bbox + polygon + a thumbnail with the + # drawn rectangle highlighted, then the actual crop fed to OCR. + xs = [p[0] for p in blk.polygon] + ys = [p[1] for p in blk.polygon] + bbox_drawn = [int(min(xs)), int(min(ys)), int(max(xs)), int(max(ys))] + cx0 = max(0, int(min(xs)) - 4) + cy0 = max(0, int(min(ys)) - 4) + cx1 = min(pil_image.size[0], int(max(xs)) + 4) + cy1 = min(pil_image.size[1], int(max(ys)) + 4) + st.text( + f"bbox(drawn) = {bbox_drawn}\n" + f"crop(ocr) = {(cx0, cy0, cx1, cy1)} (= bbox ± 4px pad)" + ) + # Thumbnail with this block's rectangle highlighted in red. + thumb = pil_image.copy() + ImageDraw.Draw(thumb).rectangle(bbox_drawn, outline="red", width=4) + st.image(thumb, caption="this block's drawn rect (red)", width=300) + # The actual crop fed to OCR + if cx1 > cx0 and cy1 > cy0: + st.image(pil_image.crop((cx0, cy0, cx1, cy1)), caption="OCR crop") + if blk.skipped: + st.info("Block skipped (visual label)") + elif blk.error: + st.error("Block OCR errored") + else: + render_ocr_html(blk.html, height=160) + st.code(blk.html, language="html") + + +if run_full_page_ocr: + annotated, page, elapsed = full_page_ocr(pil_image) + with col1: + n_blocks = len(page.blocks) + n_ok = sum(1 for b in page.blocks if not b.skipped and not b.error) + _show_timing("Full-Page OCR", elapsed, f"{n_blocks} blocks parsed, {n_ok} OK") + st.image( + annotated, + caption="Full-Page OCR (green=ok, orange=skipped, red=error)", + use_container_width=True, + ) + full_html = _assemble_page_html(page) + with st.expander("Full page HTML (rendered)", expanded=False): + render_ocr_html(full_html, height=600) + with st.expander("Full page HTML (source)", expanded=False): + st.code(full_html, language="html") + for blk in page.blocks: + with st.expander( + f"#{blk.reading_order} {blk.label} (conf {blk.confidence:.2f})" + ): + if blk.skipped: + st.info("Block skipped (visual label)") + elif blk.error: + st.error("Block OCR errored") + else: + render_ocr_html(blk.html, height=160) + st.code(blk.html, language="html") + + +if run_table_rec: + table_img, preds, t_layout, t_table = table_recognition( + pil_image, table_mode, skip_table_detection + ) + with col1: + if not skip_table_detection: + _show_timing("Table Rec — layout", t_layout, f"{len(preds)} tables found") + _show_timing(f"Table Rec — {table_mode}", t_table) + if not skip_table_detection: + _show_timing("Table Rec — total", t_layout + t_table) + st.image(table_img, caption="Table Recognition", use_container_width=True) + for pred in preds: + if pred.mode == "full" and pred.html: + with st.expander("Table HTML"): + render_ocr_html(pred.html, height=400) + st.code(pred.html, language="html") + else: + st.json(pred.model_dump(), expanded=False) + + +if run_ocr_errors: + if "pdf" not in filetype: + st.error("This feature only works with PDFs.") + else: + label, results = ocr_errors(in_file, page_count) + with col1: + st.write(label) + st.json(results) + + +with col2: + st.image(pil_image, caption="Uploaded Image", use_container_width=True) diff --git a/surya/scripts/table_recognition.py b/surya/scripts/table_recognition.py new file mode 100644 index 0000000..64fb40f --- /dev/null +++ b/surya/scripts/table_recognition.py @@ -0,0 +1,139 @@ +import os +import click +import copy +import json +from collections import defaultdict + +from surya.common.util import expand_bbox +from surya.debug.draw import draw_bboxes_on_image +from surya.inference import SuryaInferenceManager +from surya.layout import LayoutPredictor +from surya.logging import configure_logging, get_logger +from surya.scripts.config import CLILoader +from surya.table_rec import TableRecPredictor + +configure_logging() +logger = get_logger() + + +@click.command(help="Run table recognition on an input file or folder.") +@CLILoader.common_options +@click.option( + "--skip_table_detection", + is_flag=True, + help="Tables are already cropped, so don't re-detect tables.", + default=False, +) +@click.option( + "--mode", + type=click.Choice(["simple", "full"]), + default="simple", + help="simple: rows+cols only (geometric cells). full: full HTML (BLOCK_PROMPT).", +) +def table_recognition_cli( + input_path: str, skip_table_detection: bool, mode: str, **kwargs +): + # Layout runs on the low-DPI render; table crops come from the high-DPI + # image so the table_rec model sees readable cell content. + loader = CLILoader(input_path, kwargs, highres=True) + + manager = SuryaInferenceManager() + layout_predictor = LayoutPredictor(manager) + table_rec_predictor = TableRecPredictor(manager) + + pnums = [] + prev_name = None + for name in loader.names: + if prev_name is None or prev_name != name: + pnums.append(0) + else: + pnums.append(pnums[-1] + 1) + prev_name = name + + table_imgs = [] + table_counts = [] + table_counts_per_img = [] + + if skip_table_detection: + for img in loader.highres_images: + table_imgs.append(img) + table_counts.append(1) + table_counts_per_img.append(0) + else: + layout_predictions = layout_predictor( + loader.images, + target_image_sizes=[img.size for img in loader.highres_images], + ) + for layout_pred, img in zip(layout_predictions, loader.highres_images): + tables_on_page = [ + line + for line in layout_pred.bboxes + if line.label in ("Table", "TableOfContents") + ] + table_counts.append(len(tables_on_page)) + for line in tables_on_page: + bbox = expand_bbox(line.bbox) + table_imgs.append(img.crop(bbox)) + table_counts_per_img.append(line.count) + + table_preds = table_rec_predictor(table_imgs, mode=mode) + + img_idx = 0 + prev_count = 0 + table_predictions = defaultdict(list) + for i in range(sum(table_counts)): + while i >= prev_count + table_counts[img_idx]: + prev_count += table_counts[img_idx] + img_idx += 1 + + pred = table_preds[i] + orig_name = loader.names[img_idx] + pnum = pnums[img_idx] + table_img = table_imgs[i] + + out_pred = pred.model_dump() + out_pred["page"] = pnum + 1 + table_idx = i - prev_count + out_pred["table_idx"] = table_idx + table_predictions[orig_name].append(out_pred) + + if loader.save_images and pred.rows: + rows = [line.bbox for line in pred.rows] + cols = [line.bbox for line in pred.cols] + row_labels = [f"Row {line.row_id}" for line in pred.rows] + col_labels = [f"Col {line.col_id}" for line in pred.cols] + cells = [line.bbox for line in pred.cells] + + rc_image = copy.deepcopy(table_img) + rc_image = draw_bboxes_on_image( + rows, rc_image, labels=row_labels, label_font_size=20, color="blue" + ) + rc_image = draw_bboxes_on_image( + cols, rc_image, labels=col_labels, label_font_size=20, color="red" + ) + rc_image.save( + os.path.join( + loader.result_path, + f"{orig_name}_page{pnum + 1}_table{table_idx}_rc.png", + ) + ) + + cell_image = copy.deepcopy(table_img) + cell_image = draw_bboxes_on_image(cells, cell_image, color="green") + cell_image.save( + os.path.join( + loader.result_path, + f"{orig_name}_page{pnum + 1}_table{table_idx}_cells.png", + ) + ) + + with open( + os.path.join(loader.result_path, "results.json"), "w+", encoding="utf-8" + ) as f: + json.dump(table_predictions, f, ensure_ascii=False) + + logger.info(f"Wrote results to {loader.result_path}") + + +if __name__ == "__main__": + table_recognition_cli() diff --git a/surya/scripts/templates/surya_screenshot.html b/surya/scripts/templates/surya_screenshot.html new file mode 100644 index 0000000..c472f16 --- /dev/null +++ b/surya/scripts/templates/surya_screenshot.html @@ -0,0 +1,331 @@ + + + + + + Surya · Full-Page OCR + + + + + + +
+
+ {% if logo %}Datalab{% endif %} + Surya + Full-Page OCR +
+
+ + + + +
+ + — + +
+ + + + + +
+
+ +
+
+
PDF Page
+
+
+
+
Full-Page OCR
+
Load a file, scroll to a page, then run full-page OCR.
+
+
+ +
Drop a PDF or image to load
+ + + + diff --git a/surya/settings.py b/surya/settings.py new file mode 100644 index 0000000..81edfda --- /dev/null +++ b/surya/settings.py @@ -0,0 +1,165 @@ +import os +from typing import Callable, Dict, Optional + +import torch +from dotenv import find_dotenv +from pydantic import computed_field +from pydantic_settings import BaseSettings +from pathlib import Path +from platformdirs import user_cache_dir + + +class Settings(BaseSettings): + # General + TORCH_DEVICE: Optional[str] = None + IMAGE_DPI: int = 96 # used for layout + text detection (coarse structure) + IMAGE_DPI_HIGHRES: int = 192 # used for recognition + table rec (fine glyphs) + IN_STREAMLIT: bool = False + DISABLE_TQDM: bool = False + S3_BASE_URL: str = "https://models.datalab.to" + PARALLEL_DOWNLOAD_WORKERS: int = 10 + MODEL_CACHE_DIR: str = str(Path(user_cache_dir("datalab")) / "models") + LOGLEVEL: str = "INFO" + + # Paths + RESULT_DIR: str = "results" + BASE_DIR: str = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + FONT_DIR: str = os.path.join(BASE_DIR, "static", "fonts") + + @computed_field + def TORCH_DEVICE_MODEL(self) -> str: + if self.TORCH_DEVICE is not None: + return self.TORCH_DEVICE + if torch.cuda.is_available(): + return "cuda" + if torch.backends.mps.is_available(): + return "mps" + return "cpu" + + # ---- Surya2 inference (VLM-backed: vllm | llamacpp) --------------------- + SURYA_MODEL_CHECKPOINT: str = "datalab-to/surya-ocr-2" + SURYA_GGUF_REPO: str = "datalab-to/surya-ocr-2-gguf" + SURYA_GGUF_MODEL_FILE: str = "surya-2.gguf" + SURYA_GGUF_MMPROJ_FILE: str = "surya-2-mmproj.gguf" + # If set, used directly instead of HF download (handy for local-conversion testing) + SURYA_GGUF_LOCAL_MODEL_PATH: Optional[str] = None + SURYA_GGUF_LOCAL_MMPROJ_PATH: Optional[str] = None + + # Backend selection + SURYA_INFERENCE_BACKEND: Optional[str] = None # "vllm" | "llamacpp" | None (auto) + SURYA_INFERENCE_URL: Optional[str] = None # external server, skip spawn + SURYA_INFERENCE_AUTOSTART: bool = True + # Leave an auto-spawned server running after the process exits so later + # commands attach to it instead of re-spawning (avoids repeated startup / + # model-load cost). Stop it manually when done — see `surya/inference`. + SURYA_INFERENCE_KEEP_ALIVE: bool = False + SURYA_INFERENCE_HOST: str = "127.0.0.1" + SURYA_INFERENCE_PORT: Optional[int] = None # None = pick a free port + SURYA_INFERENCE_PARALLEL: int = 8 + # Max concurrent in-flight chat-completion requests to the inference server + # per batch. Cap to roughly VLLM_MAX_NUM_SEQS so block fan-out keeps the + # GPU's sequence slots full without flooding the queue. Tuned empirically. + SURYA_INFERENCE_MAX_INFLIGHT: int = 16 + # Per-parallel-slot KV-cache budget for the llama.cpp backend. Worst-case + # one OCR request: ~2k for image prefill + SURYA_MAX_TOKENS_FULL_PAGE + # (8192) generation + ~2k prompt/chat-template overhead ≈ 12k. Below this + # llama-server silently truncates outputs once a slot fills. + SURYA_INFERENCE_CTX_PER_SLOT: int = 12288 + # Optional override for the *total* ctx passed to llama-server. When None + # (default), total = max(16384, PARALLEL * CTX_PER_SLOT). Set this only + # if you've hand-tuned for a specific machine. + SURYA_INFERENCE_CTX_SIZE: Optional[int] = None + SURYA_INFERENCE_TIMEOUT_SECONDS: float = 600.0 + SURYA_INFERENCE_STARTUP_TIMEOUT: float = 600.0 + SURYA_INFERENCE_LOGPROBS: bool = True + SURYA_INFERENCE_MAX_RETRIES: int = 1 + # Force layout/table_rec output through a JSON schema via guided decoding. + # Eliminates malformed-JSON failures at small decode-throughput cost. + SURYA_GUIDED_LAYOUT: bool = True + # Disabled: with no minItems in TABLE_REC_JSON_SCHEMA, the constrained + # decoder closes the array after one element at temperature=0. The model + # produces well-formed JSON without the schema. + SURYA_GUIDED_TABLE_REC: bool = False + + # Token budgets + SURYA_MAX_TOKENS_LAYOUT: int = 3072 + SURYA_MAX_TOKENS_TABLE_REC: int = 3072 + SURYA_MAX_TOKENS_BLOCK_CEILING: int = 8192 + SURYA_MAX_TOKENS_FULL_PAGE: int = 6144 + SURYA_MAX_BLOCKS_PER_PAGE: int = 80 + + BBOX_SCALE: int = 1000 + + # vllm + VLLM_DOCKER_IMAGE: str = "vllm/vllm-openai:v0.20.1" + VLLM_API_KEY: str = "EMPTY" + VLLM_GPUS: str = "0" + VLLM_GPU_TYPE: str = "4090" + # bfloat16 needs an Ampere+ GPU (compute capability >= 8.0). On older cards + # (e.g. T4 / Turing) vllm refuses to start with bf16 — set float16 there. + VLLM_DTYPE: str = "bfloat16" + VLLM_MAX_MODEL_LEN: int = 18000 + VLLM_GPU_MEMORY_UTILIZATION: float = 0.85 + # MTP speculative decoding only feeds the nested-Docker spawn path (vllm.py), + # which is disabled in production; the real launch (start_single_container.sh) + # never passes --speculative-config. Benchmarked +27-36% SLOWER at OCR's 16-wide + # block concurrency (GPU already compute-bound) — see + # docs/quantization_benchmark_results.md §5b. Do not wire MTP into the launch. + VLLM_ENABLE_MTP: bool = True + VLLM_MTP_TOKENS: int = 2 + VLLM_EXTRA_ARGS: Optional[str] = None + DOCKER_HF_CACHE_PATH: str = "~/.cache/huggingface" + + # llama.cpp + LLAMA_CPP_BINARY: str = "llama-server" + LLAMA_CPP_NGL: int = 99 # all layers on GPU (Metal on macOS, CUDA on Linux GPU); harmless no-op on pure-CPU builds + LLAMA_CPP_NO_MMPROJ_OFFLOAD: bool = False + LLAMA_CPP_EXTRA_ARGS: Optional[str] = None + + # ---- Detection (kept) --------------------------------------------------- + DETECTOR_BATCH_SIZE: Optional[int] = None + DETECTOR_MODEL_CHECKPOINT: str = "s3://text_detection/2025_05_07" + DETECTOR_IMAGE_CHUNK_HEIGHT: int = 1400 + DETECTOR_TEXT_THRESHOLD: float = 0.6 + DETECTOR_BLANK_THRESHOLD: float = 0.35 + DETECTOR_POSTPROCESSING_CPU_WORKERS: int = min(8, os.cpu_count()) + DETECTOR_MIN_PARALLEL_THRESH: int = 3 + DETECTOR_BOX_Y_EXPAND_MARGIN: float = 0.05 + + # ---- OCR Error (kept) --------------------------------------------------- + OCR_ERROR_MODEL_CHECKPOINT: str = "s3://ocr_error_detection/2025_02_18" + OCR_ERROR_BATCH_SIZE: Optional[int] = None + + # ---- Debug / draw fonts (label rendering on annotated images) ---------- + RECOGNITION_RENDER_FONTS: Dict[str, str] = { + "all": os.path.join(FONT_DIR, "GoNotoCurrent-Regular.ttf"), + "zh": os.path.join(FONT_DIR, "GoNotoCJKCore.ttf"), + "ja": os.path.join(FONT_DIR, "GoNotoCJKCore.ttf"), + "ko": os.path.join(FONT_DIR, "GoNotoCJKCore.ttf"), + } + RECOGNITION_FONT_DL_BASE: str = ( + "https://github.com/satbyy/go-noto-universal/releases/download/v7.0" + ) + + @computed_field + def MODEL_DTYPE(self) -> torch.dtype: + if self.TORCH_DEVICE_MODEL == "cpu": + return torch.float32 + return torch.float16 + + @computed_field + def MODEL_DTYPE_BFLOAT(self) -> torch.dtype: + if self.TORCH_DEVICE_MODEL == "cpu": + return torch.float32 + return torch.bfloat16 + + @computed_field + def INFERENCE_MODE(self) -> Callable: + return torch.inference_mode + + class Config: + env_file = find_dotenv("local.env") + extra = "ignore" + + +settings = Settings() diff --git a/surya/table_rec/__init__.py b/surya/table_rec/__init__.py new file mode 100644 index 0000000..c52d19b --- /dev/null +++ b/surya/table_rec/__init__.py @@ -0,0 +1,217 @@ +"""TableRecPredictor: dual-path table structure recognition. + +- predict_simple: TABLE_REC_PROMPT → rows + columns only, cells derived + geometrically (row × column intersections). +- predict_full: BLOCK_PROMPT on the table crop → full HTML with + colspan / rowspan /
. The HTML lives on TableResult.html for marker to + consume directly. +""" + +from __future__ import annotations + +from typing import List, Optional + +from PIL import Image + +from surya.inference import SuryaInferenceManager, get_default_manager +from surya.inference.parsers import clean_block_html, denorm_bbox, parse_table_rec +from surya.inference.prompts import ( + PROMPT_TYPE_BLOCK, + PROMPT_TYPE_TABLE_REC, + TABLE_REC_JSON_SCHEMA, +) +from surya.inference.schema import BatchInputItem +from surya.inference.util import image_token_budget +from surya.logging import get_logger +from surya.settings import settings +from surya.table_rec.schema import TableCell, TableCol, TableResult, TableRow + +logger = get_logger() + + +def _polygon_from_bbox(bbox): + x0, y0, x1, y1 = bbox + return [[x0, y0], [x1, y0], [x1, y1], [x0, y1]] + + +def _intersect_bbox(a, b): + x0 = max(a[0], b[0]) + y0 = max(a[1], b[1]) + x1 = min(a[2], b[2]) + y1 = min(a[3], b[3]) + if x1 <= x0 or y1 <= y0: + return None + return (x0, y0, x1, y1) + + +class TableRecPredictor: + def __init__(self, manager: Optional[SuryaInferenceManager] = None): + self.manager = manager + self._disable_tqdm = settings.DISABLE_TQDM + + @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 to(self, *args, **kwargs): + return + + def __call__( + self, images: List[Image.Image], mode: str = "simple" + ) -> List[TableResult]: + if mode == "full": + return self.predict_full(images) + return self.predict_simple(images) + + def predict_simple(self, images: List[Image.Image]) -> List[TableResult]: + if not images: + return [] + manager = self.manager or get_default_manager() + guided = TABLE_REC_JSON_SCHEMA if settings.SURYA_GUIDED_TABLE_REC else None + batch = [ + BatchInputItem( + image=img, + prompt_type=PROMPT_TYPE_TABLE_REC, + max_tokens=settings.SURYA_MAX_TOKENS_TABLE_REC, + guided_json=guided, + ) + for img in images + ] + outputs = manager.generate(batch) + + results: List[TableResult] = [] + for img, out in zip(images, outputs): + w, h = img.size + page_bbox = [0, 0, float(w), float(h)] + if out.error or not out.raw: + results.append( + TableResult( + rows=[], + cols=[], + cells=[], + image_bbox=page_bbox, + raw=out.raw, + mode="simple", + error=True, + ) + ) + continue + try: + elements = parse_table_rec(out.raw) + except Exception as e: + logger.warning( + f"Table rec parse failed: {e}; raw[:200]={out.raw[:200]!r}" + ) + results.append( + TableResult( + rows=[], + cols=[], + cells=[], + image_bbox=page_bbox, + raw=out.raw, + mode="simple", + error=True, + ) + ) + continue + + rows: List[TableRow] = [] + cols: List[TableCol] = [] + for el in elements: + pixel_bbox = denorm_bbox(el.bbox, w, h, scale=settings.BBOX_SCALE) + poly = _polygon_from_bbox(pixel_bbox) + if el.label == "Row": + rows.append(TableRow(polygon=poly, row_id=len(rows))) + else: + cols.append(TableCol(polygon=poly, col_id=len(cols))) + + # Derive cells geometrically (row × column intersections) + cells: List[TableCell] = [] + cell_id = 0 + for row in rows: + for col in cols: + inter = _intersect_bbox(row.bbox, col.bbox) + if inter is None: + continue + cells.append( + TableCell( + polygon=_polygon_from_bbox(inter), + row_id=row.row_id, + col_id=col.col_id, + cell_id=cell_id, + ) + ) + cell_id += 1 + results.append( + TableResult( + rows=rows, + cols=cols, + cells=cells, + image_bbox=page_bbox, + raw=out.raw, + mode="simple", + error=False, + ) + ) + return results + + def predict_full( + self, images: List[Image.Image], counts: Optional[List[int]] = None + ) -> List[TableResult]: + """Full-HTML path: BLOCK_PROMPT on table crops. Use when complex + structure (spanning cells, headers) matters and ground-truth-style + HTML is preferred. `counts` (one per image) shapes max_tokens.""" + if not images: + return [] + manager = self.manager or get_default_manager() + if counts is None: + counts = [0] * len(images) + batch = [] + for img, count in zip(images, counts): + batch.append( + BatchInputItem( + image=img, + prompt_type=PROMPT_TYPE_BLOCK, + max_tokens=image_token_budget( + count, + ceiling=settings.SURYA_MAX_TOKENS_BLOCK_CEILING, + floor=1024, + ), + ) + ) + outputs = manager.generate(batch) + results: List[TableResult] = [] + for img, out in zip(images, outputs): + w, h = img.size + page_bbox = [0, 0, float(w), float(h)] + if out.error: + results.append( + TableResult( + rows=[], + cols=[], + cells=[], + image_bbox=page_bbox, + raw=out.raw, + mode="full", + error=True, + ) + ) + continue + html = clean_block_html(out.raw) + results.append( + TableResult( + rows=[], + cols=[], + cells=[], + image_bbox=page_bbox, + raw=out.raw, + html=html, + mode="full", + error=False, + ) + ) + return results diff --git a/surya/table_rec/schema.py b/surya/table_rec/schema.py new file mode 100644 index 0000000..d1cefdc --- /dev/null +++ b/surya/table_rec/schema.py @@ -0,0 +1,48 @@ +from typing import List, Optional + +from pydantic import BaseModel + +from surya.common.polygon import PolygonBox + + +class TableRow(PolygonBox): + row_id: int + + @property + def label(self) -> str: + return f"Row {self.row_id}" + + +class TableCol(PolygonBox): + col_id: int + + @property + def label(self) -> str: + return f"Column {self.col_id}" + + +class TableCell(PolygonBox): + """Geometric cell derived from row × column intersection. + + The simple-path TableRecPredictor doesn't return spanning info from the + model — colspan/rowspan/header come from the full-path HTML output if + needed.""" + + row_id: int + col_id: int + cell_id: int + + @property + def label(self) -> str: + return f"Cell {self.cell_id}" + + +class TableResult(BaseModel): + rows: List[TableRow] + cols: List[TableCol] + cells: List[TableCell] + image_bbox: List[float] + raw: Optional[str] = None # raw model output + html: Optional[str] = None # populated when full-path was used + mode: str = "simple" # "simple" | "full" + error: bool = False diff --git a/surya/timing.py b/surya/timing.py new file mode 100644 index 0000000..5c7b39c --- /dev/null +++ b/surya/timing.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +import contextvars +import logging +import os +import threading +import time +from contextlib import contextmanager +from typing import Any, Iterator + + +_current_timing: contextvars.ContextVar["TimingCollector | None"] = contextvars.ContextVar( + "surya_current_timing", + default=None, +) + + +def timing_enabled() -> bool: + return os.getenv("SUYA_TIMING_ENABLED", "true").strip().lower() not in { + "0", + "false", + "no", + "off", + } + + +class TimingCollector: + def __init__(self, *, request_id: str = "-", batch_size: int = 0) -> None: + self.request_id = request_id + self.batch_size = batch_size + self._lock = threading.Lock() + self._events: list[dict[str, Any]] = [] + + def record(self, name: str, duration_ms: float, **metadata: Any) -> None: + if not timing_enabled(): + return + event = { + "name": name, + "duration_ms": round(duration_ms, 2), + } + if metadata: + event["metadata"] = { + key: value + for key, value in metadata.items() + if value is not None + } + with self._lock: + self._events.append(event) + + def summary(self) -> list[dict[str, Any]]: + with self._lock: + return list(self._events) + + +def get_current_timing() -> TimingCollector | None: + return _current_timing.get() + + +def set_current_timing(collector: TimingCollector | None): + return _current_timing.set(collector) + + +def reset_current_timing(token) -> None: + _current_timing.reset(token) + + +@contextmanager +def timing_span(name: str, **metadata: Any) -> Iterator[None]: + collector = get_current_timing() + if collector is None or not timing_enabled(): + yield + return + start = time.perf_counter() + try: + yield + finally: + collector.record(name, (time.perf_counter() - start) * 1000, **metadata) + + +def log_timing_summary( + logger: logging.Logger, + collector: TimingCollector, + *, + message: str = "surya_timing_summary", +) -> None: + if not timing_enabled(): + return + logger.info( + "%s request_id=%s batch_size=%s events=%s", + message, + collector.request_id, + collector.batch_size, + collector.summary(), + ) diff --git a/tests/test_cer_divergence.py b/tests/test_cer_divergence.py new file mode 100644 index 0000000..2970db2 --- /dev/null +++ b/tests/test_cer_divergence.py @@ -0,0 +1,18 @@ +from scripts.cer_divergence import cer + + +def test_identical_text_is_zero(): + assert cer("hello world", "hello world") == 0.0 + + +def test_single_substitution(): + # 1 edit over 5 reference chars + assert cer("hello", "hallo") == 0.2 + + +def test_empty_reference_with_output_is_one(): + assert cer("", "abc") == 1.0 + + +def test_empty_both_is_zero(): + assert cer("", "") == 0.0 diff --git a/tests/test_endpoint_routes.py b/tests/test_endpoint_routes.py new file mode 100644 index 0000000..956c25d --- /dev/null +++ b/tests/test_endpoint_routes.py @@ -0,0 +1,21 @@ +def test_app_exposes_all_legacy_and_openai_routes(): + from surya.endpoint.app import app + + paths = {getattr(r, "path", None) for r in app.routes} + expected = { + "/home", + "/v1/api/ai/suya_ocr", + "/v1/api/ai/suya_ocr/", + "/v1/api/ai/suya_ocr_vllm", + "/v1/api/ai/suya_ocr_vllm/", + "/v1/api/ai/suya_ocr_vllm/health", + "/v1/api/ai/suya_text_det", + "/v1/api/ai/suya_text_det/", + "/v1/api/ai/suya_layout_det", + "/v1/api/ai/suya_layout_det/", + "/v1/api/ai/suya_table_rec", + "/v1/api/ai/suya_table_rec/", + "/image2text", + "/v1/chat/completions", + } + assert expected.issubset(paths), f"missing routes: {expected - paths}" diff --git a/tests/test_openai_endpoint.py b/tests/test_openai_endpoint.py new file mode 100644 index 0000000..b0543ec --- /dev/null +++ b/tests/test_openai_endpoint.py @@ -0,0 +1,96 @@ +import base64 +import io +from unittest.mock import patch + +from fastapi.testclient import TestClient +from PIL import Image + + +def _png_data_url() -> str: + buf = io.BytesIO() + Image.new("RGB", (8, 8), "white").save(buf, format="PNG") + b64 = base64.b64encode(buf.getvalue()).decode("ascii") + return f"data:image/png;base64,{b64}" + + +def _client() -> TestClient: + from surya.endpoint.app import app + + return TestClient(app) + + +def _image_message(): + return [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": _png_data_url()}}, + {"type": "text", "text": "extract the text"}, + ], + } + ] + + +def test_returns_openai_chat_completion_shape(): + body = {"model": "surya-ocr", "messages": _image_message()} + fake = {"text_lines": "hello\nworld", "ocr_text_json": {"blocks": []}, "elapsed_seconds": 1.23} + with patch("surya.endpoint.service.ocr_via_batcher", return_value=fake): + r = _client().post("/v1/chat/completions", json=body) + assert r.status_code == 200 + d = r.json() + assert d["object"] == "chat.completion" + assert d["model"] == "surya-ocr" + assert d["choices"][0]["message"]["content"] == "hello\nworld" + assert d["choices"][0]["finish_reason"] == "stop" + assert d["surya"]["ocr_text_json"] == {"blocks": []} + assert d["surya"]["elapsed_seconds"] == 1.23 + assert "usage" in d + + +def test_mode_full_page_skips_text_detection(): + body = {"model": "m", "mode": "full_page", "messages": _image_message()} + fake = {"text_lines": "x", "ocr_text_json": {}, "elapsed_seconds": None} + with patch("surya.endpoint.service.ocr_via_batcher", return_value=fake) as m: + r = _client().post("/v1/chat/completions", json=body) + assert r.status_code == 200 + assert m.call_args.kwargs["skip_text_detection"] is True + + +def test_mode_table_routes_to_table_image(): + body = {"model": "m", "mode": "table", "messages": _image_message()} + fake = {"text_lines": "tbl", "ocr_text_json": {"x": 1}, "elapsed_seconds": None} + with patch("surya.endpoint.service.table_image", return_value=fake) as m: + r = _client().post("/v1/chat/completions", json=body) + assert r.status_code == 200 + assert m.called + assert r.json()["choices"][0]["message"]["content"] == "tbl" + + +def test_ocr_with_boxes_false_omits_structured_json(): + body = {"model": "m", "ocr_with_boxes": False, "messages": _image_message()} + fake = {"text_lines": "x", "ocr_text_json": {"a": 1}, "elapsed_seconds": None} + with patch("surya.endpoint.service.ocr_via_batcher", return_value=fake): + r = _client().post("/v1/chat/completions", json=body) + assert r.json()["surya"]["ocr_text_json"] is None + + +def test_stream_true_returns_400(): + body = {"model": "m", "stream": True, "messages": _image_message()} + r = _client().post("/v1/chat/completions", json=body) + assert r.status_code == 400 + assert r.json()["error"]["type"] == "invalid_request_error" + + +def test_missing_image_returns_400(): + body = {"model": "m", "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]} + r = _client().post("/v1/chat/completions", json=body) + assert r.status_code == 400 + assert r.json()["error"]["type"] == "invalid_request_error" + + +def test_pipeline_error_returns_500(): + body = {"model": "m", "messages": _image_message()} + with patch("surya.endpoint.service.ocr_via_batcher", side_effect=RuntimeError("boom")): + r = _client().post("/v1/chat/completions", json=body) + assert r.status_code == 500 + assert r.json()["error"]["type"] == "internal_error" diff --git a/tests/test_parse_timing.py b/tests/test_parse_timing.py new file mode 100644 index 0000000..47524ff --- /dev/null +++ b/tests/test_parse_timing.py @@ -0,0 +1,30 @@ +from scripts.parse_timing import parse_line, aggregate + +SAMPLE = ( + "2026-06-11 14:35:00 - vllm_batcher - _process_jobs - line:119 - INFO - " + "surya_timing_summary request_id=a,b batch_size=2 events=[" + "{'name': 'openai_chat_completion', 'duration_ms': 100.0, 'metadata': {'token_count': 40}}, " + "{'name': 'openai_chat_completion', 'duration_ms': 300.0, 'metadata': {'token_count': 60}}, " + "{'name': 'recognition_manager_generate', 'duration_ms': 450.0}]" +) + + +def test_parse_line_extracts_events_and_batch_size(): + rec = parse_line(SAMPLE) + assert rec["batch_size"] == 2 + assert len(rec["events"]) == 3 + assert rec["events"][0]["metadata"]["token_count"] == 40 + + +def test_parse_line_returns_none_for_unrelated_line(): + assert parse_line("2026-06-11 - foo - bar - line:1 - INFO - request_start id=x") is None + + +def test_aggregate_counts_and_sums_by_span_name(): + rec = parse_line(SAMPLE) + agg = aggregate([rec]) + chat = agg["openai_chat_completion"] + assert chat["count"] == 2 + assert chat["total_ms"] == 400.0 + assert chat["mean_ms"] == 200.0 + assert chat["total_tokens"] == 100 diff --git a/tests/test_quant_aggregate.py b/tests/test_quant_aggregate.py new file mode 100644 index 0000000..77c3877 --- /dev/null +++ b/tests/test_quant_aggregate.py @@ -0,0 +1,34 @@ +from scripts.quant.aggregate import SUMMARY_FIELDS, build_row, rows_to_csv, rows_to_markdown + + +def test_build_row_defaults_unset_fields_to_none(): + row = build_row(method="awq", status="ok", t4_deployable=True, mean_cer=0.01) + assert row["method"] == "awq" + assert row["status"] == "ok" + assert row["t4_deployable"] is True + assert row["mean_cer"] == 0.01 + assert row["throughput_rps"] is None + assert set(row.keys()) == set(SUMMARY_FIELDS) + + +def test_build_row_rejects_unknown_field(): + try: + build_row(method="awq", bogus=1) + assert False, "expected KeyError" + except KeyError as exc: + assert "bogus" in str(exc) + + +def test_rows_to_csv_has_header_and_order(): + rows = [build_row(method="bf16", status="ok")] + csv_text = rows_to_csv(rows) + assert csv_text.splitlines()[0] == ",".join(SUMMARY_FIELDS) + assert "bf16" in csv_text.splitlines()[1] + + +def test_rows_to_markdown_renders_failed_row(): + rows = [build_row(method="gptq", status="failed", error="OOM at load")] + md = rows_to_markdown(rows) + assert "| gptq |" in md + assert "failed" in md + assert "OOM at load" in md diff --git a/tests/test_quant_bbox_iou.py b/tests/test_quant_bbox_iou.py new file mode 100644 index 0000000..54c537d --- /dev/null +++ b/tests/test_quant_bbox_iou.py @@ -0,0 +1,48 @@ +import json + +from scripts.quant.bbox_iou import iou, match_boxes, bbox_iou_over_dirs + + +def test_iou_identical_is_one(): + assert iou([0, 0, 10, 10], [0, 0, 10, 10]) == 1.0 + + +def test_iou_disjoint_is_zero(): + assert iou([0, 0, 10, 10], [20, 20, 30, 30]) == 0.0 + + +def test_iou_half_overlap(): + # two 10x10 boxes overlapping in a 10x5 region -> 50/150 + assert abs(iou([0, 0, 10, 10], [0, 5, 10, 15]) - (50 / 150)) < 1e-9 + + +def test_match_boxes_counts_missed_and_extra(): + ref = [[0, 0, 10, 10], [100, 100, 110, 110]] + cand = [[0, 0, 10, 10], [200, 200, 210, 210], [300, 300, 310, 310]] + + result = match_boxes(ref, cand, iou_threshold=0.5) + + assert result["matched"] == 1 + assert result["missed"] == 1 # ref box at 100,100 unmatched + assert result["extra"] == 2 # two cand boxes unmatched + assert abs(result["mean_matched_iou"] - 1.0) < 1e-9 + + +def test_match_boxes_empty(): + result = match_boxes([], [], iou_threshold=0.5) + assert result == {"matched": 0, "missed": 0, "extra": 0, "mean_matched_iou": 0.0} + + +def test_bbox_iou_over_dirs(tmp_path): + ref_dir = tmp_path / "ref" + cand_dir = tmp_path / "cand" + ref_dir.mkdir() + cand_dir.mkdir() + (ref_dir / "p1.json").write_text(json.dumps({"boxes": [[0, 0, 10, 10]]})) + (cand_dir / "p1.json").write_text(json.dumps({"boxes": [[0, 0, 10, 10]]})) + + result = bbox_iou_over_dirs(ref_dir, cand_dir, iou_threshold=0.5) + + assert abs(result["mean_bbox_iou"] - 1.0) < 1e-9 + assert result["mean_missed_lines"] == 0.0 + assert result["mean_extra_lines"] == 0.0 diff --git a/tests/test_quant_build_model.py b/tests/test_quant_build_model.py new file mode 100644 index 0000000..8a1acd6 --- /dev/null +++ b/tests/test_quant_build_model.py @@ -0,0 +1,39 @@ +from pathlib import Path + +import scripts.quant.build_model as bm + + +def test_baseline_and_online_need_no_build(tmp_path): + assert bm.needs_build("bf16") is False + assert bm.needs_build("fp8") is False + assert bm.needs_build("awq") is True + assert bm.needs_build("bnb4") is True + + +def test_build_is_idempotent_when_output_exists(tmp_path, monkeypatch): + out = tmp_path / "awq" + out.mkdir() + (out / "config.json").write_text("{}", encoding="utf-8") + called = {"n": 0} + + def fake_compressor(*a, **k): + called["n"] += 1 + + monkeypatch.setattr(bm, "_build_compressor", fake_compressor) + result = bm.build_model("awq", base_model="datalab-to/surya-ocr-2", out_dir=out, calib_images=[]) + assert result == out + assert called["n"] == 0 # skipped because config.json already present + + +def test_build_dispatches_to_compressor(tmp_path, monkeypatch): + out = tmp_path / "gptq" + seen = {} + + def fake_compressor(method, base_model, out_dir, calib_images): + seen["method"] = method + out_dir.mkdir(parents=True, exist_ok=True) + (out_dir / "config.json").write_text("{}", encoding="utf-8") + + monkeypatch.setattr(bm, "_build_compressor", fake_compressor) + bm.build_model("gptq", base_model="b", out_dir=out, calib_images=[]) + assert seen["method"] == "gptq" diff --git a/tests/test_quant_capture.py b/tests/test_quant_capture.py new file mode 100644 index 0000000..fa7782b --- /dev/null +++ b/tests/test_quant_capture.py @@ -0,0 +1,41 @@ +import json + +from scripts.quant.capture import extract_capture, write_capture + + +def test_extract_capture_pulls_text_boxes_and_latency(): + body = { + "data": { + "text_lines": "line one\nline two", + "elapsed_seconds": 5.4, + "ocr_text_json": { + "blocks": [ + {"bbox": [1, 2, 3, 4], "html": "line one"}, + {"bbox": [5, 6, 7, 8], "html": "line two"}, + ] + }, + }, + "message": "success", + "code": 200, + } + + result = extract_capture(body) + + assert result["text"] == "line one\nline two" + assert result["boxes"] == [[1, 2, 3, 4], [5, 6, 7, 8]] + assert result["elapsed_seconds"] == 5.4 + + +def test_extract_capture_tolerates_missing_fields(): + result = extract_capture({"data": {}}) + assert result == {"text": "", "boxes": [], "elapsed_seconds": None} + + +def test_write_capture_writes_txt_and_json(tmp_path): + cap = {"text": "hello", "boxes": [[0, 0, 1, 1]], "elapsed_seconds": 2.0} + + write_capture(cap, tmp_path, "page1") + + assert (tmp_path / "page1.txt").read_text(encoding="utf-8") == "hello" + loaded = json.loads((tmp_path / "page1.json").read_text(encoding="utf-8")) + assert loaded == {"boxes": [[0, 0, 1, 1]], "elapsed_seconds": 2.0} diff --git a/tests/test_quant_manifest.py b/tests/test_quant_manifest.py new file mode 100644 index 0000000..b26bbab --- /dev/null +++ b/tests/test_quant_manifest.py @@ -0,0 +1,29 @@ +from pathlib import Path + +from scripts.quant.manifest import load_manifest + + +def test_load_manifest_skips_blanks_and_comments(tmp_path): + repo = tmp_path / "repo" + repo.mkdir() + (repo / "a.png").write_bytes(b"x") + (repo / "b.jpg").write_bytes(b"y") + manifest = tmp_path / "manifest.txt" + manifest.write_text("# header\n\na.png\nb.jpg\n", encoding="utf-8") + + result = load_manifest(manifest, repo) + + assert result == [repo / "a.png", repo / "b.jpg"] + + +def test_load_manifest_raises_on_missing_image(tmp_path): + repo = tmp_path / "repo" + repo.mkdir() + manifest = tmp_path / "manifest.txt" + manifest.write_text("missing.png\n", encoding="utf-8") + + try: + load_manifest(manifest, repo) + assert False, "expected FileNotFoundError" + except FileNotFoundError as exc: + assert "missing.png" in str(exc) diff --git a/tests/test_quant_plot.py b/tests/test_quant_plot.py new file mode 100644 index 0000000..deb96c2 --- /dev/null +++ b/tests/test_quant_plot.py @@ -0,0 +1,28 @@ +from scripts.quant.aggregate import build_row +from scripts.quant.plot import pareto_points, render_all + + +def test_pareto_points_skips_failed_and_missing(): + rows = [ + build_row(method="bf16", status="ok", mean_latency_s=6.0, mean_cer=0.0, t4_deployable=True), + build_row(method="awq", status="ok", mean_latency_s=4.0, mean_cer=0.01, t4_deployable=True), + build_row(method="gptq", status="failed", t4_deployable=True), + build_row(method="fp8", status="ok", mean_latency_s=3.0, mean_cer=0.02, t4_deployable=False), + ] + points = pareto_points(rows, x_key="mean_latency_s", y_key="mean_cer") + methods = {p["method"] for p in points} + assert methods == {"bf16", "awq", "fp8"} # gptq failed -> skipped + awq = next(p for p in points if p["method"] == "awq") + assert awq["x"] == 4.0 and awq["y"] == 0.01 and awq["t4_deployable"] is True + + +def test_render_all_writes_png_files(tmp_path): + rows = [ + build_row(method="bf16", status="ok", mean_latency_s=6.0, mean_cer=0.0, + mean_bbox_iou=1.0, t4_deployable=True), + build_row(method="awq", status="ok", mean_latency_s=4.0, mean_cer=0.01, + mean_bbox_iou=0.98, t4_deployable=True), + ] + written = render_all(rows, tmp_path) + assert all(p.exists() for p in written) + assert any(p.name == "pareto_latency_cer.png" for p in written) diff --git a/tests/test_quant_recipes.py b/tests/test_quant_recipes.py new file mode 100644 index 0000000..b3354c3 --- /dev/null +++ b/tests/test_quant_recipes.py @@ -0,0 +1,37 @@ +from scripts.quant.recipes import METHOD_SPECS, method_names, vllm_serve_args + + +def test_all_seven_methods_present(): + assert method_names() == ["bf16", "fp8", "int8", "awq", "gptq", "bnb8", "bnb4"] + + +def test_t4_flags(): + assert METHOD_SPECS["fp8"]["t4_deployable"] is False + assert METHOD_SPECS["awq"]["t4_deployable"] is True + assert METHOD_SPECS["bnb8"]["t4_deployable"] is True + + +def test_baseline_serves_base_model(): + args = vllm_serve_args("bf16", model_path="/m/base", base_model="/m/base", port=8001) + assert "--model" in args and "/m/base" in args + assert "--quantization" not in args + assert args[args.index("--port") + 1] == "8001" + + +def test_fp8_uses_online_quantization_on_base(): + args = vllm_serve_args("fp8", model_path="/m/base", base_model="/m/base", port=8001) + assert args[args.index("--quantization") + 1] == "fp8" + assert "/m/base" in args + + +def test_compressor_serves_built_path_without_quant_flag(): + args = vllm_serve_args("awq", model_path="/m/awq", base_model="/m/base", port=8001) + assert "/m/awq" in args + assert "--quantization" not in args # quant config travels with the compressed checkpoint + + +def test_bnb_uses_bitsandbytes_flags(): + args = vllm_serve_args("bnb4", model_path="/m/bnb4", base_model="/m/base", port=8001) + assert args[args.index("--quantization") + 1] == "bitsandbytes" + assert args[args.index("--load-format") + 1] == "bitsandbytes" + assert "/m/bnb4" in args diff --git a/tests/test_quant_run_all.py b/tests/test_quant_run_all.py new file mode 100644 index 0000000..138712e --- /dev/null +++ b/tests/test_quant_run_all.py @@ -0,0 +1,29 @@ +import scripts.quant.run_all as ra +from scripts.quant.aggregate import SUMMARY_FIELDS + + +def test_run_method_failure_becomes_failed_row(monkeypatch): + def boom(*a, **k): + raise RuntimeError("OOM at load") + + monkeypatch.setattr(ra, "_measure_method", boom) + row = ra.run_method("gptq", base_model="b", work_dir=ra.Path("/tmp/x"), + eval_images=[], reference_dir=ra.Path("/tmp/ref")) + assert row["method"] == "gptq" + assert row["status"] == "failed" + assert "OOM at load" in row["error"] + assert set(row.keys()) == set(SUMMARY_FIELDS) + + +def test_run_method_success_passes_through_metrics(monkeypatch): + def fake_measure(method, base_model, work_dir, eval_images, reference_dir): + return {"mean_cer": 0.01, "mean_bbox_iou": 0.98, "mean_latency_s": 4.2, + "model_size_mb": 512.0} + + monkeypatch.setattr(ra, "_measure_method", fake_measure) + row = ra.run_method("awq", base_model="b", work_dir=ra.Path("/tmp/x"), + eval_images=[], reference_dir=ra.Path("/tmp/ref")) + assert row["status"] == "ok" + assert row["mean_cer"] == 0.01 + assert row["t4_deployable"] is True + assert row["mean_latency_s"] == 4.2 diff --git a/tests/test_quant_serve.py b/tests/test_quant_serve.py new file mode 100644 index 0000000..f7bce62 --- /dev/null +++ b/tests/test_quant_serve.py @@ -0,0 +1,16 @@ +from scripts.quant.serve import health_url, parse_listening_pid + + +def test_health_url_from_port(): + assert health_url(8001) == "http://127.0.0.1:8001/health" + + +def test_parse_listening_pid_extracts_from_ss_line(): + ss_output = ( + 'LISTEN 0 4096 127.0.0.1:8001 0.0.0.0:* users:(("python",pid=12345,fd=7))\n' + ) + assert parse_listening_pid(ss_output, 8001) == 12345 + + +def test_parse_listening_pid_returns_none_when_absent(): + assert parse_listening_pid("LISTEN 0 4096 127.0.0.1:9999 0.0.0.0:*\n", 8001) is None diff --git a/tests/test_resolve_max_workers.py b/tests/test_resolve_max_workers.py new file mode 100644 index 0000000..7827afe --- /dev/null +++ b/tests/test_resolve_max_workers.py @@ -0,0 +1,13 @@ +from surya.inference.backends.openai_client import resolve_max_workers + + +def test_caps_at_max_inflight(): + assert resolve_max_workers(batch_len=160, max_inflight=16) == 16 + + +def test_uses_batch_len_when_smaller(): + assert resolve_max_workers(batch_len=4, max_inflight=16) == 4 + + +def test_never_below_one(): + assert resolve_max_workers(batch_len=0, max_inflight=16) == 1 diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..9596cc8 --- /dev/null +++ b/tools.py @@ -0,0 +1,242 @@ +import io +import tempfile +from typing import List + +import pypdfium2 + +from surya.debug.draw import draw_polys_on_image, draw_bboxes_on_image + +from PIL import Image +from surya.settings import settings +from vllm_tools import predictors_vllm, ocr_vllm + + +predictors = predictors_vllm + + +def rescale_bbox(bbox, source_size, target_size): + width_ratio = target_size[0] / source_size[0] + height_ratio = target_size[1] / source_size[1] + return [ + bbox[0] * width_ratio, + bbox[1] * height_ratio, + bbox[2] * width_ratio, + bbox[3] * height_ratio, + ] + + +def expand_bbox(bbox, margin=5): + return [ + max(0, int(bbox[0]) - margin), + max(0, int(bbox[1]) - margin), + int(bbox[2]) + margin, + int(bbox[3]) + margin, + ] + + + +def page_counter(pdf_file): + doc = open_pdf(pdf_file) + doc_len = len(doc) + doc.close() + return doc_len + +def ocr_errors(pdf_file, page_count, sample_len=512, max_samples=10, max_pages=15): + from pdftext.extraction import plain_text_output + + with tempfile.NamedTemporaryFile(suffix=".pdf") as f: + f.write(pdf_file.getvalue()) + f.seek(0) + + # Sample the text from the middle of the PDF + page_middle = page_count // 2 + page_range = range( + max(page_middle - max_pages, 0), min(page_middle + max_pages, page_count) + ) + text = plain_text_output(f.name, page_range=page_range) + + sample_gap = len(text) // max_samples + if len(text) == 0 or sample_gap == 0: + return "This PDF has no text or very little text", ["no text"] + + if sample_gap < sample_len: + sample_gap = sample_len + + # Split the text into samples for the model + samples = [] + for i in range(0, len(text), sample_gap): + samples.append(text[i : i + sample_len]) + + results = predictors["ocr_error"](samples) + label = "This PDF has good text." + if results.labels.count("bad") / len(results.labels) > 0.2: + label = "This PDF may have garbled or bad OCR text." + return label, results.labels + + +def text_detection(img): + text_pred = predictors["detection"]([img])[0] + text_polygons = [p.polygon for p in text_pred.bboxes] + det_img = draw_polys_on_image(text_polygons, img.copy()) + return det_img, text_pred + + +def layout_detection(img): + pred = predictors["layout"]([img])[0] + polygons = [p.polygon for p in pred.bboxes] + labels = [ + f"{p.label}-{p.position}-{round(p.top_k[p.label], 2)}" for p in pred.bboxes + ] + layout_img = draw_polys_on_image( + polygons, img.copy(), labels=labels, label_font_size=18 + ) + return layout_img, pred + + +def table_recognition( + img, highres_img, skip_table_detection: bool +): + if skip_table_detection: + layout_tables = [(0, 0, highres_img.size[0], highres_img.size[1])] + table_imgs = [highres_img] + else: + _, layout_pred = layout_detection(img) + layout_tables_lowres = [ + line.bbox + for line in layout_pred.bboxes + if line.label in ["Table", "TableOfContents"] + ] + table_imgs = [] + layout_tables = [] + for tb in layout_tables_lowres: + highres_bbox = rescale_bbox(tb, img.size, highres_img.size) + # Slightly expand the box + highres_bbox = expand_bbox(highres_bbox) + table_imgs.append(highres_img.crop(highres_bbox)) + layout_tables.append(highres_bbox) + + table_preds = predictors["table_rec"](table_imgs) + table_img = img.copy() + + for results, table_bbox in zip(table_preds, layout_tables): + adjusted_bboxes = [] + labels = [] + colors = [] + + for item in results.cells: + adjusted_bboxes.append( + [ + (item.bbox[0] + table_bbox[0]), + (item.bbox[1] + table_bbox[1]), + (item.bbox[2] + table_bbox[0]), + (item.bbox[3] + table_bbox[1]), + ] + ) + labels.append(item.label) + if "Row" in item.label: + colors.append("blue") + else: + colors.append("red") + table_img = draw_bboxes_on_image( + adjusted_bboxes, + highres_img, + labels=labels, + label_font_size=18, + color=colors, + ) + return table_img, table_preds + + +# Function for OCR +def ocr( + img: Image.Image, + highres_img: Image.Image, + skip_text_detection: bool = False, + recognize_math: bool = True, + with_bboxes: bool = True, +): + return ocr_vllm( + img, + highres_img, + skip_text_detection=skip_text_detection, + recognize_math=recognize_math, + with_bboxes=with_bboxes, + ) + + +def open_pdf(pdf_file): + stream = io.BytesIO(pdf_file.getvalue()) + return pypdfium2.PdfDocument(stream) + + +def get_page_image(pdf_file, page_num, dpi=settings.IMAGE_DPI): + doc = open_pdf(pdf_file) + renderer = doc.render( + pypdfium2.PdfBitmap.to_pil, + page_indices=[page_num - 1], + scale=dpi / 72, + ) + png = list(renderer)[0] + png_image = png.convert("RGB") + doc.close() + return png_image + + +def page_counter(pdf_file): + doc = open_pdf(pdf_file) + doc_len = len(doc) + doc.close() + return doc_len + + +import pandas as pd + +def bbox_intersection(box1, box2): + x1 = max(box1[0], box2[0]) + y1 = max(box1[1], box2[1]) + x2 = min(box1[2], box2[2]) + y2 = min(box1[3], box2[3]) + + if x1 < x2 and y1 < y2: + return (x1, y1, x2, y2) + else: + return None + +def area_of_bbox(box): + return (box[2] - box[0]) * (box[3] - box[1]) + +def is_bbox_inside(box, parent_box): + interaction_box = bbox_intersection(box, parent_box) + if interaction_box is None: + return False + return area_of_bbox(interaction_box) / area_of_bbox(box) > 0.5 + +def center_of_bbox(box): + return ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2) + + +def extract_text_from_image(image): + layout_predictor = predictors["layout"] + recognition_predictor = predictors["recognition"] + + layout_prediction = layout_predictor([image])[0] + prediction = recognition_predictor([image], [layout_prediction], full_page=False)[0] + + items = [ + { + "text": block.html, + "position": block.reading_order, + "order_value": block.bbox[1], + } + for block in prediction.blocks + if not block.skipped and not block.error and block.html + ] + if not items: + return "" + + df = pd.DataFrame(items) + df = df.sort_values(by=['position', 'order_value']) + ds = df.groupby('position').apply(lambda x: " ".join(x['text'].tolist())) + full_text = "\n\n".join(ds.to_list()) + + return full_text diff --git a/vllm_batcher.py b/vllm_batcher.py new file mode 100644 index 0000000..58690a5 --- /dev/null +++ b/vllm_batcher.py @@ -0,0 +1,134 @@ +import logging +import os +import queue +import threading +import time +from concurrent.futures import Future +from dataclasses import dataclass +from typing import Tuple + +from PIL import Image + +from surya.timing import ( + TimingCollector, + log_timing_summary, + reset_current_timing, + set_current_timing, + timing_span, +) +from vllm_tools import ocr_vllm_batch + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class OcrOptions: + skip_text_detection: bool + recognize_math: bool + with_bboxes: bool + + +@dataclass +class OcrJob: + image: Image.Image + highres_image: Image.Image + options: OcrOptions + future: Future + request_id: str + + +class VllmOcrBatcher: + def __init__(self) -> None: + self.max_batch_size = int(os.getenv("SUYA_MAX_BATCH_SIZE", "8")) + self.batch_wait_ms = float(os.getenv("SUYA_BATCH_WAIT_MS", "25")) + self.queue_timeout_seconds = float(os.getenv("SUYA_BATCH_QUEUE_TIMEOUT_SECONDS", "900")) + self._queue: queue.Queue[OcrJob] = queue.Queue( + maxsize=int(os.getenv("SUYA_MAX_QUEUE_SIZE", "128")) + ) + self._thread = threading.Thread(target=self._run, name="vllm-ocr-batcher", daemon=True) + self._thread.start() + + def submit( + self, + image: Image.Image, + highres_image: Image.Image, + *, + skip_text_detection: bool, + recognize_math: bool, + with_bboxes: bool, + request_id: str, + ): + future: Future = Future() + job = OcrJob( + image=image, + highres_image=highres_image, + options=OcrOptions(skip_text_detection, recognize_math, with_bboxes), + future=future, + request_id=request_id, + ) + with timing_span("batcher_queue_put", request_id=request_id): + self._queue.put(job, timeout=self.queue_timeout_seconds) + with timing_span("batcher_wait_result", request_id=request_id): + return future.result(timeout=self.queue_timeout_seconds) + + def _run(self) -> None: + while True: + first = self._queue.get() + jobs = [first] + deadline = time.perf_counter() + (self.batch_wait_ms / 1000.0) + while len(jobs) < self.max_batch_size: + remaining = deadline - time.perf_counter() + if remaining <= 0: + break + try: + jobs.append(self._queue.get(timeout=remaining)) + except queue.Empty: + break + self._process_jobs(jobs) + + def _process_jobs(self, jobs: list[OcrJob]) -> None: + groups: dict[OcrOptions, list[OcrJob]] = {} + for job in jobs: + groups.setdefault(job.options, []).append(job) + + for options, grouped_jobs in groups.items(): + start = time.perf_counter() + collector = TimingCollector( + request_id=",".join(job.request_id for job in grouped_jobs), + batch_size=len(grouped_jobs), + ) + token = set_current_timing(collector) + try: + with timing_span("batcher_process_jobs", batch_size=len(grouped_jobs)): + results = ocr_vllm_batch( + [job.image for job in grouped_jobs], + [job.highres_image for job in grouped_jobs], + skip_text_detection=options.skip_text_detection, + recognize_math=options.recognize_math, + with_bboxes=options.with_bboxes, + ) + for job, result in zip(grouped_jobs, results): + job.future.set_result(result) + logger.info( + "vllm_batch_complete batch_size=%s option=%s duration_ms=%.2f request_ids=%s", + len(grouped_jobs), + options, + (time.perf_counter() - start) * 1000, + ",".join(job.request_id for job in grouped_jobs), + ) + log_timing_summary(logger, collector) + except Exception as exc: + for job in grouped_jobs: + job.future.set_exception(exc) + logger.exception( + "vllm_batch_failed batch_size=%s option=%s request_ids=%s", + len(grouped_jobs), + options, + ",".join(job.request_id for job in grouped_jobs), + exc_info=True, + ) + finally: + reset_current_timing(token) + + +vllm_ocr_batcher = VllmOcrBatcher() diff --git a/vllm_tools.py b/vllm_tools.py new file mode 100644 index 0000000..dd12c2c --- /dev/null +++ b/vllm_tools.py @@ -0,0 +1,210 @@ +import html +import os +import re +import time +from typing import Any, Dict, List, Sequence, Tuple + +from PIL import Image, ImageDraw + +# These must be set before importing surya.settings, because settings are +# instantiated at import time. +os.environ.setdefault("SURYA_INFERENCE_BACKEND", "vllm") +os.environ.setdefault("SURYA_INFERENCE_URL", "http://127.0.0.1:8000/v1") +os.environ.setdefault("SURYA_INFERENCE_AUTOSTART", "false") +os.environ.setdefault("SURYA_INFERENCE_KEEP_ALIVE", "false") +os.environ.setdefault("SURYA_INFERENCE_LOGPROBS", "false") +os.environ.setdefault("SURYA_INFERENCE_MAX_RETRIES", "1") +os.environ.setdefault("SURYA_INFERENCE_PARALLEL", "8") +os.environ.setdefault("SURYA_MAX_TOKENS_FULL_PAGE", "6144") +os.environ.setdefault("SURYA_MAX_BLOCKS_PER_PAGE", "80") +os.environ.setdefault("SUYA_VLLM_IMAGE_FORMAT", "JPEG") +os.environ.setdefault("SUYA_VLLM_JPEG_QUALITY", "92") + +from surya.debug.draw import draw_polys_on_image, draw_bboxes_on_image +from surya.inference import SuryaInferenceManager +from surya.layout import LayoutPredictor +from surya.recognition import RecognitionPredictor +from surya.table_rec import TableRecPredictor +from surya.timing import timing_span + + +_TAG_RE = re.compile(r"<[^>]+>") +OCR_MODE = os.getenv("SUYA_OCR_MODE", "block").strip().lower() + + +def _html_to_text(value: str) -> str: + return html.unescape(_TAG_RE.sub(" ", value or "")).strip() + + +def _load_predictors_vllm() -> Dict[str, Any]: + manager = SuryaInferenceManager(method="vllm") + return { + "manager": manager, + "layout": LayoutPredictor(manager), + "recognition": RecognitionPredictor(manager), + "table_rec": TableRecPredictor(manager), + } + + +predictors_vllm = _load_predictors_vllm() + + +def backend_info() -> Dict[str, Any]: + manager = predictors_vllm["manager"] + backend = manager.backend + handle = getattr(backend, "handle", None) + return { + "backend": getattr(manager, "method", "vllm"), + "started": handle is not None, + "base_url": getattr(handle, "base_url", None), + "model_name": getattr(handle, "model_name", None), + "inference_url": os.getenv("SURYA_INFERENCE_URL"), + "ocr_mode": OCR_MODE, + "parallel": os.getenv("SURYA_INFERENCE_PARALLEL"), + "max_blocks_per_page": os.getenv("SURYA_MAX_BLOCKS_PER_PAGE"), + } + + +def vllm_backend_info() -> Dict[str, Any]: + return backend_info() + + +def _annotate_page(highres_img: Image.Image, page: Any, with_bboxes: bool) -> Image.Image: + if not with_bboxes: + return highres_img + annotated = highres_img.copy() + draw = ImageDraw.Draw(annotated) + for block in page.blocks: + x0, y0, x1, y1 = block.bbox + color = "orange" if block.skipped else ("red" if block.error else "green") + draw.rectangle((x0, y0, x1, y1), outline=color, width=3) + draw.text((x0 + 4, y0 + 4), str(block.reading_order), fill=color) + return annotated + + +def ocr_vllm_batch( + images: Sequence[Image.Image], + highres_images: Sequence[Image.Image], + skip_text_detection: bool = False, + recognize_math: bool = True, + with_bboxes: bool = True, +) -> List[Tuple[Image.Image, Any, Image.Image]]: + start = time.perf_counter() + if OCR_MODE not in {"full_page", "block"}: + raise ValueError("SUYA_OCR_MODE must be 'full_page' or 'block'") + if len(images) != len(highres_images): + raise ValueError("images and highres_images must have the same length") + if not images: + return [] + + with timing_span( + "ocr_vllm_batch_total", + image_count=len(images), + ocr_mode=OCR_MODE, + skip_text_detection=skip_text_detection, + ): + if OCR_MODE == "full_page" or skip_text_detection: + with timing_span("recognition_full_page_call", image_count=len(highres_images)): + pages = predictors_vllm["recognition"](list(highres_images), full_page=True) + else: + target_sizes = [img.size for img in highres_images] + with timing_span("layout_predictor_call", image_count=len(images)): + layouts = predictors_vllm["layout"](list(images), target_image_sizes=target_sizes) + with timing_span("recognition_block_call", image_count=len(highres_images)): + pages = predictors_vllm["recognition"](list(highres_images), layouts, full_page=False) + + elapsed = time.perf_counter() - start + results = [] + with timing_span("response_assembly", image_count=len(highres_images), with_bboxes=with_bboxes): + for highres_img, page in zip(highres_images, pages): + object.__setattr__(page, "_elapsed_seconds", elapsed) + annotated = _annotate_page(highres_img, page, with_bboxes) + results.append((annotated, page, annotated)) + return results + + +def ocr_vllm( + img: Image.Image, + highres_img: Image.Image, + skip_text_detection: bool = False, + recognize_math: bool = True, + with_bboxes: bool = True, +) -> Tuple[Image.Image, Any, Image.Image]: + return ocr_vllm_batch( + [img], + [highres_img], + skip_text_detection=skip_text_detection, + recognize_math=recognize_math, + with_bboxes=with_bboxes, + )[0] + + +def page_ocr_to_response(page: Any) -> Dict[str, Any]: + page_json = page.model_dump() + text_lines = [ + text + for text in (_html_to_text(block.get("html", "")) for block in page_json.get("blocks", [])) + if text + ] + return { + "ocr_text_json": page_json, + "text_lines": "\n".join(text_lines), + "backend": backend_info(), + "elapsed_seconds": getattr(page, "_elapsed_seconds", None), + } + + +def layout_detection_vllm(img: Image.Image) -> Tuple[Image.Image, Any]: + pred = predictors_vllm["layout"]([img])[0] + polygons = [p.polygon for p in pred.bboxes] + labels = [ + f"{p.label}-{p.position}-{round(getattr(p, 'confidence', 0) or 0, 2)}" + for p in pred.bboxes + ] + layout_img = draw_polys_on_image( + polygons, img.copy(), labels=labels, label_font_size=18 + ) + return layout_img, pred + + +def table_recognition_vllm( + img: Image.Image, highres_img: Image.Image, skip_table_detection: bool +) -> Tuple[Image.Image, Any]: + if skip_table_detection: + layout_tables = [(0, 0, highres_img.size[0], highres_img.size[1])] + table_imgs = [highres_img] + else: + _, layout_pred = layout_detection_vllm(img) + layout_tables = [ + tuple(map(int, line.bbox)) + for line in layout_pred.bboxes + if line.label in ["Table", "TableOfContents"] + ] + table_imgs = [highres_img.crop(tb) for tb in layout_tables] + + table_preds = predictors_vllm["table_rec"](table_imgs) + table_img = highres_img.copy() + for result, table_bbox in zip(table_preds, layout_tables): + adjusted_bboxes = [] + labels = [] + colors = [] + for item in [*getattr(result, "rows", []), *getattr(result, "cols", []), *getattr(result, "cells", [])]: + adjusted_bboxes.append( + [ + item.bbox[0] + table_bbox[0], + item.bbox[1] + table_bbox[1], + item.bbox[2] + table_bbox[0], + item.bbox[3] + table_bbox[1], + ] + ) + labels.append(item.label) + colors.append("blue" if "Row" in item.label else "red") + if adjusted_bboxes: + table_img = draw_bboxes_on_image( + adjusted_bboxes, + highres_img, + labels=labels, + label_font_size=18, + color=colors, + ) + return table_img, table_preds