Skip to content

Replace tokenizers-cpp with our own Rust tokenizer FFI and add zero-shot classification - #4

Merged
meghsat merged 11 commits into
mainfrom
router-zero-shot
Oct 28, 2026
Merged

meghsat merged 11 commits into
mainfrom
router-zero-shot

Conversation

@meghsat

@meghsat meghsat commented Oct 1, 2026 •

Copy link
Copy Markdown
Contributor

Adds a third task to ort-server: zero-shot-classification, where the label set arrives with each request instead of being baked into the model directory at export time. Same /classify endpoint, one optional field.

Scoring a per-request label set means knowing which tokens belong to which label, which needs the tokenizer's character offsets. The HuggingFace tokenizers crate computes those on every call; the mlc-ai/tokenizers-cpp shim we used sat in between and
dropped them. So this PR replaces it with ~100 lines of our own C ABI over the same crate (third_party/tok_ffi) that passes the offsets through. Fixed-label models are unchanged as in no labels field, and supplying one is now an explicit 400 rather than a silent no-op.

Zero-shot models take the label set per request:

POST /classify
{"text":   "Write a function that merges two sorted lists",
"labels": ["code generation", "math reasoning", "creative writing"],
"top_k":  2}
Returns:
{"labels": {"code generation": 0.84, "math reasoning": 0.11}}

Scores are softmaxed over the supplied list, so they sum to 1 before top_k.

Also in this PR (details in VERSIONS.md):

  • Every /classify response carries truncated, on every task: true when the text was cut to fit the token window, false when it fitted. Servers before 0.4.0 send no field, so a missing field means "unknown".
  • /health reports version, so a client can tell the two apart.
  • Zero-shot limits: --max-length, --max-labels, --max-label-chars and --max-grid-cells, or the same fields in manifest.json; the flag wins. All four are startup errors on a fixed-label model.

Reference model: lemonade-sdk/LFM2.5-Encoder-350M-Prompt-Router-ONNX, an ONNX export of LiquidAI's LFM2.5-Encoder-350M-Prompt-Router.

Experiments:
Compared the new ort-server with the previous build (same ONNX Runtime 1.27.0) on the models Lemonade's PII router uses: mmBERT-32k PII, OpenMed privacy-filter-multilingual v2, perplexity pplx-pii-masking and two phishing classifiers, using the 20,001-case Nemotron-PII routing corpus.

  • Tokenization: identical input_ids for all 20,001 cases on all 5 models.
  • Model output: old and new binaries side by side on 500 cases per PII model (longest, most non-ASCII, random): every label score exactly equal.
  • End to end through Lemonade's router: 1,500/1,500 routing decisions match the recorded benchmark runs (500 per model).
  • Edge inputs (emoji, CJK, Arabic, decomposed accents, zero-width characters, literal special tokens, empty text): identical.
  • Latency: unchanged (within 3%).

@SlawomirNowaczyk

Copy link
Copy Markdown

Summary
The PR replaces tokenizers-cpp with a small, careful C ABI over HF tokenizers 0.22.2. This is needed to get the character offsets that zero-shot scoring depends on, and it also drops sentencepiece, protobuf, abseil and msgpack. It then adds a zero-shot-classification task that builds text_pool and category_pool (which tokens belong to the text and to each label) from those offsets. The tokenizer swap is well verified. The zero-shot path has gaps in input validation, resource limits and truncation, and so far it has only been checked against a synthetic fixture.

Verdict: Request changes (small, targeted)

