Wasabi-30B / inference /README.md
enzohagz's picture
Add inference code and usage instructions
b3baef2 verified
|
Raw History Blame Contribute Delete
6.37 kB

wasabi-read

Reads a directory of page images (.jpg, .jpeg, .png) and writes one Markdown file per page with Wasabi-30B. The supported setup uses one checkpoint, Wasabi-30B, for every model step. Per page: an orientation check by Wasabi-30B itself (a 4-way upright probe at 4,096 image tokens; a turn is applied only at margin >= 5.0), UVDoc dewarp on CPU, the whole page read at 8,192 image tokens (greedy), and a temperature ladder when a first answer is degenerate.

Install

pip install -e "./inference[dewarp]"

Python 3.10+. Endpoint and local runs both dewarp on the local CPU, so the dewarp extra (torch and OpenCV) is needed in both modes; without OpenCV the run refuses. The OpenCV, torch, numpy and Pillow versions and the UVDoc pins are part of the output identity, so a rerun with other versions is refused and one output directory never mixes them. A local GPU run needs the vLLM build of commit 6adad087 (released vLLM 0.26.0 has no muse_glimmer model); an endpoint run needs no vLLM. After installing that build, run python inference/patch_vllm_max_image_tokens.py --apply with the Python of the vLLM environment: unpatched, its processor rejects the max_image_tokens setting the servers are started with. It edits one file, pinned by sha256 (--check, --revert); a server you start yourself for endpoint mode needs it too. Download the weights with huggingface-cli download.

Tested environment: opencv-python-headless 5.0.0.93, torch 2.13.0+cu130, numpy 2.3.5, Pillow 13.0.0, vllm 0.26.1rc1.dev769+g6adad0876, transformers 5.15.0, openai 3.37.0 (the package itself talks to servers with the Python standard library; openai was present, not required).

The package sets NUMPY_MADVISE_HUGEPAGE=0 at import unless you set it yourself: with numpy's default, a Linux host with transparent huge pages on madvise can stall in the kernel while dewarping large page images. The output is the same either way.

Configuration

  • WASABI_READ_API_KEY: optional API key, sent to both the reader and the orientation endpoint. Set it in the environment and keep the value out of commands, scripts and documentation.
  • WASABI_READ_CACHE: directory for the downloaded UVDoc files. The default is ~/.cache/wasabi-read, and the files go into a pinned uvdoc-<revision> subdirectory. Page outputs and run caches stay under the output directory you pass.
  • --uvdoc-dir PATH: keep the UVDoc files in PATH in place of the default cache. In the normal fetch mode, missing pinned files are downloaded and every file is verified.
  • --uvdoc-source vendored --uvdoc-dir PATH: use an existing local UVDoc copy without downloading. The code and weights it needs must pass the built-in size and sha256 checks.

Orientation

The orientation check runs on its own server at 4,096 image tokens, before the reader starts. It uses Wasabi-30B and turns a page only at margin >= 5.0. A turn is a decision from one greedy probe per rotation: a different serving stack (attention backend, CUDA build, batching) can change it on a few pages.

Run on a CPU-only machine (endpoint mode)

Everything except the 30B forward passes runs on the local CPU. Point the package at two OpenAI-compatible servers, both serving Wasabi-30B and started with --reasoning-parser muse_glimmer: the reader, with --mm-processor-kwargs '{"max_image_tokens": 8192}', and the orientation check, with --mm-processor-kwargs '{"max_image_tokens": 4096}'. Name what each server serves with --endpoint-revision and --orient-endpoint-revision: the URLs and revisions are part of the output identity, so a rerun against another server is refused and one output directory never mixes servers. If the servers need a key, set WASABI_READ_API_KEY (see Configuration).

python inference/wasabi_read.py pages/ out/ --engine endpoint --endpoint-url http://gpu-host:8000 --endpoint-revision wasabi-30b-v1.0 --orient-endpoint-url http://gpu-host:8001 --orient-endpoint-revision wasabi-30b-v1.0

--plan prints the stages and exits without loading a model.

python inference/wasabi_read.py pages/ out/ --engine endpoint --endpoint-url http://gpu-host:8000 --endpoint-revision wasabi-30b-v1.0 --orient-endpoint-url http://gpu-host:8001 --orient-endpoint-revision wasabi-30b-v1.0 --plan

Run on a local GPU (vllm engine)

The package starts and stops the local vLLM servers itself (the orientation check first, then the reader, both from the same weights). The engine kind and its version, and the reasoning parser (vLLM's built-in muse_glimmer, or the sha256 of a --reasoning-plugin file), are part of the output and orientation identity, so a rerun on another engine or parser is refused.

python inference/wasabi_read.py pages/ out/ --engine vllm --model-dir /models/Wasabi-30B --tensor-parallel-size 1 --whole-upscale-tokens 8192

--engine transformers loads the model in-process. It was not part of the measured runs and does not follow the --engine vllm path: the image-token budget is set on the processor when the model is loaded (not through a server), the sampling seed is applied with torch.manual_seed, and the raw decoded text is used, with no server-side reasoning-channel split and no repetition-detection stop. Use it for a few pages only; its results are not comparable to the vLLM runs.

Resume and exit codes

Each page is written as soon as it finishes (<stem>.md, atomic), with a sidecar under out/.wasabi-read/. Rerun the same command to resume: finished pages are skipped, orientation and dewarp results are reused, and when nothing is left no engine is started. An existing output that this input or package settings did not write is refused unless --overwrite-stale is given. A page whose final answer is degenerate, empty or failed is not written and is listed with its reason in out/.wasabi-read/run-report.json.

exit meaning
0 every page has an output
2 refusal (bad input or options, stale or locked output directory)
4 incomplete: some pages have no usable output; rerun to retry them

Licence

Code: GNU Affero General Public License v3.0 (LICENSE, NOTICE). UVDoc code is MIT (Copyright (c) 2023 Tanguy MAGNE); its pretrained weights are downloaded from the pinned commit at first use (size and sha256 checked before they are loaded) and keep their own terms.