Line data Source code
1 : // Copyright 2026 The Authors.
2 : //
3 : // Licensed under the Apache License, Version 2.0 (the "License");
4 : // you may not use this file except in compliance with the License.
5 : // You may obtain a copy of the License at
6 : //
7 : // https://www.apache.org/licenses/LICENSE-2.0
8 : //
9 : // Unless required by applicable law or agreed to in writing, software
10 : // distributed under the License is distributed on an "AS IS" BASIS,
11 : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 : // See the License for the specific language governing permissions and
13 : // limitations under the License.
14 :
15 : import 'dart:convert';
16 : import 'dart:io';
17 : import 'dart:typed_data';
18 :
19 : import 'package:dart_sentencepiece_tokenizer/dart_sentencepiece_tokenizer.dart';
20 : import 'package:meta/meta.dart';
21 :
22 : import 'bert_tokenizer.dart' show TokenizerOutput;
23 : import 'charsmap_trie.dart';
24 : import 'model_tokenizer.dart' show ModelTokenizer;
25 :
26 : /// XLM-RoBERTa-family SentencePiece/Unigram tokenizer, e.g. for
27 : /// `multilingual-e5-small`.
28 : ///
29 : /// Composes a from-scratch [CharsmapTrie] normalizer (the one piece of
30 : /// SentencePiece's pipeline this project could not find a working Dart
31 : /// implementation of anywhere — see `CharsmapTrie`'s own doc comment) with
32 : /// `dart_sentencepiece_tokenizer`'s public API for everything else: vocab
33 : /// loading, Unigram Viterbi decoding, and BOS/EOS post-processing.
34 : ///
35 : /// ## Why this class exists instead of using `dart_sentencepiece_tokenizer`
36 : /// directly
37 : ///
38 : /// `dart_sentencepiece_tokenizer` parses HuggingFace `tokenizer.json`'s
39 : /// `Precompiled` normalizer type (the charsmap trie bytes) but never applies
40 : /// it — `HuggingFaceTokenizerLoader`'s JSON-loading path silently drops the
41 : /// charsmap, making its own `SpNormalizer` an unconditional
42 : /// (charsmap-less) pass-through for models like `multilingual-e5-small`.
43 : /// See this package's `README.md` ("Why not `dart_sentencepiece_tokenizer`
44 : /// alone") for the full write-up of this and a second, independent defect
45 : /// this class also works around.
46 : ///
47 : /// [SentencePieceTokenizer]'s real constructor is private and every public
48 : /// factory routes through internal, non-injectable construction — there is
49 : /// no way to subclass or swap in a corrected normalizer. Composition (running
50 : /// our own charsmap normalization *before* handing text to the library's
51 : /// public `encode()`) is therefore the only viable integration seam, and is
52 : /// what [encode] does.
53 : ///
54 : /// ## Pipeline
55 : ///
56 : /// [encode] applies, in order:
57 : ///
58 : /// 1. [CharsmapTrie.normalize] — the charsmap substitution described above.
59 : /// 2. Whitespace-run collapse (two-or-more plain spaces → one).
60 : /// 3. Metaspace escaping (prepend a leading space if absent; replace every
61 : /// space with `▁`, U+2581).
62 : /// 4. `dart_sentencepiece_tokenizer`'s own `SentencePieceTokenizer.encode()`.
63 : ///
64 : /// **Why steps 2–3 are done manually here rather than trusted to the
65 : /// library's own `SpNormalizer`:** `dart_sentencepiece_tokenizer`'s
66 : /// HuggingFace-JSON metadata parser derives its
67 : /// `addDummyPrefix`/`removeExtraWhitespaces`/`escapeWhitespaces` flags by
68 : /// pattern-matching the `normalizer` section of `tokenizer.json`, but
69 : /// `multilingual-e5-small`'s `tokenizer.json` puts this configuration
70 : /// entirely in `pre_tokenizer` (`{"type": "Metaspace", "replacement": "▁",
71 : /// "add_prefix_space": true}`) instead — a section the parser never reads.
72 : /// All three flags therefore come back `false`, and the library's own
73 : /// whitespace/prefix handling silently no-ops for this file. This is a
74 : /// second, independent defect from the already-known charsmap-drop one
75 : /// (confirmed empirically: feeding the library literal pre-escaped text like
76 : /// `"▁Hello"` produces the correct single-token id, while plain `"Hello"` or
77 : /// `" Hello"` does not) — so this class must own both normalization
78 : /// concerns, not just the charsmap.
79 : ///
80 : /// Implements [ModelTokenizer] so [OnnxEmbeddingModel] can select this
81 : /// tokenizer at runtime alongside [BertTokenizer], both sharing the same
82 : /// [TokenizerOutput] return shape.
83 : class XlmRobertaTokenizer implements ModelTokenizer {
84 0 : XlmRobertaTokenizer._(this._charsmapTrie, this._tokenizer, this._maxLength);
85 :
86 : final CharsmapTrie _charsmapTrie;
87 : final SentencePieceTokenizer _tokenizer;
88 : final int _maxLength;
89 :
90 : /// Loads a tokenizer from a HuggingFace `tokenizer.json` file at
91 : /// [tokenizerJsonPath].
92 : ///
93 : /// Extracts the `Precompiled` normalizer's `precompiled_charsmap` bytes
94 : /// (`normalizer.normalizers[]`, `type: "Precompiled"`) to build a
95 : /// [CharsmapTrie], and separately passes the full JSON to
96 : /// `HuggingFaceTokenizerLoader.fromJsonString` to build the underlying
97 : /// [SentencePieceTokenizer] (vocabulary, Unigram model, BOS/EOS
98 : /// configuration).
99 : ///
100 : /// [maxLength] is the maximum sequence length, including the `<s>`/`</s>`
101 : /// sentinel tokens. Defaults to 512, matching both [BertTokenizer]'s own
102 : /// default and `multilingual-e5-small`'s published `max_seq_length` (the
103 : /// model's `tokenizer.json` itself declares no `truncation`/`padding`
104 : /// section, so there is no stronger signal to defer to).
105 : ///
106 : /// Throws [FormatException] if [tokenizerJsonPath]'s JSON has no
107 : /// `Precompiled` normalizer entry, or if its `precompiled_charsmap` bytes
108 : /// are malformed (see [CharsmapTrie.parse]).
109 1 : static Future<XlmRobertaTokenizer> load(
110 : String tokenizerJsonPath, {
111 : int maxLength = 512,
112 : }) async {
113 : // coverage:ignore-start
114 : // Requires the real ~17 MB multilingual-e5-small tokenizer.json (250k-entry
115 : // vocab) to exercise meaningfully — covered by the network-gated
116 : // integration test in integration_test_app/, not make coverage. See this
117 : // package's README for why that fixture isn't committed.
118 : final raw = await File(tokenizerJsonPath).readAsString();
119 : final tokenizerJson = jsonDecode(raw) as Map<String, dynamic>;
120 : final charsmapTrie = _extractCharsmapTrie(tokenizerJson);
121 : final tokenizer = HuggingFaceTokenizerLoader.fromJsonString(raw);
122 : return XlmRobertaTokenizer._(charsmapTrie, tokenizer, maxLength);
123 : // coverage:ignore-end
124 : }
125 :
126 : /// Extracts and parses the `Precompiled` normalizer's charsmap bytes from
127 : /// a parsed `tokenizer.json` map.
128 : ///
129 : /// Throws [FormatException] if no `Precompiled` normalizer entry is
130 : /// present (the `normalizer` section is expected to be a `Sequence` whose
131 : /// entries include one with `"type": "Precompiled"`, per
132 : /// `multilingual-e5-small`'s own `tokenizer.json` shape).
133 1 : static CharsmapTrie _extractCharsmapTrie(Map<String, dynamic> tokenizerJson) {
134 1 : final normalizerSection = tokenizerJson['normalizer'];
135 1 : final normalizers = normalizerSection is Map<String, dynamic>
136 1 : ? normalizerSection['normalizers']
137 : : null;
138 1 : if (normalizers is! List) {
139 : throw const FormatException(
140 : 'tokenizer.json has no normalizer.normalizers list — expected a '
141 : 'Sequence normalizer containing a Precompiled entry.',
142 : );
143 : }
144 : final precompiled = normalizers
145 1 : .cast<Map<String, dynamic>>()
146 4 : .where((n) => n['type'] == 'Precompiled')
147 1 : .firstOrNull;
148 1 : final charsmapB64 = precompiled?['precompiled_charsmap'];
149 1 : if (charsmapB64 is! String) {
150 : throw const FormatException(
151 : 'tokenizer.json has no normalizer.normalizers[] entry with '
152 : 'type "Precompiled" and a precompiled_charsmap field — this loader '
153 : 'only supports XLM-RoBERTa-family tokenizer.json files that use a '
154 : 'Precompiled charsmap normalizer.',
155 : );
156 : }
157 2 : return CharsmapTrie.parse(base64.decode(charsmapB64));
158 : }
159 :
160 : /// Encodes [text] into a [TokenizerOutput] ready for ONNX inference.
161 : ///
162 : /// Reuses the same [TokenizerOutput] type [BertTokenizer.encode] returns
163 : /// (not a parallel type), so both tokenizers already share an identical
164 : /// concrete output shape.
165 : ///
166 : /// - [TokenizerOutput.tokenTypeIds] is all-zeros: XLM-RoBERTa/RoBERTa
167 : /// models don't use segment ids, but the field is required by
168 : /// [TokenizerOutput] — this is a direct widen of
169 : /// `Encoding.typeIds`, which is already all-zeros for single-segment
170 : /// input in `dart_sentencepiece_tokenizer`, not invented data.
171 : /// - Padding uses the loaded vocabulary's own pad id (`1` for
172 : /// `multilingual-e5-small`'s `<pad>`), sourced automatically by
173 : /// `SentencePieceTokenizer` — not the BERT-specific `padId = 0`
174 : /// [BertTokenizer] uses.
175 : /// - [TokenizerOutput.truncated] is derived via two encode passes (see
176 : /// below), because `Encoding` exposes no overflow signal and a padded
177 : /// output is always exactly [maxLength] tokens long regardless of
178 : /// whether truncation actually occurred.
179 : ///
180 : /// All three output arrays have exactly the `maxLength` passed to [load]
181 : /// elements.
182 0 : @override
183 : TokenizerOutput encode(String text) {
184 0 : final normalizedText = normalizeForTokenization(_charsmapTrie, text);
185 :
186 : // coverage:ignore-start
187 : // Requires the real multilingual-e5-small vocabulary to exercise
188 : // meaningfully — covered by the network-gated integration test in
189 : // integration_test_app/, not make coverage.
190 : //
191 : // `truncated` derivation is a two-pass process. `enablePadding`/
192 : // `enableTruncation` mutate *persistent* instance state on the shared
193 : // `SentencePieceTokenizer` (they are not per-call arguments), so an
194 : // explicit noPadding()/noTruncation() reset is mandatory on *every* call
195 : // — without it, the "unbounded" first pass below would still be bounded
196 : // by whatever config a *previous* encode() call left in place, and
197 : // `truncated` would silently stick `false` after the first call.
198 : _tokenizer.noPadding();
199 : _tokenizer.noTruncation();
200 :
201 : // Pass 1 (unbounded): get the true token count with no truncation, to
202 : // detect whether pass 2's truncation will actually cut anything.
203 : final rawLength = _tokenizer
204 : .encode(normalizedText, addSpecialTokens: true)
205 : .ids
206 : .length;
207 : final truncated = rawLength > _maxLength;
208 :
209 : // Pass 2 (bounded): the real, padded/truncated result. Both passes use
210 : // identical addSpecialTokens: true so rawLength (which includes the
211 : // <s>/</s> sentinels) is compared on equal terms against the maxLength
212 : // the bounded pass truncates to.
213 : _tokenizer
214 : ..enablePadding(direction: SpPaddingDirection.right, length: _maxLength)
215 : ..enableTruncation(
216 : maxLength: _maxLength,
217 : direction: SpTruncationDirection.right,
218 : );
219 : final encoding = _tokenizer.encode(normalizedText, addSpecialTokens: true);
220 :
221 : return TokenizerOutput(
222 : inputIds: Int64List.fromList(encoding.ids),
223 : attentionMask: Int64List.fromList(encoding.attentionMask),
224 : tokenTypeIds: Int64List.fromList(encoding.typeIds),
225 : truncated: truncated,
226 : );
227 : // coverage:ignore-end
228 : }
229 :
230 : /// Applies the charsmap substitution, whitespace-run collapse, and
231 : /// Metaspace escaping steps of [encode]'s pipeline (steps 1–3), without
232 : /// requiring a loaded [SentencePieceTokenizer]/vocabulary.
233 : ///
234 : /// Exposed (rather than kept private) purely so these offline,
235 : /// vocab-independent steps can be unit-tested directly against the small
236 : /// committed charsmap fixture, without needing the full ~17 MB
237 : /// `tokenizer.json` — see `test/xlmr_tokenizer_test.dart`. Not intended
238 : /// for use outside this package or its tests.
239 1 : @visibleForTesting
240 : static String normalizeForTokenization(CharsmapTrie trie, String text) {
241 1 : final substituted = trie.normalize(text);
242 1 : final collapsed = _collapseSpaceRuns(substituted);
243 1 : return _metaspace(collapsed);
244 : }
245 :
246 : /// Collapses runs of two or more plain space characters (U+0020) down to
247 : /// one, replicating `multilingual-e5-small`'s `tokenizer.json`
248 : /// `normalizer.normalizers[1]` stage (a literal `Replace` node for the
249 : /// pattern `" {2,}"` → `" "`).
250 : ///
251 : /// Deliberately narrower than SentencePiece C++'s own generic whitespace
252 : /// collapsing (which recognises many whitespace code points and also
253 : /// trims leading/trailing runs) — this mirrors only what this
254 : /// `tokenizer.json` itself declares.
255 1 : static String _collapseSpaceRuns(String text) =>
256 2 : text.replaceAll(RegExp(' {2,}'), ' ');
257 :
258 : /// Applies `multilingual-e5-small`'s `tokenizer.json` `pre_tokenizer`
259 : /// stage: a `Metaspace` pre-tokenizer with `replacement: "▁"` and
260 : /// `add_prefix_space: true`.
261 : ///
262 : /// Prepends a single leading space if [text] doesn't already start with
263 : /// one, then replaces every space with the metaspace symbol (U+2581).
264 : ///
265 : /// **Empty input is a special case, found by WI-4's broader parity
266 : /// corpus** (`test/fixtures/xlmr_parity_corpus.json`'s `edge_empty`
267 : /// entry): real `AutoTokenizer` output for `""` is just `[<s>, </s>]`
268 : /// (`[0, 2]`) with no content token at all, but naively adding the dummy
269 : /// prefix here would turn `""` into `"▁"`, which the real vocabulary
270 : /// treats as a *valid standalone piece* (id 6) -- producing a spurious
271 : /// extra content token (`[0, 6, 2]`) that fails byte-exact parity.
272 : /// HuggingFace's `tokenizers` Rust `Metaspace` pre-tokenizer only adds the
273 : /// prefix space when pre-tokenizing an actual word/split; a fully empty
274 : /// input has no splits to begin with, so the prefix is never added.
275 : /// Returning `''` unchanged for empty input reproduces that behaviour
276 : /// without needing to port the Rust library's split-detection machinery.
277 1 : static String _metaspace(String text) {
278 1 : if (text.isEmpty) return text;
279 2 : final withPrefix = text.startsWith(' ') ? text : ' $text';
280 1 : return withPrefix.replaceAll(' ', '▁');
281 : }
282 : }
|
