# sentencepiece-rust
A from-scratch, pure-Rust port of [SentencePiece](https://github.com/google/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](#the-rust-tokenizer-test) and
[Verification breadth](#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](#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
| 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](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/`](https://github.com/google/sentencepiece/tree/master/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](https://github.com/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
```sh
# in-tree toolchain (see global setup); plain `cargo` works too
cargo run --example inspect -- tests/models/botchan_1000_bpe.model
```
## Verify against the oracle
```sh
# 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](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](examples/bench.rs) against
[oracle/bench_oracle.py](oracle/bench_oracle.py), same per-line workload, release
build, single thread (numbers from one machine; run it yourself for yours):
| 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.
```sh
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.