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 'package:betto_onnxrt/betto_onnxrt.dart';
16 :
17 : /// Allowlist of supported embedding models for dense retrieval.
18 : ///
19 : /// [ModelCatalog] is the single place where models are registered and the
20 : /// concrete [AllowlistProvider] implementation used with [ModelDownloader].
21 : /// All models must appear here before they can be downloaded — attempting to
22 : /// look up an unregistered model ID throws [ArgumentError]. Attempting to load
23 : /// a model whose validation flag is `false` throws [UnsupportedError].
24 : ///
25 : /// ## Why AllowlistProvider
26 : ///
27 : /// [ModelCatalog] implements [AllowlistProvider] from `betto_onnxrt` so that
28 : /// [ModelDownloader] can be constructed with `allowlist: ModelCatalog()` and
29 : /// will reject any model not in this catalog before touching the network.
30 : ///
31 : /// ## Adding a new model
32 : ///
33 : /// 1. Add a private `ModelSpec` field below (as a `static final`).
34 : /// 2. Insert it into the [_catalog] map with its ID as the key.
35 : /// 3. Add `'<id>': false` to [_validated] until the model has been tested in
36 : /// CI; flip it to `true` when the validation plan is complete.
37 : ///
38 : /// ## Usage
39 : ///
40 : /// ```dart
41 : /// final spec = ModelCatalog.lookup('bge-small-en-v1.5');
42 : /// print(spec.meta['dimensions']); // 384
43 : ///
44 : /// // Use with ModelDownloader to gate downloads:
45 : /// final downloader = ModelDownloader(allowlist: ModelCatalog());
46 : /// ```
47 : final class ModelCatalog implements AllowlistProvider {
48 : /// Creates a [ModelCatalog].
49 : ///
50 : /// The catalog is stateless and lightweight — create a new instance wherever
51 : /// needed, or share a single instance.
52 2 : const ModelCatalog();
53 :
54 : // ── Registered models ──────────────────────────────────────────────────────
55 : // Note: ModelSpec / ModelFile cannot be const because Uri(...) is not const
56 : // in Dart. Use static final (lazily initialised) instead.
57 :
58 : /// BGE Small En v1.5 (BAAI).
59 : ///
60 : /// 384-dimensional English-language sentence embeddings optimised for
61 : /// retrieval. ~127 MB ONNX binary. **Validated and production-ready.**
62 9 : static final _bgeSmallEnV15 = ModelSpec(
63 : id: 'bge-small-en-v1.5',
64 3 : files: {
65 3 : 'onnx': ModelFile(
66 3 : url: Uri.parse(
67 : 'https://huggingface.co/BAAI/bge-small-en-v1.5/resolve/main/onnx/model.onnx',
68 : ),
69 : // SHA-256 of the exact model file used in CI; update if upstream
70 : // changes.
71 : sha256:
72 : '828e1496d7fabb79cfa4dcd84fa38625c0d3d21da474a00f08db0f559940cf35',
73 : ),
74 3 : 'vocab': ModelFile(
75 3 : url: Uri.parse(
76 : 'https://huggingface.co/BAAI/bge-small-en-v1.5/resolve/main/vocab.txt',
77 : ),
78 : sha256:
79 : '07eced375cec144d27c900241f3e339478dec958f92fddbc551f295c992038a3',
80 : ),
81 : },
82 3 : meta: {
83 : // Embedding vector dimension. Read by OnnxEmbeddingModel as
84 : // `spec.meta['dimensions'] as int`.
85 : 'dimensions': 384,
86 : // Selects BertTokenizer in OnnxEmbeddingModel.load(). Set explicitly
87 : // (not left absent) so a future third tokenizer family can't be
88 : // silently misresolved by a ModelSpec that forgot this key.
89 : 'tokenizerFamily': 'bert',
90 : // No queryPrefix/documentPrefix — BGE Small En v1.5 has no
91 : // passage/query prefix convention, so OnnxEmbeddingModel.embed()'s
92 : // prefix step is a no-op for this model regardless of EmbeddingKind.
93 : },
94 : );
95 :
96 : /// `multilingual-e5-small` (intfloat) — multilingual, 384-dimensional.
97 : ///
98 : /// Same embedding dimension as [_bgeSmallEnV15], so adopting it requires no
99 : /// SQ8/index-format change in a consuming database. Registered as this
100 : /// project's first cross-lingual embedding model (WI-4); XLM-RoBERTa-family
101 : /// SentencePiece/Unigram tokenisation via [XlmRobertaTokenizer].
102 : ///
103 : /// **Registers the plain fp32 `model.onnx` export (~470 MB) — not a
104 : /// quantized or GPU-oriented variant.** `intfloat/multilingual-e5-small`'s
105 : /// `onnx/` directory also publishes `model_qint8_avx512_vnni.onnx`
106 : /// (x86 AVX512-VNNI int8, a poor fit for this project's largely-ARM
107 : /// targets) and `model_O4.onnx` (an ORT graph-optimizer export with fp16
108 : /// mixed precision intended for GPU inference, not the CPU-only execution
109 : /// providers `betto_onnxrt` supports). The plain export matches
110 : /// [_bgeSmallEnV15]'s own registration precedent (also a plain fp32
111 : /// `model.onnx`) and avoids stacking an unvalidated second source of
112 : /// numerical drift underneath the storage layer's own SQ8 quantization.
113 : /// Do not "helpfully" swap in a smaller quantized/optimized variant without
114 : /// a dedicated accuracy-validation pass — see
115 : /// `docs/plans/plan_0_06_wi4_multilingual_embedding_model.md` (in the
116 : /// `kmdb` repo) for the full trade-off analysis.
117 : ///
118 : /// **Note the meaningfully larger download** compared to BGE Small En v1.5
119 : /// (~470 MB vs. ~127 MB, ~3.7×) — worth surfacing in any first-use download
120 : /// UX, especially on mobile.
121 9 : static final _multilingualE5Small = ModelSpec(
122 : id: 'multilingual-e5-small',
123 3 : files: {
124 3 : 'onnx': ModelFile(
125 3 : url: Uri.parse(
126 : 'https://huggingface.co/intfloat/multilingual-e5-small/resolve/main/onnx/model.onnx',
127 : ),
128 : // SHA-256 of the exact model file used in CI; update if upstream
129 : // changes.
130 : sha256: _multilingualE5SmallOnnxSha256,
131 : ),
132 3 : 'vocab': ModelFile(
133 3 : url: Uri.parse(
134 : 'https://huggingface.co/intfloat/multilingual-e5-small/resolve/main/tokenizer.json',
135 : ),
136 : sha256: _multilingualE5SmallTokenizerJsonSha256,
137 : ),
138 : },
139 3 : meta: {
140 : 'dimensions': 384,
141 : // Selects XlmRobertaTokenizer in OnnxEmbeddingModel.load().
142 : 'tokenizerFamily': 'xlmr',
143 : // multilingual-e5-small requires a mandatory "query: " / "passage: "
144 : // prefix per its model card — without it, retrieval quality degrades
145 : // silently (no error, just worse rankings). OnnxEmbeddingModel.embed()
146 : // applies these based on the caller's EmbeddingKind.
147 : 'queryPrefix': 'query: ',
148 : 'documentPrefix': 'passage: ',
149 : },
150 : );
151 :
152 : /// SHA-256 of `multilingual-e5-small`'s `onnx/model.onnx`, as downloaded
153 : /// and verified by `tool/register_model.dart` — see that script's own doc
154 : /// comment for how to regenerate/re-verify this value.
155 : static const _multilingualE5SmallOnnxSha256 =
156 : 'ca456c06b3a9505ddfd9131408916dd79290368331e7d76bb621f1cba6bc8665';
157 :
158 : /// SHA-256 of `multilingual-e5-small`'s `tokenizer.json`, as downloaded and
159 : /// verified by `tool/register_model.dart`.
160 : static const _multilingualE5SmallTokenizerJsonSha256 =
161 : '0b44a9d7b51c3c62626640cda0e2c2f70fdacdc25bbbd68038369d14ebdf4c39';
162 :
163 : /// Internal fixture — **not a real model.**
164 : ///
165 : /// Exists solely so tests can exercise the catalog's "registered but not
166 : /// validated" gating behaviour (throws [UnsupportedError] from [lookup])
167 : /// against a stable id that will never be flipped to validated. Its file
168 : /// URLs deliberately point at a non-resolvable host so any accidental
169 : /// real-world use (someone actually trying to download or load it) fails
170 : /// fast and obviously rather than silently. See `plan_0_06_wi4_multilingual_embedding_model.md`'s
171 : /// Q5 for why this replaced the previous `bge-m3-v1.0` stub entry (which
172 : /// had placeholder all-zero checksums and could never actually be
173 : /// downloaded or validated).
174 : ///
175 : /// **Must never be set to `true` in [_validated].**
176 9 : static final _placeholderModel = ModelSpec(
177 : id: 'placeholder-model',
178 3 : files: {
179 3 : 'onnx': ModelFile(
180 3 : url: Uri.parse('https://example.invalid/placeholder/model.onnx'),
181 : sha256:
182 : '0000000000000000000000000000000000000000000000000000000000000000',
183 : ),
184 3 : 'vocab': ModelFile(
185 3 : url: Uri.parse('https://example.invalid/placeholder/vocab.txt'),
186 : sha256:
187 : '0000000000000000000000000000000000000000000000000000000000000000',
188 : ),
189 : },
190 3 : meta: {'dimensions': 0},
191 : );
192 :
193 : // ── Internal catalog and validation state ─────────────────────────────────
194 :
195 : /// All registered models keyed by [ModelSpec.id].
196 : ///
197 : /// Uses a lazy getter rather than a const map because the values are
198 : /// `static final` (not const, due to `Uri` not being const in Dart).
199 6 : static Map<String, ModelSpec> get _catalog => {
200 3 : 'bge-small-en-v1.5': _bgeSmallEnV15,
201 3 : 'multilingual-e5-small': _multilingualE5Small,
202 3 : 'placeholder-model': _placeholderModel,
203 : };
204 :
205 : /// Validation state for each registered model.
206 : ///
207 : /// `betto_onnxrt`'s generic [ModelSpec] has no validation concept, so this
208 : /// catalog tracks it separately here. Only models explicitly set to `true`
209 : /// are permitted by [lookup]. A model absent from this map is considered
210 : /// unvalidated.
211 : static const Map<String, bool> _validated = {
212 : 'bge-small-en-v1.5': true,
213 : 'multilingual-e5-small': true,
214 : // Deliberately, permanently false — see _placeholderModel's doc comment.
215 : 'placeholder-model': false,
216 : };
217 :
218 : // ── Public API ─────────────────────────────────────────────────────────────
219 :
220 : /// The ID of the default/recommended production model.
221 : static const String defaultModelId = 'bge-small-en-v1.5';
222 :
223 : /// Returns all registered [ModelSpec]s (validated and unvalidated).
224 : ///
225 : /// Useful for listing available models in a CLI command. Check
226 : /// the model's validation state via [_validated] before presenting it as
227 : /// user-selectable.
228 6 : static Iterable<ModelSpec> get all => _catalog.values;
229 :
230 : /// Looks up the [ModelSpec] for [id].
231 : ///
232 : /// Throws [ArgumentError] if [id] is not registered in the catalog.
233 : /// Throws [UnsupportedError] if the model is registered but not yet
234 : /// validated for production use.
235 : ///
236 : /// ```dart
237 : /// final spec = ModelCatalog.lookup('bge-small-en-v1.5');
238 : /// print(spec.meta['dimensions']); // 384
239 : /// ```
240 2 : static ModelSpec lookup(String id) {
241 2 : final catalog = _catalog;
242 2 : final spec = catalog[id];
243 : if (spec == null) {
244 2 : final known = catalog.keys.join(', ');
245 2 : throw ArgumentError(
246 : "Unknown embedding model ID '$id'. "
247 : "Registered models: $known. "
248 : "Add the model to ModelCatalog to use it.",
249 : );
250 : }
251 2 : if (!(_validated[id] ?? false)) {
252 2 : throw UnsupportedError(
253 : "Embedding model '$id' is registered in the catalog but has not "
254 : "yet been validated for production use. It will be enabled in "
255 : "a future release.",
256 : );
257 : }
258 : return spec;
259 : }
260 :
261 : /// Returns `true` if [id] is a known registered model ID (validated or not).
262 : ///
263 : /// Does **not** check validation status. Useful for detecting legacy config
264 : /// files that reference a known (but maybe unvalidated) model.
265 3 : static bool isKnown(String id) => _catalog.containsKey(id);
266 :
267 : // ── AllowlistProvider ──────────────────────────────────────────────────────
268 :
269 : /// Returns `true` if [spec] is registered in this catalog.
270 : ///
271 : /// Implements [AllowlistProvider] for use with [ModelDownloader]:
272 : ///
273 : /// ```dart
274 : /// final downloader = ModelDownloader(allowlist: ModelCatalog());
275 : /// ```
276 : ///
277 : /// This permits downloading of unvalidated models (e.g.
278 : /// `multilingual-e5-small` prior to its validation pass completing during
279 : /// development). Call [lookup] (which checks validation status) before
280 : /// loading a model for inference.
281 2 : @override
282 6 : bool isAllowed(ModelSpec spec) => _catalog.containsKey(spec.id);
283 : }
|
