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_onnxrt/betto_onnxrt.dart';
19 : import 'package:betto_lexical/betto_lexical.dart' show Tokenizer;
20 : import 'package:meta/meta.dart' show visibleForTesting;
21 : import 'package:path/path.dart' as p;
22 :
23 : import 'embedding_model.dart' show EmbeddingKind, EmbeddingModel;
24 :
25 : import 'bert_tokenizer.dart';
26 : import 'math_utils.dart';
27 : import 'model_catalog.dart';
28 : import 'model_tokenizer.dart' show ModelTokenizer;
29 : import 'xlmr_tokenizer.dart' show XlmRobertaTokenizer;
30 :
31 : /// ONNX Runtime-backed embedding model for dense text retrieval.
32 : ///
33 : /// Implements [EmbeddingModel] using a model from [ModelCatalog] via the
34 : /// `betto_onnxrt` [OnnxRuntime] and [OnnxSession] API. Produces L2-normalised
35 : /// float32 embeddings suitable for cosine similarity search.
36 : ///
37 : /// ## Model identity
38 : ///
39 : /// [modelId] returns the stable [ModelSpec.id] of the loaded model (e.g.
40 : /// `bge-small-en-v1.5`). This should be persisted alongside the vector index
41 : /// so that a model change can be detected and the index rebuilt.
42 : ///
43 : /// ## Loading with download-on-demand (preferred)
44 : ///
45 : /// Supply a [cacheDir] (and optionally a [ModelSpec] via [spec]) to fetch the
46 : /// model on first use via [ModelDownloader]. If the model files are already
47 : /// cached and their SHA-256 checksums match, they are used immediately.
48 : /// Otherwise [ModelDownloader] fetches the files before opening the ORT
49 : /// session:
50 : ///
51 : /// ```dart
52 : /// final spec = ModelCatalog.lookup('bge-small-en-v1.5');
53 : /// final model = await OnnxEmbeddingModel.load(
54 : /// spec: spec,
55 : /// cacheDir: '/path/to/cache',
56 : /// onProgress: (received, total) {
57 : /// stderr.writeln('Downloading: $received / $total bytes');
58 : /// },
59 : /// );
60 : /// ```
61 : ///
62 : /// ## Loading from an explicit path
63 : ///
64 : /// The [modelPath] parameter loads a model from a specific filesystem path,
65 : /// bypassing the catalog and downloader. Specifying [modelPath] without [spec]
66 : /// uses [ModelCatalog.defaultModelId] for the identity.
67 : ///
68 : /// ```dart
69 : /// final model = await OnnxEmbeddingModel.load(
70 : /// modelPath: '/path/to/bge_small.onnx',
71 : /// );
72 : /// ```
73 : ///
74 : /// **Important:** Either [modelPath] or [cacheDir] must be supplied.
75 : /// Calling [load] without either throws [ArgumentError] synchronously — there
76 : /// is no bundled model asset path. See [ModelCatalog] and [ModelDownloader].
77 : ///
78 : /// ## Lifecycle
79 : ///
80 : /// [load] opens the native ORT session via [OnnxRuntime.load]. [embed] runs
81 : /// synchronously on the calling isolate — do **not** call from the UI thread
82 : /// in Flutter without isolate offloading. [dispose] releases native resources;
83 : /// always call it (use `try/finally`).
84 : ///
85 : /// ## Thread safety
86 : ///
87 : /// ORT sessions are thread-affine. All [embed] and [dispose] calls must come
88 : /// from the same isolate that called [load].
89 : class OnnxEmbeddingModel implements EmbeddingModel {
90 : /// Internal constructor — use [load].
91 0 : OnnxEmbeddingModel._(
92 : this._runtime,
93 : this._session,
94 : this._tokenizer,
95 : this._spec,
96 : );
97 :
98 : final OnnxRuntime _runtime;
99 : final OnnxSession _session;
100 : final ModelTokenizer _tokenizer;
101 :
102 : /// The [ModelSpec] of the loaded model.
103 : ///
104 : /// Provides [modelId] and [dimensions] for the [EmbeddingModel] interface.
105 : final ModelSpec _spec;
106 :
107 : // ── EmbeddingModel interface ───────────────────────────────────────────────
108 :
109 : /// Stable identifier of the loaded model, matching a [ModelCatalog] entry.
110 : ///
111 : /// Should be persisted alongside the vector index so a later model swap can
112 : /// be detected and the index rebuilt. Example: `bge-small-en-v1.5`.
113 0 : @override
114 : String get modelId => _spec.id; // coverage:ignore-line
115 :
116 : /// Embedding vector length produced by this model.
117 : ///
118 : /// Sourced from `spec.meta['dimensions']`. This is the single source of
119 : /// truth for SQ8 byte lengths and score-path length guards.
120 : /// Example: 384 for BGE Small En v1.5.
121 0 : @override
122 : int get dimensions => _spec.meta['dimensions'] as int; // coverage:ignore-line
123 :
124 : // ── Factory ────────────────────────────────────────────────────────────────
125 :
126 : /// Loads an embedding model and returns an [OnnxEmbeddingModel].
127 : ///
128 : /// **Either [modelPath] or [cacheDir] must be supplied.** Calling [load]
129 : /// without either throws [ArgumentError] synchronously — there is no default
130 : /// asset path or bundled model. Use [cacheDir] to download the model on
131 : /// demand (preferred), or [modelPath] to load from an explicit filesystem
132 : /// path. See [ModelCatalog] and [ModelDownloader].
133 : ///
134 : /// ## Download-on-demand path (preferred)
135 : ///
136 : /// When [cacheDir] is provided, [ModelDownloader] is invoked to ensure the
137 : /// model files are present and checksummed before opening the ORT session.
138 : /// Files already in the cache are reused without downloading. Pass [spec] to
139 : /// select a specific catalog model; if omitted, [ModelCatalog.defaultModelId]
140 : /// (`'bge-small-en-v1.5'`) is used.
141 : ///
142 : /// [onProgress] is forwarded to [ModelDownloader.ensure] and receives
143 : /// incremental download progress. It is not called when files are cached.
144 : ///
145 : /// ```dart
146 : /// final model = await OnnxEmbeddingModel.load(
147 : /// cacheDir: '/path/to/cache',
148 : /// onProgress: (received, total) => print('$received / $total'),
149 : /// );
150 : /// ```
151 : ///
152 : /// ## Explicit-path
153 : ///
154 : /// When [modelPath] is provided, the file at that path is loaded directly
155 : /// (no download, no checksum). The model identity is set to
156 : /// [ModelCatalog.defaultModelId] unless [spec] is also supplied.
157 : ///
158 : /// **The second asset file must be named `vocab.txt`, beside [modelPath],
159 : /// regardless of `tokenizerFamily`.** This explicit-path mode is a
160 : /// lower-level escape hatch than the [cacheDir] path below and does not
161 : /// branch on tokenizer family for the asset filename — so loading an
162 : /// `'xlmr'`-family model this way still requires its `tokenizer.json` to
163 : /// be named/placed as `vocab.txt` beside [modelPath]. The [cacheDir] path
164 : /// handles per-family asset naming correctly and is the recommended way
165 : /// to load an XLM-R-family model.
166 : ///
167 : /// [tokenizer] overrides the word-segmentation step inside [BertTokenizer].
168 : /// Defaults to [RegExpTokenizer]. Supply `IcuTokenizer()` from
169 : /// `package:betto_icu` for superior Unicode coverage. Ignored for models
170 : /// whose `tokenizerFamily` is `'xlmr'` (word segmentation there is handled
171 : /// entirely by [XlmRobertaTokenizer]'s SentencePiece/Unigram pipeline, which
172 : /// has no equivalent seam).
173 : ///
174 : /// ## Tokenizer family selection
175 : ///
176 : /// The concrete [ModelTokenizer] implementation is chosen from
177 : /// `spec.meta['tokenizerFamily']`: `'bert'` loads [BertTokenizer],
178 : /// `'xlmr'` loads [XlmRobertaTokenizer]. This key must be present and
179 : /// recognised — an absent or unrecognised value throws [ArgumentError]
180 : /// rather than silently defaulting, so a future third tokenizer family
181 : /// can't be misresolved by accident.
182 : ///
183 : /// Throws [ArgumentError] if neither [modelPath] nor [cacheDir] is supplied,
184 : /// or if [spec]`.meta['tokenizerFamily']` is missing or unrecognised.
185 : /// Throws [UnsupportedError] if the model file does not exist on disk.
186 : /// Throws [Exception] if the ORT library cannot be loaded or the model is
187 : /// corrupt.
188 2 : static Future<OnnxEmbeddingModel> load({
189 : ModelSpec? spec,
190 : String? cacheDir,
191 : String? modelPath,
192 : Tokenizer? tokenizer,
193 : DownloadProgress? onProgress,
194 : }) async {
195 : // Guard: at least one of modelPath or cacheDir must be supplied.
196 : // This is a required-argument check: throw synchronously before any I/O
197 : // so callers get a fast, clear failure rather than a confusing downstream
198 : // error. The bundled LFS asset path has been removed — download-on-demand
199 : // (cacheDir) or an explicit modelPath is the only supported mechanism.
200 : if (modelPath == null && cacheDir == null) {
201 1 : throw ArgumentError(
202 : 'Either modelPath or cacheDir must be supplied. '
203 : 'Pass an explicit modelPath, or pass cacheDir (with an optional spec) '
204 : 'to download the model on demand. '
205 : 'See ModelCatalog and ModelDownloader.',
206 : );
207 : }
208 :
209 : // Resolve the model spec. When no spec is given, use the default model.
210 : // When a raw modelPath is supplied without a spec, we still need an id for
211 : // model identity tracking — use the default catalog ID.
212 : final resolvedSpec =
213 1 : spec ?? ModelCatalog.lookup(ModelCatalog.defaultModelId);
214 :
215 : // Validate tokenizerFamily synchronously, before any I/O — same
216 : // fail-fast rationale as the modelPath/cacheDir guard above. This is a
217 : // pure, fast check (no file or network access), so it is fully covered
218 : // by unit tests rather than requiring live model assets.
219 2 : final tokenizerFamily = _tokenizerFamily(resolvedSpec);
220 :
221 : final String resolvedModelPath;
222 : final String resolvedVocabPath;
223 :
224 : if (modelPath != null) {
225 : // Explicit path — bypass catalog and downloader. The second asset file
226 : // is named 'vocab.txt' regardless of tokenizer family for this path —
227 : // an explicit modelPath is a lower-level escape hatch than the
228 : // catalog/downloader path below, so it keeps the original, simpler
229 : // convention rather than branching on tokenizerFamily too.
230 : resolvedModelPath = modelPath;
231 4 : resolvedVocabPath = p.join(p.dirname(modelPath), 'vocab.txt');
232 : } else {
233 : // cacheDir != null is guaranteed by the guard above.
234 : // Download-on-demand path: let ModelDownloader ensure the files are
235 : // present and their checksums match before opening the ORT session.
236 : // The ModelCatalog allowlist gates which models may be downloaded.
237 0 : final downloader = ModelDownloader(allowlist: ModelCatalog());
238 0 : final resolved = await downloader.ensure(
239 : resolvedSpec,
240 : cacheDir: cacheDir!,
241 : onProgress: onProgress,
242 : );
243 : // File names in ResolvedModel.filePaths match the keys in ModelSpec.files:
244 : // 'onnx' → absolute path to the .onnx model, 'vocab' → the tokenizer
245 : // asset (vocab.txt for BERT-family models, tokenizer.json for
246 : // XLM-R-family models — the key name is a generic "second asset" slot,
247 : // not tied to any one file format).
248 0 : resolvedModelPath = resolved.filePaths['onnx']!;
249 0 : resolvedVocabPath = resolved.filePaths['vocab']!;
250 : }
251 :
252 2 : _assertFileExists(resolvedModelPath, 'model file');
253 1 : _assertFileExists(resolvedVocabPath, 'tokenizer asset');
254 :
255 : // coverage:ignore-start
256 : // The lines below require a live ORT native library and model assets.
257 : // They are tested through integration tests when model assets are present.
258 : final runtime = await OnnxRuntime.load();
259 : final session = runtime.createSessionFromFile(resolvedModelPath);
260 : // tokenizerFamily was already validated above (fail-fast, before I/O) —
261 : // this switch only selects which concrete loader to invoke.
262 : final ModelTokenizer tok = tokenizerFamily == 'xlmr'
263 : ? await XlmRobertaTokenizer.load(resolvedVocabPath)
264 : : await BertTokenizer.load(resolvedVocabPath, tokenizer: tokenizer);
265 : return OnnxEmbeddingModel._(runtime, session, tok, resolvedSpec);
266 : // coverage:ignore-end
267 : }
268 :
269 : /// Validates and returns `spec.meta['tokenizerFamily']`.
270 : ///
271 : /// Pure and synchronous — safe to call before any I/O, so a misconfigured
272 : /// [ModelSpec] fails fast with a clear message rather than surfacing as a
273 : /// confusing downstream error once the ORT session is already open.
274 : ///
275 : /// Throws [ArgumentError] if the key is missing or is not one of the
276 : /// recognised values (`'bert'`, `'xlmr'`) — deliberately not defaulted, so
277 : /// a future third tokenizer family can't be silently misresolved by a
278 : /// [ModelSpec] that forgot to set it.
279 2 : static String _tokenizerFamily(ModelSpec spec) {
280 4 : final family = spec.meta['tokenizerFamily'];
281 3 : if (family != 'bert' && family != 'xlmr') {
282 2 : throw ArgumentError(
283 1 : "ModelSpec '${spec.id}' has meta['tokenizerFamily'] = "
284 1 : "${family == null ? 'null (missing)' : "'$family'"}, but only "
285 : "'bert' and 'xlmr' are recognised. Add a valid tokenizerFamily "
286 : 'entry to the ModelSpec.meta map.',
287 : );
288 : }
289 : return family as String;
290 : }
291 :
292 : /// Prepends the `kind`-appropriate prefix from `spec.meta` to [text], if
293 : /// one is configured.
294 : ///
295 : /// Looks up `spec.meta['queryPrefix']` for [EmbeddingKind.query] or
296 : /// `spec.meta['documentPrefix']` for [EmbeddingKind.document]. If the
297 : /// corresponding key is absent, [text] is returned unchanged — this is a
298 : /// deliberate no-op default so models that don't need a prefix (e.g.
299 : /// `bge-small-en-v1.5`) are byte-for-byte unaffected by [kind].
300 : ///
301 : /// Exposed (rather than kept private) so this pure, spec-driven logic can
302 : /// be unit-tested directly without needing a live ORT session — see this
303 : /// package's coverage notes for why [embed] itself is `coverage:ignore`d.
304 1 : @visibleForTesting
305 : static String applyPrefix(String text, EmbeddingKind kind, ModelSpec spec) {
306 1 : final key = kind == EmbeddingKind.query ? 'queryPrefix' : 'documentPrefix';
307 2 : final prefix = spec.meta[key];
308 2 : return prefix is String ? '$prefix$text' : text;
309 : }
310 :
311 : // ── EmbeddingModel.embed ──────────────────────────────────────────────────
312 :
313 : /// Embeds [text] into an L2-normalised float32 vector of [dimensions] elements.
314 : ///
315 : /// Runs synchronously on the calling isolate. For large batches or UI
316 : /// applications, wrap calls in [Isolate.run] — but note that ORT sessions
317 : /// are thread-affine, so the session must be created inside the same isolate
318 : /// that calls [embed].
319 : ///
320 : /// An empty or whitespace-only [text] produces a `[CLS][SEP]`-only
321 : /// embedding (two real tokens) and returns `truncated = false`.
322 : ///
323 : /// [kind] selects [spec.meta]'s `'queryPrefix'` (for
324 : /// [EmbeddingKind.query]) or `'documentPrefix'` (for
325 : /// [EmbeddingKind.document], the default) and prepends it to [text] before
326 : /// tokenisation — see [applyPrefix]. Models with neither key (e.g.
327 : /// `bge-small-en-v1.5`) are unaffected: the prepend is a no-op, so
328 : /// behaviour is byte-for-byte unchanged regardless of [kind].
329 : ///
330 : /// Returns `(embedding, truncated)`:
331 : /// - `embedding` — [dimensions]-element [Float32List] with unit L2 norm.
332 : /// - `truncated` — `true` if [text] (after prefixing) exceeded the usable
333 : /// token budget and was silently cut before embedding.
334 : // coverage:ignore-start
335 : // This entire method requires a live ORT session to exercise meaningfully
336 : // — covered by integration tests with model assets, not make coverage.
337 : // (applyPrefix itself, the pure prefix-selection logic embed() delegates
338 : // to below, is unit-tested directly without needing a live session — see
339 : // its own doc comment and test/onnx_embedding_model_test.dart.)
340 : @override
341 : Future<(Float32List, bool)> embed(
342 : String text, {
343 : EmbeddingKind kind = EmbeddingKind.document,
344 : }) async {
345 : final prefixedText = applyPrefix(text, kind, _spec);
346 : final tokens = _tokenizer.encode(prefixedText);
347 : final seqLen = tokens.inputIds.length;
348 : final hiddenDim = dimensions; // sourced from spec.meta['dimensions']
349 :
350 : // Build int64 input tensors shaped [1, seqLen].
351 : // The BGE model requires three inputs: input_ids, attention_mask,
352 : // and token_type_ids — all shaped [1, seqLen] with int64 elements.
353 : final shape = [1, seqLen];
354 : final inputIds = OnnxTensor.fromInt64(shape, tokens.inputIds);
355 : final attentionMask = OnnxTensor.fromInt64(shape, tokens.attentionMask);
356 : final tokenTypeIds = OnnxTensor.fromInt64(shape, tokens.tokenTypeIds);
357 :
358 : // Run ONNX inference. The output 'last_hidden_state' has shape
359 : // [1, seqLen, hiddenDim]. We rely on OnnxSession.run() populating
360 : // the shape from the native OrtValue via the output-shape-readback
361 : // slots (31/32/33) added in the generic betto_onnxrt API.
362 : final outputs = _session.run(
363 : inputs: {
364 : 'input_ids': inputIds,
365 : 'attention_mask': attentionMask,
366 : 'token_type_ids': tokenTypeIds,
367 : },
368 : outputNames: ['last_hidden_state'],
369 : );
370 : final outputTensor = outputs.first;
371 :
372 : // Extract the flat float32 output. The tensor shape is [1, seqLen, D].
373 : // asFloat32() throws StateError if the output element type is not float32,
374 : // which would indicate a mismatched model.
375 : final raw = outputTensor.asFloat32().toList();
376 :
377 : // Mean-pool over non-padding token positions, then L2-normalise.
378 : final pooled = meanPool(
379 : raw,
380 : tokens.attentionMask.toList(),
381 : seqLen: seqLen,
382 : hiddenDim: hiddenDim,
383 : );
384 : final embedding = l2Normalize(pooled);
385 :
386 : return (embedding, tokens.truncated);
387 : // coverage:ignore-end
388 : }
389 :
390 : /// Releases the native ORT session and runtime resources.
391 : ///
392 : /// Must be called exactly once when the model is no longer needed.
393 : /// After [dispose], [embed] must not be called.
394 0 : @override
395 : void dispose() {
396 : // coverage:ignore-start
397 : // Requires a live ORT session — covered by integration tests with model assets.
398 : _session.dispose();
399 : _runtime.dispose();
400 : // coverage:ignore-end
401 : }
402 :
403 : // ── Private helpers ────────────────────────────────────────────────────────
404 :
405 2 : static void _assertFileExists(String path, String label) {
406 4 : if (!File(path).existsSync()) {
407 4 : throw UnsupportedError(
408 : '$label not found at: $path\n'
409 : 'Ensure model assets are present or configure a cacheDir for '
410 : 'download-on-demand. See ModelCatalog and ModelDownloader.',
411 : );
412 : }
413 : }
414 : }
|