Scope reviewed: 27 files, +2403/−864. I read src/main.cpp, third_party/tok_ffi/*, CMakeLists.txt, release.yml, README, VERSIONS.md and docs line by line. I did not review Cargo.lock, the fixture tokenizer.json/model.onnx files, or the deleted license files. I did not build or run anything.

Major

  1. src/main.cpp:745 — labels can break the prompt template. Labels are pasted unescaped into "Categories:\n- <label>\n…\n\nText:\n". The current check (:734) only rejects blank labels. A label containing \n, such as "a\n- b" or "x\n\nText:\nfoo", makes the model read a different list of categories than the pooling grid describes, and the scores come back silently wrong. Suggested fix: reject labels containing \r, \n or other control characters with a 400, and add a smoke test for it.
  2. src/main.cpp:782 — no limits on label count or input size, and the big allocation happens before the window check. category_pool allocates R × n floats before the per-label loop that detects window overrun. Tokenizer truncation is turned off, so the whole text is tokenized first. A request with tens of thousands of labels, or a text of several MB, costs a large allocation plus an R × n loop before it fails. Suggested fix: cap labels.size() and label length. Count the label block's tokens before building the grids. Cut the text by characters (for example ~16 × max_length) before tokenizing. Build both grids in a single pass over the sorted ranges.
  3. src/main.cpp:762 — truncation silently drops the end of the text. Labels come first, so the "plain cut" keeps the start of the text and drops its end. In practice, the end of a long user message often holds the actual question, for example code pasted above "why does this crash?". Nothing in the response says the text was cut. Suggested fix: budget the text as max_length − label-block tokens, consider keeping head + tail, and return "truncated": true so Lemonade can show it in route_trace.
  4. Zero-shot correctness has not been checked against the real model. The tiny-router golden values come from a Python transcription of route() running on a synthetic graph. If that transcription and main.cpp share a mistake, both agree and the test still passes. The 20k-case and 1,500-decision experiments in the description cover only the tokenizer swap on fixed-label PII models. Suggested fix: add a parity check on the real LFM2.5 ONNX export against the PyTorch reference in tools/classifier_catalog (a few dozen prompts, including non-ASCII labels and long text).

Minor

  1. src/main.cpp:516 / :531 — tokenizer leak. If Ort::Session or bind_inputs() throws, the destructor never runs and tokenizer_ leaks. It only happens at startup, but the fix is trivial: std::unique_ptr<TokHandle, decltype(&tok_free)>.
  2. docs/ZERO-SHOT-ROUTER.md:15, :71 — unfinished placeholders. <Model checkpoint. WIP> and model_type: , architectures: are still in the doc.
  3. VERSIONS.md — version bump. The /classify contract changed (labels is now required or rejected depending on the model) and the tokenizer stack was replaced. The release should be a minor bump (0.4.0) rather than 0.3.8, so Lemonade's pin bump clearly signals the change.

Nits

  • CMakeLists.txt: on non-Windows platforms, no Rust native libraries are linked (Threads::Threads, ${CMAKE_DL_LIBS}, m). This works on glibc ≥ 2.34, but it's fragile on older or musl toolchains.

Praise

  • catch_unwind on every FFI entry point, --locked cargo builds, and default-features = false, which keeps the http feature, and with it OpenSSL, out of the build (the lesson from 0.3.4).
  • Separate architecture allowlists per task, which are refused in both directions.
  • The output-size check against the label count catches exports where the label count was baked in.
  • The test offsets_are_character_indexed pins down the exact byte-vs-character pitfall.

@iswaryaalex iswaryaalex left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice work on this one.

The good

  • tokenizers-cpp was already a wrapper around the HF tokenizers crate. This drops the middle layer and calls the crate directly, so parity with transformers comes by design.
  • The 20,001-case parity run (same ids, same scores, same routing decisions) is the evidence I wanted.
  • Char offsets via encode_char_offsets, with a test that would catch a byte-offset build. Nice!
  • tokenizers-cpp also pulled in SentencePiece, Abseil, protobuf and msgpack for .model files, but ort-server only reads tokenizer.json. Dropping that stack shrinks both the binary and the license list.
  • Panics are caught in the Rust layer, so a bad tokenizer can't take the server down.

The only tradeoff is that we now own the unsafe FFI layer and the crate bumps ourselves, which seems like a fair price for getting offsets.

Medium

  • Cut labels still get scored
    • We only refuse a label when it has zero tokens left. Test M even expects it: "gamma is cut mid-label but still pooled".
    • The README says we refuse this case. Can we refuse any label that ends past enc.ends.back() and update test M?
  • Line breaks in labels
    • "a\n- b" shows up as two categories, and "\n\nText:\n" adds a fake text header.
    • Can we reject \n and \r in labels and add a smoke test?
  • Release builds
    • No workflow changes needed. release.yml already builds the four archives Lemonade downloads, with matching names.
    • But nothing has been built for this branch yet. Could you kick off a workflow_dispatch run and link it? I'd like to see all four green, including smoke.py and check_deps.py. linux-arm64 is the one I'd watch.
  • A model to run
    • There's no ONNX export of the router out there yet, so nobody can actually use this task. Can we add a full export script, and note whether it needs trust_remote_code?
    • ZERO-SHOT-ROUTER.md still has "" and a blank model_type / architectures.
  • Real-model golden
    • Today's golden comes from our own copy of route() on random weights, so it can't catch a mismatch with the reference.
    • A small golden from the real checkpoint would close that, even as an optional job.

Minor

  • Long text gets cut quietly. Text goes last, so the end of a long input is dropped and the caller can't tell. A log line plus a note in the docs would be enough.
  • Put the 20,001-case parity script under test/, so we can rerun it next time the tokenizer moves.
  • Version and release note
    • VERSIONS.md still says "unreleased". What version will this be? Lemonade uses one exact pin for all platforms, so #3952 needs the number and a tagged release right after merge.
    • Add a line that labels on a fixed-label model is now a 400.
  • Licenses: THIRD_PARTY_NOTICES.md lists 4 Rust bits, but Cargo.lock has 81 crates. Not new to this PR. Suggest cargo about to generate the list and cargo deny check licenses in the release job. Keep the hand-written entries for ORT, httplib, json and Oniguruma.
  • A pull_request trigger for build and test, without the upload, so the next PR gets checked automatically.

@meghsat

meghsat commented Oct 24, 2026

Copy link
Copy Markdown
Contributor Author

Thanks @iswaryaalex and @SlawomirNowaczyk for the feedback. The two reviews overlap so I merged those and answer each one once.

Fixed

# Issue change
1 A label can break the prompt template The server refuses any control character in a label. This covers \n, \r, \t, NUL and DEL. The answer is a 400. Smoke tests cover each one, and a label that forges a Text: block.
2 Cut labels still get scored. No limits on request size Labels now draw from the window first. A label set that does not fit is refused whole, so the server never pools a label from half its tokens. Three caps bound a request: max_labels, max_label_chars and max_grid_cells. Each one derives from max_length, so it scales with the model.
3 A long text is cut silently The cut keeps the head and the tail, so a question at the end survives. The response carries "truncated": true, and omits the field when nothing was cut.
4 Doc placeholders. No way to get a model The placeholders are gone. The model is gated so can't add the tests for now since ONNX model is not provided. Will raise a PR once the model is opensourced
5 Version and release note VERSIONS.md records 0.4.0 and the /classify contract change.
6 Tokenizer leak on a startup throw The handle is a unique_ptr now.
7 No Rust native libraries linked CMakeLists.txt links Threads::Threads, ${CMAKE_DL_LIBS} and m.
8 Parity script not in the repo It is test/parity_tokenizer.py. It compares ids and character offsets against HuggingFace.
9 License list is hand-written The release job runs cargo deny check licenses and generates the notices with cargo about. It refuses to package an archive without them.
10 No automatic check on a PR release.yml now runs on pull_request. It builds and tests, and skips the upload.

Open: no model to test against

The only ONNX export of a router is gated. So CI cannot download it, so an optional job is not possible. Hence, we cannot commit a golden. Scores from the checkpoint are a derived artifact of a licensed model. I will update the tests and export scripts once the model is made public.

We ran the real check locally instead, against the reference implementation that ships with the checkpoint, on the int8 graph:

  • 16 cases: worst score delta 2.7e-08, same argmax everywhere.
  • Tokenizer: 3000 cases, 494,286 tokens, ids and offsets exact.

@meghsat
meghsat requested a review from iswaryaalex October 24, 2026 10:33

@fl0rianr fl0rianr left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice work addressing the previous review findings! The tokenizer parity checks, real-model validation, improved input handling, and green CI across all four platforms give me good confidence in this change.

I found just two small remaining issues:

  1. CMake dependency tracking: Cargo.lock is missing from the tok_ffi build dependencies. Updating a locked Rust dependency alone may therefore reuse a stale static library. Please add it to DEPENDS and consider cargo build --locked for reproducibility.

  2. Integer overflow: --max-labels and --max-label-chars, as well as their manifest equivalents, convert long long to int without checking the upper bound. Values above INT_MAX should be rejected rather than potentially wrapping.

Integration note for Lemonade #3952: The new truncated response field is currently dropped when Lemonade constructs its /v1/classify response. We should preserve it so callers can recognize when a routing decision was based on truncated input. This can be addressed in #3952 rather than here.

The real-model parity results are reassuring, and I understand why the gated model cannot currently be included in CI.

Overall, this looks very close to ready. I wouldn't hold up the PR for further refactoring or additional test infrastructure. Just the two small fixes above, followed by the 0.4.0 release so Lemonade can update its backend pin.

@meghsat

meghsat commented Oct 24, 2026 •

Copy link
Copy Markdown
Contributor Author

Thanks for the review @fl0rianr Added cargo.lock to the tok_ffi and added an upper bound to the long long variables. It refuses a value above that bound before the cast. Will add the truncated on the Lemonade side in #3952

@meghsat
meghsat requested a review from fl0rianr October 25, 2026 00:13

@fl0rianr fl0rianr left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for addressing the previous findings.
One last edge case I noticed in src/main.cpp:1049–1075:

Unicode truncation may split a multi-token character. The head/tail logic uses token offsets to select character boundaries. When several tokens represent the same Unicode character (e.g. an emoji), this can retain more tokens than budgeted. The final fallback then calls enc.ids.resize(max_len), potentially cutting through that character's token sequence.

Could we add a small regression test using a narrow token window and a long sequence of emojis (e.g. "🚀" repeated 100 times), and ensure the final truncation respects character boundaries? Ideally, shorten the text and re-tokenize rather than trimming raw token IDs.

This is the only remaining concern I found. No broader refactoring needed; with this addressed, I'm happy to approve and suggest to merge.

@meghsat

meghsat commented Oct 25, 2026

Copy link
Copy Markdown
Contributor Author

Thank you for highlighting the edge case @fl0rianr :)
Handled it in my latest commit

@SlawomirNowaczyk

Copy link
Copy Markdown

Thanks for the thorough round. I re-reviewed at 8fd2863, and nearly everything from the previous reviews is addressed. A few remaining points:

  1. Padding regression on the zero-shot path. tok_new no longer calls with_padding(None) (and padding_is_disabled was removed), on the basis that pads are filtered downstream. Only classify() filters them (main.cpp:884). route() doesn't, so a zero-shot tokenizer.json with a padding section inflates label_tokens and the text length, and feeds pad IDs to the model under an all-ones mask, which is the 0.3.1 bug class. The classify() filter is also right-side only. Suggest restoring with_padding(None) (it matches transformers' default padding=False), keeping the filter as a backstop, and adding a padded tiny-router variant to the smoke tests.
  2. classify() truncates to max_length but never sets truncated (rank() at :960). Lemonade #3952 needs that flag for fixed-label PII classifiers as well, so adding it now saves a second release.
  3. The comment on route() ("mirrors the reference route() exactly") only holds when the input fits; head+tail and \n...\n deliberately diverge. Did the 16-case real-model check include truncated inputs? Consider leaving the marker's tokens out of text_pool.
  4. The 0.4.0 notes in VERSIONS.md should also list the truncated field, head+tail truncation, the three limits and their flags, and the new 400s.
  5. Why the global mutex around tok_encode? Tokenizer is Sync, and route() encodes several times per request. I suggest either remove it or document the reason.
  6. Still open from @iswaryaalex's review: an export script and a note on whether it needs trust_remote_code. Even if the model is gated, the script would help anyone who has access. @iswaryaalex, could you re-review? Your CHANGES_REQUESTED is still the active state.

Nits: the first character-level cut (utf8_head_tail) splits ½/½ without \n...\n, unlike the documented ⅓/⅔ cut; the "internal: label got no tokens" error should be a 400; the Rust tests should fail rather than skip when the committed fixture is missing.

@meghsat

meghsat commented Oct 26, 2026

Copy link
Copy Markdown
Contributor Author

Thank you for the review @SlawomirNowaczyk
@iswaryaalex I still feel we should push back on the export script as docs/ZERO-SHOT-ROUTER.md now carries the full contract: the four inputs, the dynamic_axes block, the route() wrapper trap, trust_remote_code, the fast-tokenizer requirement, and max_length. Anyone with access can write the export against it. I want to release the script in a follow up PR once the model goes public.

Rest of the issues are handled.

@iswaryaalex

Copy link
Copy Markdown

@meghsat Re-reviewed at 925f99d, everything from my Oct 23 review and both of Slawomir's rounds is fixed, verified in the code.

  • One new blocker before we tag 0.4.0: The 0.3.7 doesn't emit truncated, so lemond's on_truncation: fail silently never fires and the PII policy in #3952 forwards the full prompt to the cloud..

  • Please report the ort-server version on /health so lemond can refuse on_truncation: fail against a backend that can't honour it.

on-blocking: commit the export script you used for the 16-case run, and decide whetherlfm2-bidirectional belongs in the allowlist."

@iswaryaalex iswaryaalex left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@meghsat

meghsat commented Oct 28, 2026 •

Copy link
Copy Markdown
Contributor Author

@iswaryaalex @SlawomirNowaczyk
Added an older LFM2 zero-shot variant to Lemonade HF and included it as the reference model here

@SlawomirNowaczyk

Copy link
Copy Markdown

Thanks, I re-reviewed at 78e1e1a. All six points from my last round are fixed and verified in the code.

The remaining nit (the first character-level cut doesn't insert \n...\n) is non-blocking. LGTM for the C++/FFI side.

One thing to fix before Lemonade ships the built-in model: the published manifest.json sets max_length: 128000 (max_position_embeddings, the export.py default), while tokenizer.json truncates at 448. With 128K, practically nothing is ever cut, and the export uses eager attention (16 heads, 6 full-attention layers). An 8K-token message needs about 4 GB of temporary memory per attention layer, and 16K about 17 GB, so a large paste into chat would make the router's classifier run out of memory or stall. It's also 288× beyond the training window with no real-model validation past about 3K characters, and the head+tail truncation path never runs against the real model.

Suggest:

  • pick a measured serving window (448, or e.g. 2048 after measuring CPU latency and memory), set it in the HF manifest, and make it export.py's default;
  • add a real-model golden case longer than that window;
  • optionally, log a startup warning when the manifest's max_length is far above the tokenizer's truncation.max_length.

For the record: model_type is lfm2 in the published config, so no lfm2-bidirectional allowlist entry is needed. After merging, please tag 0.4.0 so #3952 can bump onnxruntime.cpu off 0.3.7.

@meghsat

meghsat commented Oct 28, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for review @SlawomirNowaczyk

  • Added a --max-length flag to override the zero-shot token window giving the control to the user.
  • Made the truncated flag a three state.
  • Added the zero-shot label similar to the systemone

@meghsat
meghsat merged commit c9f193f into main Oct 28, 2026
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants