sentencepiece-rust 0.1.1

From-scratch pure-Rust SentencePiece inference (encode/decode), differentially verified against the upstream C++/Python implementation.
Documentation
  • Coverage
  • 30.77%
    8 out of 26 items documented1 out of 1 items with examples
  • Size
  • Source code size: 848.7 kB This is the summed size of all the files inside the crates.io package for this release.
  • Documentation size: 627.0 kB This is the summed size of all files generated by rustdoc for all configured targets
  • Ø build duration
  • this release: 2s Average build duration of successful builds.
  • all releases: 3s Average build duration of successful builds in releases after 2024-10-23.
  • Links
  • Homepage
  • VoiceLessQ/sentencepiece-rust
    1 0 0
  • crates.io
  • Dependencies
  • Versions
  • Owners
  • VoiceLessQ

sentencepiece-rust

A from-scratch, pure-Rust port of SentencePiece inference (encode/decode), differentially verified against the upstream C++/Python implementation.

This is not a binding. The existing sentencepiece crate wraps Google's C++ library; this crate reimplements the tokenisation algorithms in Rust from the reference source, the same way the author's other library ports/reimplement their Python originals and verify against them.

It is also independent of the similarly named sentencepiece-rs crate. This code was written from scratch against the upstream C++/Python reference. It is not a copy, fork, or derivative of that crate or of any other Rust port; the shared subject matter is all they have in common.

Training is out of scope. A trained .model is loaded and used to tokenise.

Correctness goal

The point of this crate is to match upstream SentencePiece exactly:

  • Every module is a line-by-line port of the reference file it names in its header, so its behaviour can be checked against the original.
  • Output is verified against the Python sentencepiece oracle over thousands of real-text lines, both encode and decode. See The Rust tokenizer test and Verification breadth. A difference from the oracle is treated as a bug until it turns out to be an equally-optimal tie (the one documented precision fence).

Two things fall out of that:

  • Zero dependencies, so there is nothing to audit but this crate and std.
  • Speed is measured, not assumed (see Speed), and never traded against matching the oracle.

Why

The scary parts of the upstream repo don't apply to inference:

  • The 138k-line normalization_rule.h is compile-time charsmap generation; a trained model embeds its own precompiled_charsmap, which we just replay.
  • The suffix-array trainer, builder, and esaxx are training-only.
  • The .proto is tiny, so a hand-written proto2 reader replaces prost/protoc and keeps the crate dependency-free.

Milestones

Version Scope Verified against
v0.1 BPE segmentation; ASCII / whitespace normalisation; byte-fallback encode Python sentencepiece, ASCII corpora
v0.2 Darts charsmap normaliser to full Unicode input + Unicode corpora (full-width, ligatures, CJK, …)
v0.3 (current) Unigram Viterbi segmentation; byte-fallback and byte-piece decode 3 models (BPE, Unigram, BPE+byte_fallback), encode and decode

Inference is feature-complete for both BPE and Unigram models, with faithful decode (byte-piece reassembly, unk surface, dummy-prefix handling). Both directions are verified against the Python oracle.

Verification breadth

Beyond the small committed corpora, the encoders were run over the upstream training corpora (botchan.txt, 4.3k English lines; wagahaiwa_nekodearu.txt, 2.3k Japanese lines) through all three models: 13,264 lines, encode and decode. The result was 13,258 exact, 6 equal-score reorderings, and 0 real mismatches.

The 6 differences are all runs of repeated punctuation (......., long ----) on the Unigram model. Each is a same-multiset reordering of pieces, that is, an equally-optimal Viterbi path with the same total score. This is a documented precision fence: the compiled oracle binary's optimised float arithmetic resolves these exact ties differently than portable f32, and both segmentations are optimal. See the note in src/unigram.rs. The test counts these separately and only fails on a genuine (different-piece) mismatch.

To reproduce the broad run, point gen_oracle.py at a larger corpus instead of the small corpus_*.txt files. For example, use botchan.txt and wagahaiwa_nekodearu.txt from the upstream SentencePiece repo's data/ directory. Those corpora are not vendored here; download them from the original source.

Layout

