OPENNLP-1885: Pure-Java SentencePiece inference with exact original-text spans (opennlp-subword)#1165
OPENNLP-1885: Pure-Java SentencePiece inference with exact original-text spans (opennlp-subword)#1165krickert wants to merge 6 commits into
Conversation
…ginal-text spans New opennlp-extensions module implementing SentencePiece model inference without native code: the ModelProto reader, the model-embedded normalizer (precompiled character map over a Darts-clone double-array trie, whitespace collapsing and escaping, the dummy word-boundary marker), unigram best-path segmentation, BPE agenda merging, byte fallback, and user-defined symbol handling. The public contract is SubwordTokenizer/SubwordPiece; every piece reports the exact UTF-16 span of the caller's original text it came from, and the model normalizer is also exposed as an OffsetAwareNormalizer producing AlignedText. Parity with the reference implementation is asserted, not assumed: five tiny bundled models (unigram, unigram with byte fallback, BPE, identity normalization, whitespace-as-suffix) carry fixtures generated by the sentencepiece Python package over 40 inputs each, checked piece for piece, id for id, span for span, plus each model's embedded self-test samples. An opt-in test (-Dopennlp.subword.eval.dir) runs the same assertions against real downloaded models; T5-small and ALBERT-base-v2 pass exactly, including mixed scripts, emoji ZWJ sequences, BOM, and CRLF inputs.
The vocabulary trie dispatches wide nodes (the root and first level of a real vocabulary) through a 256-entry direct table, one load per byte, and scans narrow nodes' short label slices linearly instead of binary searching; a randomized differential test holds both layouts against a map-backed reference, and moving the duplicate-piece detection into the counting pass fixes the index error it previously produced. Non-unknown segments reuse the vocabulary's piece string instead of decoding their bytes, since the trie match means the bytes are identical. The normalizer precomputes, per possible first byte, whether any character-map rule or user-defined symbol starts with it; a clear bit proves the prefix machinery would pass the byte through raw, so plain ASCII text skips it entirely. The per-chunk record became a per-call scratch, the input view keeps its oversized buffers with an explicit length instead of trimming (pure-ASCII text gets an identity offset map and no map array at all), the Viterbi scratch is one interleaved array with scores as raw float bits, and the character-map trie walk relies on the JVM's own bounds checks with the fail-loud translation on the cold path. All 37 bundled parity tests and the T5-small and ALBERT real-model fixtures pass byte-identically. Single-thread throughput on the T5-small vocabulary goes from 2.83M to 6.47M pieces per second, from 0.62x to 1.42x of the reference implementation measured through its Python binding.
|
Single-thread throughput measurement after 3fb8b6b, for the record. Machine: AMD Ryzen 9 9950X3D (16 cores, one thread used), Linux, OpenJDK 25.0.3. Workload: 100,100 short texts (77 distinct lines cycled), t5-small unigram model, 32k vocabulary, 1.17M pieces total, 3 warmup passes over the corpus, then 5 timed passes.
That is 1.42x the reference on the same corpus and model, measured call for call from a host language, so the binding's per-call overhead is included in the reference number; the raw C++ core inside a batch loop is faster than that number. The Java side also does more work per token, since it maps every piece back to a UTF-16 span of the original input, which the reference does not produce against the original string. Output is unchanged: the bundled parity fixtures and the T5-small and ALBERT real-model fixtures assert identical pieces, ids, spans, and normalized forms against the reference before and after the optimization commit. |
SubwordTokenizer and SubwordPiece move to opennlp.tools.tokenize, next to Tokenizer and WordpieceTokenizer, matching where every other seam of this round lives. The opennlp-subword module keeps only the SentencePiece implementation.
…zer into it WordpieceEncoder in opennlp-api runs the full BERT tokenization pipeline as a SubwordTokenizer: every piece carries its vocabulary id and the span of the original text, surviving the normalization steps that change, insert, and remove characters. Content is computed with the same library calls the previous pipeline made; offsets come from a per-code-point rerun, with contextual case mappings (Greek final sigma) falling back to word-wide spans that widen but never misplace. List and map constructors cover line-number and explicit-id vocabularies. BertTokenizer, unreleased and superseded, is removed. The dl tokenizer creation builds on the encoder behind the existing Tokenizer plumbing via a package-private adapter with unchanged special-token selection, and a vocabulary missing its special tokens now fails at construction instead of at the first id mapping, pinned by a test. Parity is enforced twice: a differential suite against the reference pipeline (kept test-only as ReferenceBertPipeline) over a curated corpus plus 800 randomized inputs, and the removed class's reference token sequences ported case for case. WordpieceTokenizer is untouched.
Adds a new opennlp-extensions/opennlp-subword module implementing SentencePiece model inference purely in Java: the model file reader (hand-written protobuf wire parsing, no new dependency), the model-embedded normalizer (precompiled character map over the double-array trie format, whitespace collapsing and escaping, word-boundary marker), unigram best-path segmentation, BPE agenda merging, byte fallback, and user-defined symbol handling. The public contract is SubwordTokenizer/SubwordPiece; every piece reports the exact UTF-16 span of the caller's original text it came from, and the model normalizer is also exposed as an OffsetAwareNormalizer producing AlignedText.
Parity with the reference implementation is asserted, not assumed. Five tiny trained models are bundled with fixtures generated by the sentencepiece Python package (unigram, unigram with byte fallback, BPE, identity normalization, whitespace-as-suffix), checked piece for piece, id for id, span for span, plus each model's embedded self-test samples. An opt-in test (-Dopennlp.subword.eval.dir) runs the same assertions against real downloaded models; t5-small (32k vocabulary) and albert-base-v2 (30k) pass exactly, including mixed scripts, emoji ZWJ sequences, BOM, and CRLF inputs. Fixture regeneration scripts live in the test resources.
Measured single-thread throughput on the t5-small vocabulary is about 2.8M pieces/s; the native reference is about 1.6x faster. The value here is zero native dependencies, one shareable thread-safe instance, and exact original-text offsets. A non-breaking performance follow-up with identified wins is planned separately.