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:io';
16 : import 'dart:typed_data';
17 :
18 : import 'package:betto_lexical/betto_lexical.dart'
19 : show Tokenizer, RegExpTokenizer;
20 :
21 : import 'model_tokenizer.dart' show ModelTokenizer;
22 :
23 : /// A BERT WordPiece tokenizer backed by a `vocab.txt` file.
24 : ///
25 : /// Converts arbitrary text into BERT token IDs suitable for feeding into the
26 : /// BGE Small En v1.5 ONNX model via [OnnxSession.run].
27 : ///
28 : /// ## Pipeline
29 : ///
30 : /// 1. **Normalise** — lower-case and strip combining accent characters.
31 : /// 2. **Word segmentation** — delegate to the [Tokenizer] supplied at
32 : /// construction time. [RegExpTokenizer] is used by default; `IcuTokenizer`
33 : /// from `package:betto_icu` can be substituted as a drop-in
34 : /// replacement for superior Unicode coverage.
35 : /// 3. **WordPiece** — split each word into sub-word pieces and look up IDs in
36 : /// the vocabulary loaded from `vocab.txt`. Unknown pieces map to `[UNK]`.
37 : /// 4. **Assemble** — prepend `[CLS]` (101), append `[SEP]` (102), and pad to
38 : /// [maxLength] with `[PAD]` (0).
39 : ///
40 : /// ## Token ID space
41 : ///
42 : /// BERT token IDs are entirely distinct from the stemmed token strings
43 : /// produced by the lexical search pipeline (FtsManager / BM25). They must
44 : /// not be interchanged.
45 : ///
46 : /// ## IcuTokenizer
47 : ///
48 : /// ```dart
49 : /// import 'package:betto_icu/betto_icu.dart';
50 : /// final tokenizer = await BertTokenizer.load(vocabPath,
51 : /// tokenizer: IcuTokenizer());
52 : /// ```
53 : class BertTokenizer implements ModelTokenizer {
54 : final Map<String, int> _vocab;
55 : final int _maxLength;
56 : final Tokenizer _tokenizer;
57 :
58 : /// `[CLS]` token ID — always the first token in BERT input sequences.
59 : static const int clsId = 101;
60 :
61 : /// `[SEP]` token ID — marks the end of a segment in BERT input sequences.
62 : static const int sepId = 102;
63 :
64 : /// `[UNK]` token ID — substituted for vocabulary entries not found in WordPiece
65 : /// decomposition.
66 : static const int unkId = 100;
67 :
68 : /// `[PAD]` token ID — used to fill sequences shorter than [maxLength].
69 : static const int padId = 0;
70 :
71 2 : BertTokenizer._(this._vocab, this._maxLength, this._tokenizer);
72 :
73 : /// Loads the vocabulary from [vocabPath] and returns a [BertTokenizer].
74 : ///
75 : /// [vocabPath] must point to a `vocab.txt` file where each line is a
76 : /// vocabulary token and the line index (0-based) is the token ID. The BGE
77 : /// Small En v1.5 vocabulary has 30,522 entries.
78 : ///
79 : /// [maxLength] is the maximum sequence length including the `[CLS]` and
80 : /// `[SEP]` sentinel tokens (default 512 per the BERT specification).
81 : ///
82 : /// [tokenizer] controls word segmentation before WordPiece splitting.
83 : /// Defaults to [RegExpTokenizer]. Supply `IcuTokenizer()` from
84 : /// `package:betto_icu` for improved Unicode coverage.
85 2 : static Future<BertTokenizer> load(
86 : String vocabPath, {
87 : int maxLength = 512,
88 : Tokenizer? tokenizer,
89 : }) async {
90 4 : final lines = await File(vocabPath).readAsLines();
91 2 : final vocab = <String, int>{};
92 6 : for (var i = 0; i < lines.length; i++) {
93 6 : vocab[lines[i].trim()] = i;
94 : }
95 4 : return BertTokenizer._(vocab, maxLength, tokenizer ?? RegExpTokenizer());
96 : }
97 :
98 : /// Encodes [text] into a [TokenizerOutput] ready for ONNX inference.
99 : ///
100 : /// The output always starts with `[CLS]` (101) and ends with `[SEP]` (102).
101 : /// If [text] contains more WordPiece tokens than `maxLength - 2` (510 usable
102 : /// tokens), the excess is silently discarded and
103 : /// [TokenizerOutput.truncated] is `true`.
104 : ///
105 : /// An empty or whitespace-only [text] produces a two-token sequence
106 : /// `[CLS][SEP]` with all remaining positions padded — [TokenizerOutput.truncated]
107 : /// is `false`.
108 : ///
109 : /// All three output arrays ([TokenizerOutput.inputIds],
110 : /// [TokenizerOutput.attentionMask], [TokenizerOutput.tokenTypeIds]) have
111 : /// exactly [maxLength] elements.
112 2 : @override
113 : TokenizerOutput encode(String text) {
114 2 : final normalized = _normalize(text);
115 4 : final words = _tokenizer.tokenise(normalized);
116 :
117 : // Build the token ID list starting with [CLS].
118 : // Leave one slot for the closing [SEP] token.
119 2 : final tokenIds = <int>[clsId];
120 : var wasTruncated = false;
121 :
122 : outer:
123 4 : for (final word in words) {
124 2 : if (word.isEmpty) continue;
125 4 : for (final id in _wordPiece(word)) {
126 : // Reserve the last slot for [SEP].
127 8 : if (tokenIds.length >= _maxLength - 1) {
128 : wasTruncated = true;
129 : break outer;
130 : }
131 2 : tokenIds.add(id);
132 : }
133 : }
134 2 : tokenIds.add(sepId);
135 :
136 : // Build attention mask: 1 for real tokens, 0 for padding.
137 4 : final attentionMask = List<int>.filled(tokenIds.length, 1, growable: true);
138 6 : while (tokenIds.length < _maxLength) {
139 2 : tokenIds.add(padId);
140 2 : attentionMask.add(0);
141 : }
142 :
143 2 : return TokenizerOutput(
144 2 : inputIds: Int64List.fromList(tokenIds),
145 2 : attentionMask: Int64List.fromList(attentionMask),
146 : // BERT token_type_ids are all-zeros for single-segment input.
147 6 : tokenTypeIds: Int64List.fromList(List.filled(_maxLength, 0)),
148 : truncated: wasTruncated,
149 : );
150 : }
151 :
152 : /// Decodes a list of token IDs back to vocabulary strings.
153 : ///
154 : /// Unknown IDs are mapped to `'[UNK]'`. Primarily for diagnostics.
155 1 : List<String> decode(List<int> ids) {
156 1 : final inverse = <int, String>{};
157 4 : _vocab.forEach((k, v) => inverse[v] = k);
158 4 : return ids.map((id) => inverse[id] ?? '[UNK]').toList();
159 : }
160 :
161 : // ── Private helpers ─────────────────────────────────────────────────────────
162 :
163 : /// Lower-cases [text] and strips Unicode combining accent characters
164 : /// (U+0300–U+036F) which BERT treats as noise.
165 2 : String _normalize(String text) {
166 2 : final buf = StringBuffer();
167 6 : for (final char in text.toLowerCase().runes) {
168 : // Strip combining diacritical marks (accents, cedillas, etc.)
169 3 : if (char >= 0x0300 && char <= 0x036F) continue;
170 2 : buf.writeCharCode(char);
171 : }
172 2 : return buf.toString();
173 : }
174 :
175 : /// Splits [word] into sub-word pieces using the WordPiece algorithm and
176 : /// returns the corresponding vocabulary token IDs.
177 : ///
178 : /// All sub-word pieces after the first are prefixed with `##` per the BERT
179 : /// convention. Returns `[unkId]` if any position cannot be decomposed.
180 2 : List<int> _wordPiece(String word) {
181 : // Fast path: the whole word is in the vocabulary.
182 10 : if (_vocab.containsKey(word)) return [_vocab[word]!];
183 :
184 1 : final ids = <int>[];
185 : var start = 0;
186 2 : while (start < word.length) {
187 1 : var end = word.length;
188 : int? found;
189 1 : while (start < end) {
190 1 : final sub = start == 0
191 1 : ? word.substring(start, end)
192 2 : : '##${word.substring(start, end)}';
193 2 : if (_vocab.containsKey(sub)) {
194 2 : found = _vocab[sub];
195 : break;
196 : }
197 1 : end--;
198 : }
199 : // If no sub-word piece was found, map the entire word to [UNK].
200 1 : if (found == null) return [unkId];
201 1 : ids.add(found);
202 : start = end;
203 : }
204 : return ids;
205 : }
206 : }
207 :
208 : /// The output of [BertTokenizer.encode]: three parallel int64 arrays ready for
209 : /// ONNX Runtime inference.
210 : ///
211 : /// All three arrays have exactly `maxLength` elements. Padding positions have
212 : /// `inputIds = 0`, `attentionMask = 0`, `tokenTypeIds = 0`.
213 : final class TokenizerOutput {
214 : /// Creates a [TokenizerOutput].
215 2 : const TokenizerOutput({
216 : required this.inputIds,
217 : required this.attentionMask,
218 : required this.tokenTypeIds,
219 : required this.truncated,
220 : });
221 :
222 : /// BERT token IDs, starting with `[CLS]` (101) and ending with `[SEP]`
223 : /// (102), then zero-padded to [BertTokenizer.maxLength].
224 : final Int64List inputIds;
225 :
226 : /// 1 for real tokens (including `[CLS]` and `[SEP]`), 0 for padding.
227 : final Int64List attentionMask;
228 :
229 : /// Segment IDs — all zeros for single-segment BERT input.
230 : final Int64List tokenTypeIds;
231 :
232 : /// `true` if the input text exceeded the usable token budget and was
233 : /// silently truncated before the `[SEP]` token.
234 : final bool truncated;
235 : }
|