src/
  proto.rs       hand-written proto2 wire reader (no deps)
  model.rs       ModelProto -> typed struct (pieces, flags, normalizer spec)
  vocab.rs       piece <-> id maps, scores, byte-fallback ids
  normalizer.rs  whitespace/▁ normalisation + Darts charsmap replay
  bpe.rs         BPE merge encoder
  unigram.rs     Unigram Viterbi best-path (+ piece prefix-trie)
  trie.rs        Darts double-array (common_prefix_search / traverse)
  processor.rs   SentencePieceProcessor: load + encode/decode orchestration
examples/inspect.rs   load a model and dump a summary
tests/oracle.rs       differential test vs the Python oracle
oracle/gen_oracle.py  generates oracle/cases.tsv from Python sentencepiece

Each module names the reference file it ports in its header comment, from the upstream google/sentencepiece source. That source is not vendored here. During development it sits as a local clone outside the repo; clone it from the original if you want to follow the ports.

Try it

# in-tree toolchain (see global setup); plain `cargo` works too
cargo run --example inspect -- tests/models/botchan_1000_bpe.model

Verify against the oracle

# reference implementation, used only to produce the oracle
python -m pip install sentencepiece

# For each (model, output) pair, run over the ASCII corpus (>) then the
# Unicode corpus (>>):
#   botchan_1000_bpe.model          -> cases.tsv          (BPE)
#   test_oss_model_unigram.model    -> cases_unigram.tsv  (Unigram)
#   wagahaiwa_2000_bpe_byte.model   -> cases_byte.tsv     (BPE + byte_fallback)
python oracle/gen_oracle.py tests/models/botchan_1000_bpe.model \
    oracle/corpus_ascii.txt > oracle/cases.tsv
python oracle/gen_oracle.py tests/models/botchan_1000_bpe.model \
    oracle/corpus_unicode.txt >> oracle/cases.tsv
# ... likewise for the Unigram and byte-fallback fixtures.

cargo test

The Rust tokenizer test

tests/oracle.rs is a plain cargo test integration test. It is pure Rust, with no Python at run time. For each fixture model it:

  1. loads the .model with SentencePieceProcessor (the same public API a user calls);
  2. reads the generated oracle/cases*.tsv (hex(text)\t<ids>\thex(decode(ids)));
  3. encodes each text and compares the ids to the oracle's, and decodes the oracle's ids and compares the text, so both directions are checked;
  4. classifies any encode difference: a same-multiset reordering counts as an equally-optimal Viterbi tie (accepted, counted separately), and anything else is a genuine mismatch that fails the test.

It runs across all three fixtures (BPE, Unigram, BPE+byte_fallback) and prints a breakdown, for example:

checked 117 cases against the Python oracle: 116 exact, 1 equal-score reorderings, 0 real mismatches

The test self-skips green on a fresh checkout until the oracle/cases*.tsv files exist, so cargo test never fails just because the oracle hasn't been generated yet. The oracle (Python sentencepiece) is only used to produce those files; the test itself, and the crate, need nothing but Rust.

Speed

A rough throughput cross-check: examples/bench.rs against oracle/bench_oracle.py, same per-line workload, release build, single thread (numbers from one machine; run it yourself for yours):

model / corpus encode, Rust encode, Py sp decode, Rust decode, Py sp
BPE / English 6.7 MB/s 2.9 MB/s 28 MB/s 5.9 MB/s
Unigram / English 10.0 MB/s 4.8 MB/s 26 MB/s 5.5 MB/s
BPE + byte / Japanese 24.5 MB/s 6.2 MB/s 50 MB/s 6.9 MB/s

Read these as a sign that the port runs at a healthy native throughput, not as proof that Rust beats C++. The Python column includes per-call CPython/FFI overhead (the C++ core in batch mode would be faster), and this crate is straightforward single-threaded inference tuned for matching the oracle rather than for speed (for example, the normaliser re-runs common_prefix_search per position). Even so, native Rust encode runs about 2 to 4 times the per-line Python path, and decode about 5 to 7 times, with no native dependency to link.

cargo run --release --example bench -- tests/models/botchan_1000_bpe.model <corpus> 50
python oracle/bench_oracle.py        tests/models/botchan_1000_bpe.model <corpus> 50

License

Apache-2.0, matching upstream SentencePiece. The bundled tests/models/*.model fixtures originate from the SentencePiece repository.