Tofu Snippet Tool

Tofu Snippet Tool #

Classification: Confidential (describes source code, model architecture and infrastructure).

Detects how many documents are on a page image, and splits a page that holds more than one. It replaced two earlier approaches: the vision-LLM classifier’s own bounding boxes, and the OCR-anchored refinement explored in ENG-6272.

Two deployments run from one image. tofu-snippet-gate sits in front of extraction and answers a routing question in under a second. tofu-snippet-main runs the expensive precise pass off that path, only when a user asks for it. Keeping those in separate pods is the whole design: the precise model needs ~5Gi and ~37s per page, which cannot sit in front of every upload.

Where it sits #

bonsai-doc-convert decides, at the end of a split, whether an upload goes through detection or straight to its extraction queue. Detection then routes:

upload → doc-convert → tofu-snippet-v1 → tofu-snippet-gate
                                              ├── one document  → extraction queue
                                              └── two or more   → held for review

The gate forwards the payload it was given rather than answering a question about it. bonsai-doc-convert builds the extraction job up front and hands it over inside the snippet message, so the gate either publishes that payload onward or holds the document. Nothing polls, and nothing waits on a timeout.

Bank statements skip detection entirely — one statement spanning many pages is a single document, so there is nothing to split.

An entity opts in through its snippet_type, and the organization through the auto_suggest_multi_snippet feature flag. NO_SNIPPET and SKIP_MULTI_SNIPPET both mean no. Documents past a page cap skip detection too.

The two models #

Gate (Sam3SnippetGate) Predictor (Sam3SnippetPredictor)
Deployment tofu-snippet-gate tofu-snippet-main
Queue tofu-snippet-v1 tofu-snippet-improve-v1
Per page ~0.9s ~37s
Memory (request/limit) 1Gi / 1280Mi 5Gi / 6Gi
Output plain rectangles mask-fitted polygons

The gate is a distilled EfficientSAM3 (TinyViT-11M backbone) with no segmentation head, so it can only draw rectangles. That is enough to decide how many documents are on a page, which is the only question on the critical path. Redrawing them precisely is the predictor’s job, triggered by the user’s “Improve” action.

Both are OpenVINO exports of ours, fetched from the tofu-snippet-model-${env} S3 bucket at startup — no weights in the image, nothing pulled from HuggingFace at build or run time. The upstream lineage is an ONNX export of facebook/sam3 under the Meta SAM License, converted offline. The ultralytics SAM3 wrapper is deliberately avoided: it is AGPL-3.0.

Detection uses SAM3’s Promptable Concept Segmentation with the prompt fixed at "document". The prompt embedding is committed as a 33KB .npz asset rather than computed at startup, which is why the prompt is a constant and not an env var — the 1.6GB language encoder that produces it is not a dependency.

Known limitation: flatbed scans #

PCS works on photographs, where each document is an object against a contrasting background. It does not split a flatbed scan of several receipts on one white sheet: with no background to segment against, the whole sheet legitimately matches “document” and returns one page-sized box, which the area-ratio guard drops. Scanned pages are therefore a no-op rather than a wrong answer. Splitting them needs a background-independent signal — grouping OCR text clusters is the natural next attempt.

Failing open #

Detection is on the critical path, so a dead snippet service must not stop extraction. A gate that will not load, a page that will not read, and a bonsapi that will not answer all end the same way: log it, forward anyway. An undetected page is a page without suggestions, never a document that stops moving. Only a broker or bonsapi failure during routing itself raises, because that is the case a retry can fix.

This is why a pod whose model fails to load still reports ready — refusing readiness would stall uploads to protect a feature that is allowed to be absent.

Operations #

  • Scaling. KEDA on queue depth, floor of one replica, no scale-to-zero: at zero, nothing extracts except bank statements. The gate’s depth counts documents (its routing decision needs every page, so a document cannot be split across workers); tofu-snippet-main counts pages, since each Improve message carries one.
  • Drains. Every PDB uses maxUnavailable: 1 rather than minAvailable: 1, so a node can be drained at a single replica (ENG-7305).
  • Rollouts. Both deployments have a readinessProbe checking a marker file written once the model has loaded, so a rollout cannot retire the old pod before the new one can consume.
  • Duplicate work. Jobs are claimed in Redis with an ownership token, so a worker that runs past its claim TTL cannot extend or delete the claim its successor now holds.
  • Failures. A failed page does not abort its siblings; the job then raises so bonsai-mq dead-letters it for replay.
  • Threads. The inference pool is sized from CONTAINER_CPU_LIMIT divided by RABBITMQ_MAX_CONCURRENT_JOBS, so changing an overlay’s cpu limit is enough. SNIPPET_INFERENCE_THREADS overrides it.

Re-measuring memory #

tests/measure_memory.py produces the peak-RSS numbers the envelopes above are sized from. Re-run it after changing a model, a threshold, or the thread count — the envelopes are deliberately just above measured peak, so an unmeasured model change risks an OOM kill 37 seconds into an inference.

See also #

  • apps/tofu-snippet/README.md — implementation detail, model provenance, and the prompt-embedding regeneration recipe.
  • ENG-7130 (detection engine), ENG-7248 (gate-first), ENG-7250 (Improve worker).