Repository navigation
Replace tokenizers-cpp with our own Rust tokenizer FFI and add zero-shot classification - #4
Conversation
|
Summary Verdict: Request changes (small, targeted) Scope reviewed: 27 files, +2403/−864. I read Major
Minor
Nits
Praise
|
iswaryaalex
left a comment
There was a problem hiding this comment.
Nice work on this one.
The good
tokenizers-cppwas already a wrapper around the HFtokenizerscrate. This drops the middle layer and calls the crate directly, so parity withtransformerscomes 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-cppalso pulled in SentencePiece, Abseil, protobuf and msgpack for.modelfiles, but ort-server only readstokenizer.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
\nand\rin labels and add a smoke test?
- Release builds
- No workflow changes needed.
release.ymlalready builds the four archives Lemonade downloads, with matching names. - But nothing has been built for this branch yet. Could you kick off a
workflow_dispatchrun and link it? I'd like to see all four green, includingsmoke.pyandcheck_deps.py. linux-arm64 is the one I'd watch.
- No workflow changes needed.
- 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.mdstill has "" and a blankmodel_type/architectures.
- 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
- 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.
- Today's golden comes from our own copy of
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.mdstill 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
labelson a fixed-label model is now a 400.
- Licenses:
THIRD_PARTY_NOTICES.mdlists 4 Rust bits, butCargo.lockhas 81 crates. Not new to this PR. Suggestcargo aboutto generate the list andcargo deny check licensesin the release job. Keep the hand-written entries for ORT, httplib, json and Oniguruma. - A
pull_requesttrigger for build and test, without the upload, so the next PR gets checked automatically.
99501bf to
a1aea42
Compare
|
Thanks @iswaryaalex and @SlawomirNowaczyk for the feedback. The two reviews overlap so I merged those and answer each one once. Fixed
Open: no model to test againstThe 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:
|
fl0rianr
left a comment
There was a problem hiding this comment.
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:
-
CMake dependency tracking:
Cargo.lockis missing from thetok_ffibuild dependencies. Updating a locked Rust dependency alone may therefore reuse a stale static library. Please add it toDEPENDSand considercargo build --lockedfor reproducibility. -
Integer overflow:
--max-labelsand--max-label-chars, as well as their manifest equivalents, convertlong longtointwithout checking the upper bound. Values aboveINT_MAXshould 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.
|
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 |
fl0rianr
left a comment
There was a problem hiding this comment.
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.
|
Thank you for highlighting the edge case @fl0rianr :) |
|
Thanks for the thorough round. I re-reviewed at
Nits: the first character-level cut ( |
|
Thank you for the review @SlawomirNowaczyk Rest of the issues are handled. |
aae8393 to
925f99d
Compare
|
@meghsat Re-reviewed at
on-blocking: commit the export script you used for the 16-case run, and decide whether |
|
@iswaryaalex @SlawomirNowaczyk |
|
Thanks, I re-reviewed at The remaining nit (the first character-level cut doesn't insert One thing to fix before Lemonade ships the built-in model: the published Suggest:
For the record: |
|
Thanks for review @SlawomirNowaczyk
|
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
/classifyendpoint, 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
tokenizerscrate computes those on every call; themlc-ai/tokenizers-cppshim we used sat in between anddropped 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 nolabelsfield, and supplying one is now an explicit 400 rather than a silent no-op.Zero-shot models take the label set per request:
Scores are softmaxed over the supplied list, so they sum to 1 before top_k.
Also in this PR (details in VERSIONS.md):
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.
input_idsfor all 20,001 cases on all 5 models.