From b4492376547b78c0542c71c25ce0300f3055b23c Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Mon, 6 Jul 2026 18:01:15 +0300 Subject: [PATCH 01/20] =?UTF-8?q?feat!:=20v2.0.0=20=E2=80=94=20deepen=20th?= =?UTF-8?q?e=20FFI=20seam,=20extract=20worker=20modules,=20typed=20errors?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BREAKING CHANGE: Model2Vec is now a stateless namespace of static methods. Native library seam (Candidates 1 & 4) - Switch ffigen to @Native code-assets mode (assetId package:model2vec/model2vec.so); the SDK resolver — fed by the build hook — now locates the native library. Delete _resolveLibPath, boot(), the singleton instance, both DynamicLibrary constructors, and the _cachedDimension cache. Fixes Flutter bundling and the boot()-across-isolates bug. Native ABI and error seam (Candidate 3) - Rust now owns output allocation: generate_* return ptr+dim(+count), freed via free_floats — removing the dim/model-switch buffer-overflow race by construction. Every fallible function returns a stable code plus a char** out_error, and every #[no_mangle] is wrapped in catch_unwind so a panic (including from a malformed model) becomes a typed error instead of UB. Length params use size_t (fixes a 64-bit-Windows usize mismatch). Dart surfaces a single Model2VecException with an exhaustive Model2VecErrorKind. Streaming worker (Candidate 2) - Rebuild generateEmbeddingStream on small, tested modules: batched() (pure transformer), Channel (isolate-port seam + in-memory fake), WorkerProtocol (serialized-queue state machine) and a private EmbeddingWorker. Worker errors cross the isolate boundary as typed exceptions; a dying worker fails pending requests via onExit instead of hanging. Catalog (Candidate 5) - getRecommendedModels() -> typed Model2Vec.recommendedModels constant. Tests - 60 tests. The model-backed suite is migrated to the static API; new no-model/no-network suites cover batching, the worker protocol (fake channel), the worker over a real isolate (incl. death/hang robustness) and native error paths. dart_test.yaml pins concurrency: 1 because the native model is a single process-global. Full migration guide is in CHANGELOG.md. --- .gitignore | 2 + CHANGELOG.md | 44 ++ README.md | 39 +- benchmark/embedding_benchmark.dart | 23 +- dart_test.yaml | 4 + example/main.dart | 52 ++- lib/model2vec.dart | 3 +- lib/src/batcher.dart | 25 ++ lib/src/channel.dart | 147 +++++++ lib/src/embedding_worker.dart | 73 ++++ lib/src/exception.dart | 86 +++- lib/src/model2vec_base.dart | 660 +++++++++++------------------ lib/src/model2vec_bindings.g.dart | 419 ++++++------------ lib/src/recommended_model.dart | 45 ++ lib/src/worker_protocol.dart | 123 ++++++ native/Cargo.lock | 1 - native/Cargo.toml | 1 - native/model2vec.h | 70 ++- native/src/lib.rs | 455 ++++++++++++-------- native/src/model.rs | 15 +- pubspec.yaml | 6 +- test/batcher_test.dart | 53 +++ test/embedding_worker_test.dart | 123 ++++++ test/error_path_test.dart | 65 +++ test/model2vec_test.dart | 91 ++-- test/utils_test.dart | 20 +- test/worker_protocol_test.dart | 138 ++++++ 27 files changed, 1750 insertions(+), 1033 deletions(-) create mode 100644 dart_test.yaml create mode 100644 lib/src/batcher.dart create mode 100644 lib/src/channel.dart create mode 100644 lib/src/embedding_worker.dart create mode 100644 lib/src/recommended_model.dart create mode 100644 lib/src/worker_protocol.dart create mode 100644 test/batcher_test.dart create mode 100644 test/embedding_worker_test.dart create mode 100644 test/error_path_test.dart create mode 100644 test/worker_protocol_test.dart diff --git a/.gitignore b/.gitignore index d8f2fb8..9cc6b98 100644 --- a/.gitignore +++ b/.gitignore @@ -29,3 +29,5 @@ doc/api/ .flutter-plugins-dependencies native/target +docs +CONTEXT.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ab0531..a362b5d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,48 @@ +# 2.0.0 + +Major release reworking the FFI boundary and public surface for testability and +correctness. **This release is breaking** — see migration below. + +**Breaking changes:** + +- **Static API.** `Model2Vec` is now a stateless namespace of static methods. + `Model2Vec.instance`, the `Model2Vec(DynamicLibrary)` constructor and + `Model2Vec.boot(...)` were removed — the native library is resolved + automatically through Native Assets (`@Native` code assets). Replace + `Model2Vec.instance.foo(...)` with `Model2Vec.foo(...)`. +- **Recommended models.** `getRecommendedModels()` (returning + `List>`) is replaced by the typed constant + `Model2Vec.recommendedModels` (`List`). +- **Typed errors.** `Model2VecException` now carries a `Model2VecErrorKind kind`; + its constructor is `(kind, message, [code])` and the `fromCode` factory is + replaced by `fromNative(code, message)`. Native failures surface the message + produced by the Rust layer, each with an exhaustively-switchable `kind`. + +**Improvements:** + +- **Native memory safety.** The `generate_*` FFI functions now allocate their + output inside the native call (returned as a pointer the caller frees), + removing a dimension/model-switch race that could overflow the output buffer. + Every native entry point is wrapped in `catch_unwind`, so a panic (including + from a malformed model) surfaces as a typed error instead of undefined + behaviour. +- **Windows ABI fix.** FFI length parameters use `size_t` (was `unsigned long`, + 32-bit on 64-bit Windows and mismatched against Rust's `usize`). +- **Streaming rework.** `generateEmbeddingStream` is rebuilt on small, tested + modules — a batching transformer, a transport-agnostic worker protocol, and a + worker isolate. Worker errors cross the isolate boundary as typed + `Model2VecException`s (kind + code preserved) rather than stringified errors. + +**Migration:** + +| 1.x | 2.0.0 | +| --- | --- | +| `Model2Vec.instance.generateEmbedding(t)` | `Model2Vec.generateEmbedding(t)` | +| `Model2Vec.boot(lib)` / `Model2Vec(lib)` | removed — resolution is automatic | +| `Model2Vec.instance.getRecommendedModels()` | `Model2Vec.recommendedModels` (typed) | +| `catch (e) { e.code }` | still works; add `e.kind` for exhaustive handling | + # 1.2.0 - Lowered minimum Dart SDK requirement to `3.10.0` to support a wider range of environments. diff --git a/README.md b/README.md index b6b24e0..ff5ba42 100644 --- a/README.md +++ b/README.md @@ -71,19 +71,26 @@ dart pub add model2vec import 'package:model2vec/model2vec.dart'; void main() { - final m2v = Model2Vec.instance; - + // Model2Vec is a stateless namespace of static methods — there is one + // active model per process, loaded automatically via Native Assets. + // Initialize with a model from Hugging Face - m2v.initEmbedder('minishlab/potion-base-2M'); - + Model2Vec.initEmbedder('minishlab/potion-base-2M'); + // Generate an embedding - final embedding = m2v.generateEmbedding('Dart FFI is blazingly fast 🚀'); - - print('Vector dimension: ${m2v.embeddingDimension}'); - print('Vocabulary size: ${m2v.vocabularySize}'); + final embedding = Model2Vec.generateEmbedding('Dart FFI is blazingly fast 🚀'); + + print('Vector dimension: ${Model2Vec.embeddingDimension}'); + print('Vocabulary size: ${Model2Vec.vocabularySize}'); } ``` +> **Migrating from 1.x?** The instance API is gone. Replace +> `Model2Vec.instance.foo(...)` with `Model2Vec.foo(...)`, drop any +> `Model2Vec.boot(...)` / `Model2Vec(lib)` calls (library resolution is now +> automatic), and read `Model2Vec.recommendedModels` instead of calling +> `getRecommendedModels()`. See the [CHANGELOG](CHANGELOG.md) for the full list. + ## Recipes & Patterns ### 1. Advanced Batch Processing @@ -93,7 +100,7 @@ Process multiple strings at once for maximum hardware utilization. You can contr ```dart final texts = ['Dart', 'Rust', 'Flutter']; -final embeddings = m2v.generateBatchEmbeddings( +final embeddings = Model2Vec.generateBatchEmbeddings( texts, maxLength: 256, // Truncate strings longer than 256 tokens batchSize: 1024, // Internal chunks sent to the FFI layer @@ -115,7 +122,7 @@ Future processHugeFile() async { .transform(const LineSplitter()); // Converts a Stream into a Stream - final embeddingStream = m2v.generateEmbeddingStream( + final embeddingStream = Model2Vec.generateEmbeddingStream( fileStream, batchSize: 500, // Process 500 strings at a time useIsolate: true, // Run math in background threads @@ -132,8 +139,8 @@ Future processHugeFile() async { Never block the main thread. If you are building a Flutter app, always use the `Async` variants to perform generation in a background `Isolate`. ```dart -final embedding = await m2v.generateEmbeddingAsync('A very long text...'); -final batch = await m2v.generateBatchEmbeddingsAsync(['A', 'B', 'C']); +final embedding = await Model2Vec.generateEmbeddingAsync('A very long text...'); +final batch = await Model2Vec.generateBatchEmbeddingsAsync(['A', 'B', 'C']); ``` ### 4. Vector Math & Quantization @@ -141,10 +148,10 @@ final batch = await m2v.generateBatchEmbeddingsAsync(['A', 'B', 'C']); The library ships with `Model2VecUtils` — a powerful suite of math operations tuned for embeddings. ```dart -final query = m2v.generateEmbedding('cat'); +final query = Model2Vec.generateEmbedding('cat'); final candidates = [ - m2v.generateEmbedding('dog'), - m2v.generateEmbedding('space'), + Model2Vec.generateEmbedding('dog'), + Model2Vec.generateEmbedding('space'), ]; // 1. Semantic Similarity (Cosine) @@ -174,7 +181,7 @@ final base64String = Model2VecUtils.toBase64(query); | `initEmbedder(path)` | Initializes the model from a Hugging Face repo ID or local path. | | `initEmbedderAdvanced(...)` | Advanced initialization with custom `cacheDirectory`, `hfToken`, or `normalize` overrides. | | `initEmbedderFromBytes(...)` | Initializes the model directly from raw `Uint8List` bytes (`model.safetensors`, `tokenizer.json`, etc). | -| `getRecommendedModels()` | Returns a list of officially supported models. | +| `recommendedModels` | A typed `List` catalog of officially recommended Potion models (offline). | | `tokenize(text)` | Runs the internal BPE tokenizer and returns a `List`. | | `generateEmbedding(text)` | Synchronously generates a `Float32List` embedding vector. | | `generateBatchEmbeddings(texts)` | Synchronously generates embeddings for a `List` using Rust SIMD. | diff --git a/benchmark/embedding_benchmark.dart b/benchmark/embedding_benchmark.dart index 413446b..5dedfef 100644 --- a/benchmark/embedding_benchmark.dart +++ b/benchmark/embedding_benchmark.dart @@ -12,32 +12,18 @@ late List batchTexts; class EmbeddingBenchmark extends BenchmarkBase { EmbeddingBenchmark() : super('Single'); - late Model2Vec m2v; - - @override - void setup() { - m2v = Model2Vec.instance; - } - @override void run() { - m2v.generateEmbedding(testText); + Model2Vec.generateEmbedding(testText); } } class BatchEmbeddingBenchmark extends BenchmarkBase { BatchEmbeddingBenchmark() : super('Batch (32)'); - late Model2Vec m2v; - - @override - void setup() { - m2v = Model2Vec.instance; - } - @override void run() { - m2v.generateBatchEmbeddings(batchTexts); + Model2Vec.generateBatchEmbeddings(batchTexts); } } @@ -58,11 +44,10 @@ void main() { 'Warming up models (Downloading to cache, this may take a while)...', ); - final m2v = Model2Vec.instance; for (final model in models) { stdout.write('Warming up $model... '); final watch = Stopwatch()..start(); - m2v.initEmbedder(model); + Model2Vec.initEmbedder(model); watch.stop(); stdout.writeln('Done in ${watch.elapsedMilliseconds} ms.'); } @@ -77,7 +62,7 @@ void main() { for (final model in models) { // Measure Load Time (Already cached) final watch = Stopwatch()..start(); - m2v.initEmbedder(model); + Model2Vec.initEmbedder(model); watch.stop(); final loadTime = '${watch.elapsedMilliseconds} ms'; diff --git a/dart_test.yaml b/dart_test.yaml new file mode 100644 index 0000000..7bc9d1b --- /dev/null +++ b/dart_test.yaml @@ -0,0 +1,4 @@ +# The native layer keeps a single active model per process. The model-switching +# test mutates that global, so suites must run one at a time — otherwise a +# parallel suite can read a model of the wrong dimension mid-switch. +concurrency: 1 diff --git a/example/main.dart b/example/main.dart index 75644cd..b6d2d12 100644 --- a/example/main.dart +++ b/example/main.dart @@ -5,47 +5,43 @@ import 'package:model2vec/model2vec.dart'; /// Example demonstrating production-ready usage of the Model2Vec package. Future main() async { try { - // 1. Initialize the API via shared instance - final m2v = Model2Vec.instance; - stdout.writeln('✅ Model2Vec library loaded successfully.'); + stdout.writeln('=== Model2Vec Example ==='); - // 3. Explore Recommended Models - final models = m2v.getRecommendedModels(); + // 1. Explore Recommended Models + const models = Model2Vec.recommendedModels; stdout.writeln('\n📦 Available models:'); for (final m in models) { - final name = m['name']! as String; - - stdout.writeln(' - ${name.padRight(25)} ID: ${m['id']}'); + stdout.writeln(' - ${m.name.padRight(25)} ID: ${m.id}'); } - // 4. Initialize a Model + // 2. Initialize a Model const modelId = 'minishlab/potion-base-2M'; stdout.writeln('\n🚀 Initializing $modelId...'); final sw = Stopwatch()..start(); - m2v.initEmbedder(modelId); + Model2Vec.initEmbedder(modelId); stdout ..writeln('✨ Initialized in ${sw.elapsedMilliseconds}ms') - // 5. Inspect Model Metadata + // 3. Inspect Model Metadata ..writeln('\n📊 Model Metadata:') - ..writeln(' - Dimension: ${m2v.embeddingDimension}') - ..writeln(' - Vocabulary Size: ${m2v.vocabularySize}') - ..writeln(' - Is Normalized: ${m2v.isNormalized}') - ..writeln(' - Median Token Length: ${m2v.medianTokenLength}'); + ..writeln(' - Dimension: ${Model2Vec.embeddingDimension}') + ..writeln(' - Vocabulary Size: ${Model2Vec.vocabularySize}') + ..writeln(' - Is Normalized: ${Model2Vec.isNormalized}') + ..writeln(' - Median Token Length: ${Model2Vec.medianTokenLength}'); - // 6. Tokenization Demo + // 4. Tokenization Demo const text = 'Model2Vec is incredibly fast!'; - final tokens = m2v.tokenize(text); + final tokens = Model2Vec.tokenize(text); stdout ..writeln('\n🔍 Tokenization: "$text"') ..writeln(' - Tokens: $tokens') - // 7. Single Embedding + // 5. Single Embedding ..writeln('\n🧠 Generating single embedding...'); - final embedding = m2v.generateEmbedding(text); + final embedding = Model2Vec.generateEmbedding(text); stdout ..writeln(' - Vector (first 3): ${embedding.take(3).toList()}') ..writeln(' - Total Length: ${embedding.length}'); - // 8. Batch Embedding (Production Optimization) + // 6. Batch Embedding (Production Optimization) final texts = [ 'The first sentence.', 'A second, slightly longer sentence for the batch.', @@ -55,7 +51,7 @@ Future main() async { '\n⚡ Generating batch embeddings for ${texts.length} sentences...', ); final batchStartTime = DateTime.now(); - final batch = m2v.generateBatchEmbeddings(texts); + final batch = Model2Vec.generateBatchEmbeddings(texts); final batchDuration = DateTime.now().difference(batchStartTime); stdout.writeln(' - Processed in ${batchDuration.inMicroseconds}μs'); @@ -63,13 +59,13 @@ Future main() async { stdout.writeln(' - Result $i length: ${batch[i].length}'); } - // 9. Vector Math & Semantic Search + // 7. Vector Math & Semantic Search stdout.writeln('\n🧠 Vector Math & Semantic Search:'); - final query = m2v.generateEmbedding('A cute little kitten'); + final query = Model2Vec.generateEmbedding('A cute little kitten'); final db = [ - m2v.generateEmbedding('A small cat'), - m2v.generateEmbedding('A big dog'), - m2v.generateEmbedding('Space exploration'), + Model2Vec.generateEmbedding('A small cat'), + Model2Vec.generateEmbedding('A big dog'), + Model2Vec.generateEmbedding('Space exploration'), ]; final simCat = Model2VecUtils.cosineSimilarity(query, db[0]); @@ -85,10 +81,10 @@ Future main() async { final topMatch = Model2VecUtils.similaritySearch(query, db, topK: 1); stdout ..writeln(' - Best match index: ${topMatch.first}') - // 10. Streaming API for Huge Datasets + // 8. Streaming API for Huge Datasets ..writeln('\n🌊 Streaming API (1000 items):'); final stream = Stream.fromIterable(List.generate(1000, (i) => 'Item $i')); - final resultStream = m2v.generateEmbeddingStream( + final resultStream = Model2Vec.generateEmbeddingStream( stream, batchSize: 200, ); diff --git a/lib/model2vec.dart b/lib/model2vec.dart index 183b49d..6659d6e 100644 --- a/lib/model2vec.dart +++ b/lib/model2vec.dart @@ -2,6 +2,7 @@ /// library for generating text embeddings. library; -export 'src/exception.dart' show Model2VecException; +export 'src/exception.dart' show Model2VecErrorKind, Model2VecException; export 'src/model2vec_base.dart' show Model2Vec; +export 'src/recommended_model.dart' show RecommendedModel; export 'src/utils.dart' show Model2VecUtils; diff --git a/lib/src/batcher.dart b/lib/src/batcher.dart new file mode 100644 index 0000000..2c10a24 --- /dev/null +++ b/lib/src/batcher.dart @@ -0,0 +1,25 @@ +/// Groups a [source] stream into fixed-size batches. +/// +/// Each batch holds [size] items and is emitted as soon as it fills; the final +/// batch may be smaller. This is a pure stream transformer with no isolate or +/// model dependency, so it is trivially testable in isolation. +/// +/// Throws [ArgumentError] if [size] is less than 1. +Stream> batched(Stream source, int size) async* { + if (size < 1) { + throw ArgumentError.value(size, 'size', 'must be >= 1'); + } + + var buffer = []; + await for (final item in source) { + buffer.add(item); + if (buffer.length >= size) { + yield buffer; + buffer = []; + } + } + + if (buffer.isNotEmpty) { + yield buffer; + } +} diff --git a/lib/src/channel.dart b/lib/src/channel.dart new file mode 100644 index 0000000..b3dac33 --- /dev/null +++ b/lib/src/channel.dart @@ -0,0 +1,147 @@ +import 'dart:async'; +import 'dart:isolate'; + +/// A bidirectional message channel — the transport a worker protocol runs on. +/// +/// This is the seam that makes the worker testable: production uses +/// [IsolateChannel] over real isolate ports, while tests drive an in-memory +/// fake. The protocol on top of it never mentions isolates or ports. +/// +/// [incoming] closes when the other end goes away — for [IsolateChannel] that +/// includes the worker isolate dying unexpectedly, so a protocol listening for +/// `onDone` can fail its pending requests instead of hanging forever. +abstract interface class Channel { + /// Sends [message] to the other end. + void send(Object? message); + + /// Messages arriving from the other end (excluding any handshake). + Stream get incoming; + + /// Tears the channel down, releasing the other end. + Future close(); +} + +/// A [Channel] backed by a dedicated worker [Isolate]. +/// +/// The worker's `entryPoint` receives a [SendPort] and must reply with its own +/// [SendPort] as the first message (the handshake); every later message is +/// surfaced on [incoming]. If the isolate exits — cleanly or by crashing — +/// [incoming] is closed so callers are never left waiting on a dead worker. +class IsolateChannel implements Channel { + IsolateChannel._( + this._isolate, + this._workerSendPort, + this._receivePort, + this._exitPort, + this._subscription, + this._exitSubscription, + this._incoming, + ); + + final Isolate _isolate; + final SendPort _workerSendPort; + final ReceivePort _receivePort; + final ReceivePort _exitPort; + final StreamSubscription _subscription; + final StreamSubscription _exitSubscription; + final StreamController _incoming; + var _closed = false; + + /// Spawns a worker running [entryPoint] and completes the handshake. + static Future start( + void Function(SendPort) entryPoint, + ) async { + final receivePort = ReceivePort(); + final exitPort = ReceivePort(); + final isolate = await Isolate.spawn( + entryPoint, + receivePort.sendPort, + onExit: exitPort.sendPort, + ); + + final incoming = StreamController(); + final handshake = Completer(); + + final subscription = receivePort.listen((message) { + if (!handshake.isCompleted) { + if (message is SendPort) { + handshake.complete(message); + } else { + handshake.completeError( + StateError('worker isolate did not hand back a SendPort'), + ); + } + return; + } + // The reply port and the exit port are independent; if the exit + // notification already closed [incoming], drop this late message rather + // than adding to a closed controller. + if (!incoming.isClosed) { + incoming.add(message); + } + }); + + // The isolate exited. Fail a pending handshake, and close [incoming] so a + // protocol awaiting a reply is released rather than hanging forever. + final exitSubscription = exitPort.listen((_) { + if (!handshake.isCompleted) { + handshake.completeError( + StateError('worker isolate exited during startup'), + ); + } + if (!incoming.isClosed) { + unawaited(incoming.close()); + } + }); + + final SendPort workerSendPort; + try { + workerSendPort = await handshake.future; + } on Object { + await subscription.cancel(); + await exitSubscription.cancel(); + receivePort.close(); + exitPort.close(); + if (!incoming.isClosed) { + await incoming.close(); + } + isolate.kill(priority: Isolate.immediate); + rethrow; + } + + return IsolateChannel._( + isolate, + workerSendPort, + receivePort, + exitPort, + subscription, + exitSubscription, + incoming, + ); + } + + @override + void send(Object? message) => _workerSendPort.send(message); + + @override + Stream get incoming => _incoming.stream; + + @override + Future close() async { + if (_closed) { + return; + } + _closed = true; + // Cancel the exit listener first so our own kill() below doesn't re-enter + // the unexpected-exit path. + await _exitSubscription.cancel(); + _workerSendPort.send('close'); + await _subscription.cancel(); + _isolate.kill(); + _receivePort.close(); + _exitPort.close(); + if (!_incoming.isClosed) { + await _incoming.close(); + } + } +} diff --git a/lib/src/embedding_worker.dart b/lib/src/embedding_worker.dart new file mode 100644 index 0000000..0bbfc6d --- /dev/null +++ b/lib/src/embedding_worker.dart @@ -0,0 +1,73 @@ +import 'dart:async'; +import 'dart:isolate'; +import 'dart:typed_data'; + +import 'channel.dart'; +import 'exception.dart'; +import 'model2vec_base.dart'; +import 'worker_protocol.dart'; + +/// A background worker that produces embeddings off the caller's thread, one +/// batch at a time. +/// +/// Private to the package: callers reach it through +/// [Model2Vec.generateEmbeddingStream]. It wires an [IsolateChannel] to a +/// [WorkerProtocol]; the isolate lifecycle and the request protocol are the +/// only concerns here. +class EmbeddingWorker { + EmbeddingWorker._(this._protocol); + + final WorkerProtocol _protocol; + + /// Spawns the worker isolate and completes its handshake. + /// + /// [entryPoint] defaults to the model-backed [embeddingWorkerEntryPoint]; + /// tests inject a fake top-level entry point to exercise the protocol + /// without a model. + static Future start({ + void Function(SendPort) entryPoint = embeddingWorkerEntryPoint, + }) async { + final channel = await IsolateChannel.start(entryPoint); + return EmbeddingWorker._(WorkerProtocol(channel)); + } + + /// Embeds [batch] on the worker isolate. + Future> embedBatch( + List batch, { + int maxLength = 512, + }) => _protocol.embedBatch(batch, maxLength: maxLength); + + /// Tears the worker down, failing any outstanding requests. + Future close() => _protocol.close(); +} + +/// Default worker isolate entry point. +/// +/// Embeds each requested batch against the process-global model and replies +/// with the vectors, or with a typed [Model2VecException] on failure so the +/// error's [Model2VecErrorKind] and code survive the isolate boundary. +void embeddingWorkerEntryPoint(SendPort mainSendPort) { + final receivePort = ReceivePort(); + mainSendPort.send(receivePort.sendPort); + + receivePort.listen((message) { + if (message == 'close') { + receivePort.close(); + return; + } + + final request = message as EmbedRequest; + try { + final results = Model2Vec.generateBatchEmbeddings( + request.batch, + maxLength: request.maxLength, + batchSize: request.batch.length, + ); + mainSendPort.send(results); + } on Model2VecException catch (e) { + mainSendPort.send(e); + } on Object catch (e) { + mainSendPort.send(Model2VecException(Model2VecErrorKind.unknown, '$e')); + } + }); +} diff --git a/lib/src/exception.dart b/lib/src/exception.dart index 4b36d32..a4bacad 100644 --- a/lib/src/exception.dart +++ b/lib/src/exception.dart @@ -1,30 +1,72 @@ +/// Classification of a failure that crossed the native (FFI) boundary. +/// +/// Values mirror the stable error codes the Rust layer returns, so callers can +/// branch exhaustively with a `switch` instead of matching magic numbers. +enum Model2VecErrorKind { + /// No model has been initialized yet (call `initEmbedder` first). + notInitialized, + + /// A model could not be loaded from the given repo id or path. + modelLoadFailed, + + /// A model could not be initialized from the provided bytes. + initFromBytesFailed, + + /// The native model lock was poisoned by a previous panic. + lockPoisoned, + + /// A required argument was null or otherwise invalid. + nullArgument, + + /// Tokenization of the input failed. + tokenizationFailed, + + /// The native layer produced no result where one was expected. + emptyResult, + + /// A panic was caught at the native boundary. + panic, + + /// An error code the current Dart version does not recognize. + unknown, +} + /// Exception thrown when a Model2Vec operation fails. class Model2VecException implements Exception { - const Model2VecException(this.message, [this.code]); - - /// Creates a [Model2VecException] from a native error code. - factory Model2VecException.fromCode(int code, String context) { - final msg = switch (code) { - -1 => 'Model not initialized or invalid path.', - -2 => 'Failed to load model from the provided path.', - -3 => 'Failed to initialize model from memory bytes.', - -4 => 'Native library internal error (Lock poisoned).', - _ => 'Unknown native error.', - }; - return Model2VecException('$context: $msg', code); - } - - /// Error message describing the failure. + /// Creates an exception with an explicit [kind] and [message]. + const Model2VecException(this.kind, this.message, [this.code]); + + /// Builds an exception from a native status [code] and the native [message]. + factory Model2VecException.fromNative(int code, String message) => + Model2VecException(_kindFromCode(code), message, code); + + /// The category of failure, suitable for exhaustive handling. + final Model2VecErrorKind kind; + + /// Human-readable description. When the failure originated in the native + /// layer, this is the message that layer produced. final String message; - /// Optional error code returned by the native library. + /// The raw native status code, when the failure came from the native layer. final int? code; @override - String toString() { - if (code != null) { - return 'Model2VecException (Code $code): $message'; - } - return 'Model2VecException: $message'; - } + String toString() => code != null + ? 'Model2VecException(${kind.name}, code $code): $message' + : 'Model2VecException(${kind.name}): $message'; } + +/// Maps a stable native status code to its [Model2VecErrorKind]. +/// +/// Keep in sync with the `CODE_*` constants in `native/src/lib.rs`. +Model2VecErrorKind _kindFromCode(int code) => switch (code) { + 1 => Model2VecErrorKind.notInitialized, + 2 => Model2VecErrorKind.modelLoadFailed, + 3 => Model2VecErrorKind.initFromBytesFailed, + 4 => Model2VecErrorKind.lockPoisoned, + 5 => Model2VecErrorKind.nullArgument, + 6 => Model2VecErrorKind.tokenizationFailed, + 7 => Model2VecErrorKind.emptyResult, + 8 => Model2VecErrorKind.panic, + _ => Model2VecErrorKind.unknown, +}; diff --git a/lib/src/model2vec_base.dart b/lib/src/model2vec_base.dart index 78408a7..0b82667 100644 --- a/lib/src/model2vec_base.dart +++ b/lib/src/model2vec_base.dart @@ -1,277 +1,210 @@ -import 'dart:async'; import 'dart:convert'; import 'dart:ffi'; -import 'dart:io'; import 'dart:isolate'; import 'dart:typed_data'; import 'package:ffi/ffi.dart'; -import 'package:path/path.dart' as p; +import 'batcher.dart'; +import 'embedding_worker.dart'; import 'exception.dart'; -import 'model2vec_bindings.g.dart'; +import 'model2vec_bindings.g.dart' as native; +import 'recommended_model.dart'; /// The main entry point for the Model2Vec library. /// -/// This class provides methods to initialize the embedder, generate text -/// embeddings, and access model metadata like vocabulary size and dimensions. -class Model2Vec { - /// Creates a new instance of [Model2Vec] using the provided [library]. - /// - /// In most cases, you should use the shared [instance] instead of - /// creating a new one manually. - Model2Vec(DynamicLibrary library) : _bindings = Model2VecBindings(library); - - static Model2Vec? _instance; - - final Model2VecBindings _bindings; - int? _cachedDimension; - - /// Manually initializes the shared [instance] with a specific [library]. - /// - /// This is useful if you need to load the native library from a custom - /// location or if automatic resolution fails. - static void boot(DynamicLibrary library) { - _instance = Model2Vec(library); - } - - /// A shared singleton instance of [Model2Vec]. - /// - /// When accessed for the first time, it attempts to automatically resolve - /// and load the native library using [Platform.isLinux], [Platform.isMacOS], - /// and [Platform.isWindows] to determine the correct library filename. - /// - /// Throws a [Model2VecException] if the library cannot be found. - static Model2Vec get instance { - if (_instance != null) { - return _instance!; - } - - final libPath = _resolveLibPath(); - try { - final library = DynamicLibrary.open(libPath); - _instance = Model2Vec(library); - - return _instance!; - } catch (e) { - throw Model2VecException( - 'Failed to auto-load Model2Vec native library "$libPath".\n' - 'Please ensure you have built the Rust library ' - '(cd native && cargo build --release)\n' - 'or that the shared object is placed in a system library path.\n' - 'Underlying error: $e', - ); - } - } - +/// The native layer keeps a single active model per process, so this type is a +/// stateless namespace of static operations rather than something you +/// instantiate. The native library is located automatically through the Dart +/// SDK's code-asset resolver — there is nothing to boot or inject. +// ignore: avoid_classes_with_only_static_members +abstract final class Model2Vec { /// Returns the embedding dimension of the currently loaded model. /// /// Throws a [Model2VecException] if no model has been initialized yet. - int get embeddingDimension { - if (_cachedDimension != null) { - return _cachedDimension!; - } - - final dim = _bindings.get_embedding_dimension(); - if (dim <= 0) { - throw const Model2VecException( - 'No model initialized. Call initEmbedder() before accessing dimension.', - ); - } - - _cachedDimension = dim; - - return dim; - } + static int get embeddingDimension => + _readInt(native.get_embedding_dimension); /// Returns the total number of unique tokens in the model's vocabulary. /// /// Throws a [Model2VecException] if no model has been initialized yet. - int get vocabularySize { - final size = _bindings.get_vocabulary_size(); - if (size < 0) { - throw const Model2VecException( - 'Failed to get vocabulary size. Is a model initialized?', - ); - } + static int get vocabularySize => _readInt(native.get_vocabulary_size); - return size; - } - - /// Returns `true` if the model is configured to L2-normalize output - /// embeddings. - bool get isNormalized => _bindings.is_normalized() == 1; + /// Returns `true` if the model L2-normalizes its output embeddings. + /// + /// Throws a [Model2VecException] if no model has been initialized yet. + static bool get isNormalized => _readInt(native.is_normalized) == 1; /// Returns the median length (in characters) of tokens in the vocabulary. /// /// Throws a [Model2VecException] if no model has been initialized yet. - int get medianTokenLength { - final length = _bindings.get_median_token_length(); - - if (length < 0) { - throw const Model2VecException('Failed to get median token length.'); - } - - return length; - } + static int get medianTokenLength => + _readInt(native.get_median_token_length); + + static int _readInt( + int Function(Pointer, Pointer>) fn, + ) => using((arena) { + final outValue = arena(); + final outError = arena>(); + _check(fn(outValue, outError), outError); + return outValue.value; + }); /// Tokenizes the input [text] into a list of strings. /// - /// Throws a [Model2VecException] if tokenization fails or if no model is - /// initialized. - List tokenize(String text) => using((arena) { - final textPtr = text.toNativeUtf8(allocator: arena); - final resPtr = _bindings.tokenize(textPtr.cast()); - - if (resPtr == nullptr) { - throw const Model2VecException('Tokenization failed.'); - } + /// Throws a [Model2VecException] if tokenization fails or no model is loaded. + static List tokenize(String text) => using((arena) { + final textPtr = text.toNativeUtf8(allocator: arena).cast(); + final outJson = arena>(); + final outError = arena>(); + _check(native.tokenize(textPtr, outJson, outError), outError); + + final jsonPtr = outJson.value; try { - final jsonString = resPtr.cast().toDartString(); - return .from(json.decode(jsonString) as List); + final jsonString = jsonPtr.cast().toDartString(); + return List.from(json.decode(jsonString) as List); } finally { - _bindings.free_string(resPtr); + native.free_string(jsonPtr); } }); - /// Initializes a model from a Hugging Face repo ID or a local directory path. + /// Initializes a model from a Hugging Face repo id or a local directory path. /// - /// [modelPath] can be either a repo ID like 'minishlab/potion-base-8M' - /// or a path to a directory containing `model.safetensors`, `config.json`, - /// and `tokenizer.json`. - void initEmbedder(String modelPath) => + /// [modelPath] can be a repo id like `minishlab/potion-base-8M` or a path to + /// a directory containing `model.safetensors`, `config.json` and + /// `tokenizer.json`. + static void initEmbedder(String modelPath) => initEmbedderAdvanced(modelPath: modelPath); /// Advanced model initialization with additional options. /// - /// - [modelPath]: Repo ID or local path. + /// - [modelPath]: Repo id or local path. /// - [hfToken]: Optional Hugging Face API token for private repos. /// - [cacheDirectory]: Optional path to store downloaded models. /// - [normalize]: Whether to L2-normalize output embeddings. /// - [subfolder]: Optional subfolder within the repo/path. - void initEmbedderAdvanced({ + static void initEmbedderAdvanced({ required String modelPath, String? hfToken, String? cacheDirectory, bool? normalize, String? subfolder, - }) { - using((arena) { - final pathPtr = modelPath.toNativeUtf8(allocator: arena); - final tokenPtr = hfToken?.toNativeUtf8(allocator: arena) ?? nullptr; - final cachePtr = - cacheDirectory?.toNativeUtf8(allocator: arena) ?? nullptr; - final subPtr = subfolder?.toNativeUtf8(allocator: arena) ?? nullptr; - final normInt = normalize == null ? -1 : (normalize ? 1 : 0); - - final res = _bindings.init_embedder_advanced( - pathPtr.cast(), - tokenPtr.cast(), - cachePtr.cast(), + }) => using((arena) { + final pathPtr = modelPath.toNativeUtf8(allocator: arena).cast(); + final tokenPtr = hfToken == null + ? nullptr + : hfToken.toNativeUtf8(allocator: arena).cast(); + final cachePtr = cacheDirectory == null + ? nullptr + : cacheDirectory.toNativeUtf8(allocator: arena).cast(); + final subPtr = subfolder == null + ? nullptr + : subfolder.toNativeUtf8(allocator: arena).cast(); + final normInt = normalize == null ? -1 : (normalize ? 1 : 0); + final outError = arena>(); + + _check( + native.init_embedder_advanced( + pathPtr, + tokenPtr, + cachePtr, normInt, - subPtr.cast(), - ); - - if (res != 0) { - throw Model2VecException.fromCode(res, 'Initialization failed'); - } - - _cachedDimension = null; - }); - } + subPtr, + outError, + ), + outError, + ); + }); /// Initializes a model using raw bytes from memory. /// - /// Requires the content of the three main configuration files: /// - [tokenizerBytes]: Content of `tokenizer.json`. /// - [modelBytes]: Content of `model.safetensors`. /// - [configBytes]: Content of `config.json`. - void initEmbedderFromBytes({ + static void initEmbedderFromBytes({ required Uint8List tokenizerBytes, required Uint8List modelBytes, required Uint8List configBytes, - }) { - using((arena) { - final tPtr = arena(tokenizerBytes.length); - final mPtr = arena(modelBytes.length); - final cPtr = arena(configBytes.length); - - tPtr.asTypedList(tokenizerBytes.length).setAll(0, tokenizerBytes); - mPtr.asTypedList(modelBytes.length).setAll(0, modelBytes); - cPtr.asTypedList(configBytes.length).setAll(0, configBytes); - - final res = _bindings.init_embedder_from_bytes( + }) => using((arena) { + final tPtr = arena(tokenizerBytes.length); + final mPtr = arena(modelBytes.length); + final cPtr = arena(configBytes.length); + + tPtr.asTypedList(tokenizerBytes.length).setAll(0, tokenizerBytes); + mPtr.asTypedList(modelBytes.length).setAll(0, modelBytes); + cPtr.asTypedList(configBytes.length).setAll(0, configBytes); + + final outError = arena>(); + _check( + native.init_embedder_from_bytes( tPtr.cast(), tokenizerBytes.length, mPtr.cast(), modelBytes.length, cPtr.cast(), configBytes.length, - ); - - if (res != 0) { - throw Model2VecException.fromCode( - res, - 'Initialization from bytes failed', - ); - } - - _cachedDimension = null; - }); - } + outError, + ), + outError, + ); + }); - /// Returns a list of officially recommended Potion models from Hugging Face. - List> getRecommendedModels() => [ - { - 'id': 'minishlab/potion-base-2M', - 'name': 'Potion Base 2M', - 'lang': 'English', - 'params': '1.8M', - 'description': 'Smallest English model, very fast.', - }, - { - 'id': 'minishlab/potion-base-4M', - 'name': 'Potion Base 4M', - 'lang': 'English', - 'params': '3.7M', - 'description': 'Small and efficient English model.', - }, - { - 'id': 'minishlab/potion-base-8M', - 'name': 'Potion Base 8M', - 'lang': 'English', - 'params': '7.5M', - 'description': 'Balanced English model.', - }, - { - 'id': 'minishlab/potion-base-32M', - 'name': 'Potion Base 32M', - 'lang': 'English', - 'params': '32.3M', - 'description': 'Large and accurate English model.', - }, - { - 'id': 'minishlab/potion-retrieval-32M', - 'name': 'Potion Retrieval 32M', - 'lang': 'English', - 'params': '32.3M', - 'description': 'Optimized specifically for RAG and retrieval tasks.', - }, - { - 'id': 'minishlab/potion-code-16M', - 'name': 'Potion Code 16M', - 'lang': 'Code', - 'params': '16M', - 'description': 'Optimized for code retrieval and analysis.', - }, - { - 'id': 'minishlab/potion-multilingual-128M', - 'name': 'Potion Multilingual 128M', - 'lang': 'Multilingual (101)', - 'params': '128M', - 'description': 'Best for multi-language tasks.', - }, + /// Officially recommended Potion models to start from. + /// + /// A curated, offline catalog — not fetched from Hugging Face (see + /// `docs/adr/0003`), since the editorial fields here cannot be fetched. + // ignore: omit_obvious_property_types explicit type documents the public API + static const List recommendedModels = [ + RecommendedModel( + id: 'minishlab/potion-base-2M', + name: 'Potion Base 2M', + lang: 'English', + params: '1.8M', + description: 'Smallest English model, very fast.', + ), + RecommendedModel( + id: 'minishlab/potion-base-4M', + name: 'Potion Base 4M', + lang: 'English', + params: '3.7M', + description: 'Small and efficient English model.', + ), + RecommendedModel( + id: 'minishlab/potion-base-8M', + name: 'Potion Base 8M', + lang: 'English', + params: '7.5M', + description: 'Balanced English model.', + ), + RecommendedModel( + id: 'minishlab/potion-base-32M', + name: 'Potion Base 32M', + lang: 'English', + params: '32.3M', + description: 'Large and accurate English model.', + ), + RecommendedModel( + id: 'minishlab/potion-retrieval-32M', + name: 'Potion Retrieval 32M', + lang: 'English', + params: '32.3M', + description: 'Optimized specifically for RAG and retrieval tasks.', + ), + RecommendedModel( + id: 'minishlab/potion-code-16M', + name: 'Potion Code 16M', + lang: 'Code', + params: '16M', + description: 'Optimized for code retrieval and analysis.', + ), + RecommendedModel( + id: 'minishlab/potion-multilingual-128M', + name: 'Potion Multilingual 128M', + lang: 'Multilingual (101)', + params: '128M', + description: 'Best for multi-language tasks.', + ), ]; /// Generates a dense vector embedding for the provided [text]. @@ -279,38 +212,42 @@ class Model2Vec { /// - [maxLength]: Maximum number of tokens to keep before truncating. /// /// Returns a [Float32List] representing the text embedding. - /// Throws a [Model2VecException] if generation fails or no model is - /// initialized. - Float32List generateEmbedding(String text, {int maxLength = 512}) { - final dim = embeddingDimension; - return using((arena) { - final textPtr = text.toNativeUtf8(allocator: arena).cast(); - final outVector = arena(dim); - - final res = _bindings.generate_embedding( - textPtr, - outVector, - maxLength, - ); - - if (res != 0) { - throw Model2VecException.fromCode(res, 'Embedding generation failed'); - } + /// Throws a [Model2VecException] if generation fails or no model is loaded. + static Float32List generateEmbedding(String text, {int maxLength = 512}) => + using((arena) { + final textPtr = text.toNativeUtf8(allocator: arena).cast(); + final outData = arena>(); + final outDim = arena(); + final outError = arena>(); + + _check( + native.generate_embedding( + textPtr, + maxLength, + outData, + outDim, + outError, + ), + outError, + ); - return .fromList(outVector.asTypedList(dim)); - }); - } + final dim = outDim.value; + final dataPtr = outData.value; + try { + return Float32List.fromList(dataPtr.asTypedList(dim)); + } finally { + native.free_floats(dataPtr, dim); + } + }); - /// Generates embeddings for multiple [texts] in a highly optimized batch - /// call. + /// Generates embeddings for multiple [texts] in a single batch call. /// /// - [maxLength]: Maximum number of tokens to keep before truncating. - /// - [batchSize]: Size of the internal batches sent to the model for - /// inference. + /// - [batchSize]: Size of the internal batches sent to the model. /// - /// Returns a list of [Float32List], one for each input text. + /// Returns a list of [Float32List], one per input text. /// Throws a [Model2VecException] if generation fails. - List generateBatchEmbeddings( + static List generateBatchEmbeddings( List texts, { int maxLength = 512, int batchSize = 1024, @@ -319,61 +256,66 @@ class Model2Vec { return []; } - final dim = embeddingDimension; - final count = texts.length; - return using((arena) { - final outVectors = arena(count * dim); + final count = texts.length; final textPointers = arena>(count); - for (var i = 0; i < count; i++) { textPointers[i] = texts[i].toNativeUtf8(allocator: arena).cast(); } - final res = _bindings.generate_batch_embeddings_advanced( - textPointers, - count, - outVectors, - maxLength, - batchSize, + final outData = arena>(); + final outDim = arena(); + final outCount = arena(); + final outError = arena>(); + + _check( + native.generate_batch_embeddings_advanced( + textPointers, + count, + maxLength, + batchSize, + outData, + outDim, + outCount, + outError, + ), + outError, ); - if (res != 0) { - throw Model2VecException.fromCode( - res, - 'Batch embedding generation failed', - ); - } - - final results = []; - final flatData = outVectors.asTypedList(count * dim); - for (var i = 0; i < count; i++) { - final start = i * dim; - results.add(.fromList(flatData.sublist(start, start + dim))); + final dim = outDim.value; + final resultCount = outCount.value; + final dataPtr = outData.value; + try { + final flatData = dataPtr.asTypedList(resultCount * dim); + final results = []; + for (var i = 0; i < resultCount; i++) { + final start = i * dim; + // sublist already copies out of native memory into a Dart-owned + // Float32List, so it stays valid after free_floats below. + results.add(flatData.sublist(start, start + dim)); + } + return results; + } finally { + native.free_floats(dataPtr, resultCount * dim); } - - return results; }); } /// Generates an embedding vector asynchronously in a background Isolate. - Future generateEmbeddingAsync( + static Future generateEmbeddingAsync( String text, { int maxLength = 512, }) => Isolate.run( - () => Model2Vec.instance.generateEmbedding( - text, - maxLength: maxLength, - ), + () => Model2Vec.generateEmbedding(text, maxLength: maxLength), ); /// Generates batch embeddings asynchronously in a background Isolate. - Future> generateBatchEmbeddingsAsync( + static Future> generateBatchEmbeddingsAsync( List texts, { int maxLength = 512, int batchSize = 1024, }) => Isolate.run( - () => Model2Vec.instance.generateBatchEmbeddings( + () => Model2Vec.generateBatchEmbeddings( texts, maxLength: maxLength, batchSize: batchSize, @@ -384,177 +326,67 @@ class Model2Vec { /// /// The stream is buffered into batches of [batchSize] to maximize throughput. /// - /// By default, [useIsolate] is `true`, which runs the heavy FFI computations + /// By default, [useIsolate] is `true`, which runs the heavy FFI computation /// in a single background worker isolate, preventing the main thread from - /// stuttering. This is ideal for Flutter applications or processing millions - /// of rows. - /// - /// **Performance Warning for CLI / Server:** - /// If you are writing a pure Dart CLI or server application where blocking - /// the main thread is acceptable (or if you already manage your own isolate - /// pool), setting `useIsolate: false` will eliminate Inter-Process - /// Communication (IPC) overhead, avoiding unnecessary serialization of - /// strings and Float32Lists between isolates, making generation up to 2x - /// faster for small batches. - Stream generateEmbeddingStream( + /// stuttering. Set `useIsolate: false` in a pure CLI/server context where + /// blocking the main thread is acceptable, to avoid isolate IPC overhead. + static Stream generateEmbeddingStream( Stream texts, { int batchSize = 1024, int maxLength = 512, bool useIsolate = true, }) async* { + final batches = batched(texts, batchSize); + if (!useIsolate) { - var buffer = []; - Stream flushBuffer(List batch) async* { + await for (final batch in batches) { final results = generateBatchEmbeddings( batch, maxLength: maxLength, batchSize: batch.length, ); - - for (final res in results) { - yield res; + for (final result in results) { + yield result; } } - - await for (final text in texts) { - buffer.add(text); - - if (buffer.length >= batchSize) { - final currentBatch = buffer; - buffer = []; - yield* flushBuffer(currentBatch); - } - } - - if (buffer.isNotEmpty) { - yield* flushBuffer(buffer); - } - return; } - final mainReceivePort = ReceivePort(); - Isolate? isolate; - SendPort? workerSendPort; - StreamSubscription? sub; - + EmbeddingWorker? worker; try { - isolate = await .spawn(_streamWorker, mainReceivePort.sendPort); - var pendingRequest = Completer(); - - sub = mainReceivePort.listen((message) { - pendingRequest.complete(message); - }); - - final firstMessage = await pendingRequest.future; - if (firstMessage is SendPort) { - workerSendPort = firstMessage; - } else { - throw StateError('Worker isolate failed to initialize'); - } - - var buffer = []; - - Stream flushBuffer(List batch) async* { - pendingRequest = Completer(); - workerSendPort!.send(( - batch: batch, - maxLength: maxLength, - )); - - final message = await pendingRequest.future; - if (message is List) { - for (final res in message) { - yield res; - } - } else if (message is String && message.startsWith('ERROR:')) { - throw Model2VecException(message); - } else { - throw StateError('Unexpected response from worker isolate: $message'); - } - } - - await for (final text in texts) { - buffer.add(text); - - if (buffer.length >= batchSize) { - final currentBatch = buffer; - buffer = []; - yield* flushBuffer(currentBatch); + worker = await EmbeddingWorker.start(); + await for (final batch in batches) { + final results = await worker.embedBatch(batch, maxLength: maxLength); + for (final result in results) { + yield result; } } - - if (buffer.isNotEmpty) { - yield* flushBuffer(buffer); - } } finally { - if (workerSendPort != null) { - workerSendPort.send('close'); - } - - isolate?.kill(); - await sub?.cancel(); - mainReceivePort.close(); + await worker?.close(); } } - static void _streamWorker(SendPort mainSendPort) { - final workerReceivePort = ReceivePort(); - mainSendPort.send(workerReceivePort.sendPort); - - workerReceivePort.listen((message) { - if (message == 'close') { - workerReceivePort.close(); - return; - } - - try { - final data = message as ({List batch, int maxLength}); - - final results = Model2Vec.instance.generateBatchEmbeddings( - data.batch, - maxLength: data.maxLength, - batchSize: data.batch.length, - ); - mainSendPort.send(results); - } on Object catch (e) { - mainSendPort.send('ERROR: $e'); - } - }); + /// Throws a typed [Model2VecException] when [code] is a native failure + /// (non-zero), reading and freeing the message written to [outError]. + static void _check(int code, Pointer> outError) { + if (code != 0) { + throw _nativeException(code, outError.value); + } } - static String _resolveLibPath() { - final libName = Platform.isLinux - ? 'libm2v_ffi.so' - : Platform.isMacOS - ? 'libm2v_ffi.dylib' - : 'm2v_ffi.dll'; - - final path = Directory.current.path; - - final searchPaths = [ - p.join(path, libName), - p.join(path, 'lib', libName), - p.join(path, '.dart_tool', 'lib', libName), - p.join(path, 'packages', 'model2vec', '.dart_tool', 'lib', libName), - p.join(path, 'native', 'target', 'release', libName), - p.join( - path, - 'packages', - 'model2vec', - 'native', - 'target', - 'release', - libName, - ), - ]; - - for (final path in searchPaths) { - if (File(path).existsSync()) { - return path; - } + /// Reads the native error message at [errPtr] (freeing it) and builds a + /// typed [Model2VecException] for the given native [code]. + static Model2VecException _nativeException(int code, Pointer errPtr) { + if (errPtr == nullptr) { + return Model2VecException.fromNative(code, 'native error (code $code)'); + } + try { + return Model2VecException.fromNative( + code, + errPtr.cast().toDartString(), + ); + } finally { + native.free_string(errPtr); } - - return libName; } } diff --git a/lib/src/model2vec_bindings.g.dart b/lib/src/model2vec_bindings.g.dart index 58cbc54..86abc50 100644 --- a/lib/src/model2vec_bindings.g.dart +++ b/lib/src/model2vec_bindings.g.dart @@ -6,286 +6,143 @@ // // Generated by `package:ffigen`. // ignore_for_file: type=lint, unused_import -import 'dart:ffi' as ffi; - -/// Bindings for model2vec-rs. -class Model2VecBindings { - /// Holds the symbol lookup function. - final ffi.Pointer Function(String symbolName) - _lookup; - - /// The symbols are looked up in [dynamicLibrary]. - Model2VecBindings(ffi.DynamicLibrary dynamicLibrary) - : _lookup = dynamicLibrary.lookup; - - /// The symbols are looked up with [lookup]. - Model2VecBindings.fromLookup( - ffi.Pointer Function(String symbolName) lookup, - ) : _lookup = lookup; - - int init_embedder( - ffi.Pointer model_path, - ) { - return _init_embedder( - model_path, - ); - } - - late final _init_embedderPtr = - _lookup)>>( - 'init_embedder', - ); - late final _init_embedder = _init_embedderPtr - .asFunction)>(); - - int init_embedder_advanced( - ffi.Pointer model_path, - ffi.Pointer hf_token, - ffi.Pointer cache_dir, - int normalize, - ffi.Pointer subfolder, - ) { - return _init_embedder_advanced( - model_path, - hf_token, - cache_dir, - normalize, - subfolder, - ); - } - - late final _init_embedder_advancedPtr = - _lookup< - ffi.NativeFunction< - ffi.Int Function( - ffi.Pointer, - ffi.Pointer, - ffi.Pointer, - ffi.Int, - ffi.Pointer, - ) - > - >('init_embedder_advanced'); - late final _init_embedder_advanced = _init_embedder_advancedPtr - .asFunction< - int Function( - ffi.Pointer, - ffi.Pointer, - ffi.Pointer, - int, - ffi.Pointer, - ) - >(); - - int init_embedder_from_bytes( - ffi.Pointer tokenizer_ptr, - int tokenizer_len, - ffi.Pointer model_ptr, - int model_len, - ffi.Pointer config_ptr, - int config_len, - ) { - return _init_embedder_from_bytes( - tokenizer_ptr, - tokenizer_len, - model_ptr, - model_len, - config_ptr, - config_len, - ); - } - - late final _init_embedder_from_bytesPtr = - _lookup< - ffi.NativeFunction< - ffi.Int Function( - ffi.Pointer, - ffi.UnsignedLong, - ffi.Pointer, - ffi.UnsignedLong, - ffi.Pointer, - ffi.UnsignedLong, - ) - > - >('init_embedder_from_bytes'); - late final _init_embedder_from_bytes = _init_embedder_from_bytesPtr - .asFunction< - int Function( - ffi.Pointer, - int, - ffi.Pointer, - int, - ffi.Pointer, - int, - ) - >(); - - int get_embedding_dimension() { - return _get_embedding_dimension(); - } - - late final _get_embedding_dimensionPtr = - _lookup>( - 'get_embedding_dimension', - ); - late final _get_embedding_dimension = _get_embedding_dimensionPtr - .asFunction(); - - int get_vocabulary_size() { - return _get_vocabulary_size(); - } - - late final _get_vocabulary_sizePtr = - _lookup>('get_vocabulary_size'); - late final _get_vocabulary_size = _get_vocabulary_sizePtr - .asFunction(); +@ffi.DefaultAsset('package:model2vec/model2vec.so') +library; - int is_normalized() { - return _is_normalized(); - } - - late final _is_normalizedPtr = - _lookup>('is_normalized'); - late final _is_normalized = _is_normalizedPtr.asFunction(); - - int get_median_token_length() { - return _get_median_token_length(); - } - - late final _get_median_token_lengthPtr = - _lookup>( - 'get_median_token_length', - ); - late final _get_median_token_length = _get_median_token_lengthPtr - .asFunction(); - - ffi.Pointer tokenize( - ffi.Pointer text, - ) { - return _tokenize( - text, - ); - } - - late final _tokenizePtr = - _lookup< - ffi.NativeFunction< - ffi.Pointer Function(ffi.Pointer) - > - >('tokenize'); - late final _tokenize = _tokenizePtr - .asFunction Function(ffi.Pointer)>(); - - int generate_embedding( - ffi.Pointer text, - ffi.Pointer out_vector, - int max_len, - ) { - return _generate_embedding( - text, - out_vector, - max_len, - ); - } - - late final _generate_embeddingPtr = - _lookup< - ffi.NativeFunction< - ffi.Int Function( - ffi.Pointer, - ffi.Pointer, - ffi.UnsignedLong, - ) - > - >('generate_embedding'); - late final _generate_embedding = _generate_embeddingPtr - .asFunction< - int Function(ffi.Pointer, ffi.Pointer, int) - >(); - - int generate_batch_embeddings( - ffi.Pointer> texts_ptr, - int count, - ffi.Pointer out_vectors, - ) { - return _generate_batch_embeddings( - texts_ptr, - count, - out_vectors, - ); - } - - late final _generate_batch_embeddingsPtr = - _lookup< - ffi.NativeFunction< - ffi.Int Function( - ffi.Pointer>, - ffi.UnsignedLong, - ffi.Pointer, - ) - > - >('generate_batch_embeddings'); - late final _generate_batch_embeddings = _generate_batch_embeddingsPtr - .asFunction< - int Function( - ffi.Pointer>, - int, - ffi.Pointer, - ) - >(); - - int generate_batch_embeddings_advanced( - ffi.Pointer> texts_ptr, - int count, - ffi.Pointer out_vectors, - int max_length, - int batch_size, - ) { - return _generate_batch_embeddings_advanced( - texts_ptr, - count, - out_vectors, - max_length, - batch_size, - ); - } - - late final _generate_batch_embeddings_advancedPtr = - _lookup< - ffi.NativeFunction< - ffi.Int Function( - ffi.Pointer>, - ffi.UnsignedLong, - ffi.Pointer, - ffi.UnsignedLong, - ffi.UnsignedLong, - ) - > - >('generate_batch_embeddings_advanced'); - late final _generate_batch_embeddings_advanced = - _generate_batch_embeddings_advancedPtr - .asFunction< - int Function( - ffi.Pointer>, - int, - ffi.Pointer, - int, - int, - ) - >(); - - void free_string( - ffi.Pointer s, - ) { - return _free_string( - s, - ); - } +import 'dart:ffi' as ffi; - late final _free_stringPtr = - _lookup)>>( - 'free_string', - ); - late final _free_string = _free_stringPtr - .asFunction)>(); -} +@ffi.Native< + ffi.Int Function( + ffi.Pointer, + ffi.Pointer, + ffi.Pointer, + ffi.Int, + ffi.Pointer, + ffi.Pointer>, + ) +>() +external int init_embedder_advanced( + ffi.Pointer model_path, + ffi.Pointer hf_token, + ffi.Pointer cache_dir, + int normalize, + ffi.Pointer subfolder, + ffi.Pointer> out_error, +); + +@ffi.Native< + ffi.Int Function( + ffi.Pointer, + ffi.Size, + ffi.Pointer, + ffi.Size, + ffi.Pointer, + ffi.Size, + ffi.Pointer>, + ) +>() +external int init_embedder_from_bytes( + ffi.Pointer tokenizer_ptr, + int tokenizer_len, + ffi.Pointer model_ptr, + int model_len, + ffi.Pointer config_ptr, + int config_len, + ffi.Pointer> out_error, +); + +@ffi.Native< + ffi.Int Function(ffi.Pointer, ffi.Pointer>) +>() +external int get_embedding_dimension( + ffi.Pointer out_value, + ffi.Pointer> out_error, +); + +@ffi.Native< + ffi.Int Function(ffi.Pointer, ffi.Pointer>) +>() +external int get_vocabulary_size( + ffi.Pointer out_value, + ffi.Pointer> out_error, +); + +@ffi.Native< + ffi.Int Function(ffi.Pointer, ffi.Pointer>) +>() +external int is_normalized( + ffi.Pointer out_value, + ffi.Pointer> out_error, +); + +@ffi.Native< + ffi.Int Function(ffi.Pointer, ffi.Pointer>) +>() +external int get_median_token_length( + ffi.Pointer out_value, + ffi.Pointer> out_error, +); + +@ffi.Native< + ffi.Int Function( + ffi.Pointer, + ffi.Pointer>, + ffi.Pointer>, + ) +>() +external int tokenize( + ffi.Pointer text, + ffi.Pointer> out_json, + ffi.Pointer> out_error, +); + +@ffi.Native< + ffi.Int Function( + ffi.Pointer, + ffi.Size, + ffi.Pointer>, + ffi.Pointer, + ffi.Pointer>, + ) +>() +external int generate_embedding( + ffi.Pointer text, + int max_length, + ffi.Pointer> out_data, + ffi.Pointer out_dim, + ffi.Pointer> out_error, +); + +@ffi.Native< + ffi.Int Function( + ffi.Pointer>, + ffi.Size, + ffi.Size, + ffi.Size, + ffi.Pointer>, + ffi.Pointer, + ffi.Pointer, + ffi.Pointer>, + ) +>() +external int generate_batch_embeddings_advanced( + ffi.Pointer> texts_ptr, + int count, + int max_length, + int batch_size, + ffi.Pointer> out_data, + ffi.Pointer out_dim, + ffi.Pointer out_count, + ffi.Pointer> out_error, +); + +@ffi.Native)>() +external void free_string( + ffi.Pointer s, +); + +@ffi.Native, ffi.Size)>() +external void free_floats( + ffi.Pointer ptr, + int len, +); diff --git a/lib/src/recommended_model.dart b/lib/src/recommended_model.dart new file mode 100644 index 0000000..89788ea --- /dev/null +++ b/lib/src/recommended_model.dart @@ -0,0 +1,45 @@ +/// A curated Potion model recommendation from the Model2Vec catalog. +/// +/// This is editorial data (a friendly name, language and "best for" blurb) that +/// only makes sense hand-curated — see `docs/adr/0003` for why the catalog is a +/// hardcoded constant rather than fetched from Hugging Face. +final class RecommendedModel { + /// Creates a catalog entry. + const RecommendedModel({ + required this.id, + required this.name, + required this.lang, + required this.params, + required this.description, + }); + + /// Hugging Face repo id, e.g. `minishlab/potion-base-8M`. + final String id; + + /// Human-friendly display name. + final String name; + + /// Language coverage, e.g. `English` or `Multilingual (101)`. + final String lang; + + /// Approximate parameter count, e.g. `7.5M`. + final String params; + + /// One-line description of what the model is best for. + final String description; + + @override + bool operator ==(Object other) => + other is RecommendedModel && + other.id == id && + other.name == name && + other.lang == lang && + other.params == params && + other.description == description; + + @override + int get hashCode => Object.hash(id, name, lang, params, description); + + @override + String toString() => 'RecommendedModel($id)'; +} diff --git a/lib/src/worker_protocol.dart b/lib/src/worker_protocol.dart new file mode 100644 index 0000000..4f0d355 --- /dev/null +++ b/lib/src/worker_protocol.dart @@ -0,0 +1,123 @@ +import 'dart:async'; +import 'dart:collection'; +import 'dart:typed_data'; + +import 'channel.dart'; +import 'exception.dart'; + +/// The request payload sent to an embedding worker over a [Channel]. +typedef EmbedRequest = ({List batch, int maxLength}); + +class _Request { + _Request(this.batch, this.maxLength) : completer = Completer(); + + final List batch; + final int maxLength; + final Completer> completer; +} + +/// A serialized request/response protocol over a [Channel]. +/// +/// One request is in flight at a time; concurrent [embedBatch] calls queue and +/// are serviced in submission order. This mirrors the single active native +/// model — there is nothing to gain from overlapping requests, and serializing +/// keeps any model switch deterministically ordered against embedding calls. +/// +/// The protocol is pure with respect to transport: it only talks to a +/// [Channel], so it can be driven by an in-memory fake in tests with no +/// isolate. If the channel closes on its own (the worker died), every +/// outstanding request is failed rather than left hanging. +class WorkerProtocol { + WorkerProtocol(this._channel) { + _subscription = _channel.incoming.listen( + _onMessage, + onDone: _onChannelClosed, + ); + } + + final Channel _channel; + late final StreamSubscription _subscription; + final _queue = Queue<_Request>(); + var _inFlight = false; + var _acceptingRequests = true; + var _disposed = false; + + /// Submits [batch] for embedding. Safe to call concurrently; calls are + /// serviced one at a time, in submission order. + Future> embedBatch( + List batch, { + required int maxLength, + }) { + if (!_acceptingRequests) { + return Future>.error( + const Model2VecException( + Model2VecErrorKind.unknown, + 'embedding worker is closed', + ), + ); + } + final request = _Request(batch, maxLength); + _queue.add(request); + _pump(); + return request.completer.future; + } + + void _pump() { + if (_inFlight || _queue.isEmpty) { + return; + } + _inFlight = true; + final head = _queue.first; + _channel.send((batch: head.batch, maxLength: head.maxLength)); + } + + void _onMessage(Object? message) { + if (!_inFlight || _queue.isEmpty) { + return; // stray reply; nothing is awaiting it + } + final request = _queue.removeFirst(); + _inFlight = false; + + if (message is List) { + request.completer.complete(message); + } else if (message is Model2VecException) { + request.completer.completeError(message); + } else { + request.completer.completeError( + StateError('unexpected worker reply: $message'), + ); + } + + _pump(); + } + + /// The channel closed without us asking — the worker went away. Reject new + /// requests and fail everything outstanding. + void _onChannelClosed() { + if (!_acceptingRequests) { + return; + } + _acceptingRequests = false; + _failAll('embedding worker terminated unexpectedly'); + } + + /// Closes the protocol, failing any queued or in-flight requests, then tears + /// down the underlying [Channel]. Idempotent. + Future close() async { + if (_disposed) { + return; + } + _disposed = true; + _acceptingRequests = false; + await _subscription.cancel(); + _failAll('embedding worker closed before the request completed'); + await _channel.close(); + } + + void _failAll(String message) { + final error = Model2VecException(Model2VecErrorKind.unknown, message); + while (_queue.isNotEmpty) { + _queue.removeFirst().completer.completeError(error); + } + } +} diff --git a/native/Cargo.lock b/native/Cargo.lock index 1146714..2a4342d 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -711,7 +711,6 @@ dependencies = [ "hf-hub", "libc", "ndarray", - "once_cell", "safetensors", "serde", "serde_json", diff --git a/native/Cargo.toml b/native/Cargo.toml index 03aed30..e40c966 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -9,7 +9,6 @@ crate-type = ["staticlib", "cdylib"] [dependencies] libc = "0.2.186" anyhow = "1.0.102" -once_cell = "1.21.4" serde_json = "1.0.150" serde = { version = "1.0.228", features = ["derive"] } tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } diff --git a/native/model2vec.h b/native/model2vec.h index 0186711..765aec5 100644 --- a/native/model2vec.h +++ b/native/model2vec.h @@ -1,50 +1,76 @@ -int init_embedder(const char* model_path); +#include + +/* + * Error model + * ----------- + * Every fallible function returns an int status code: + * 0 success + * > 0 failure; the value is a stable Model2VecErrorKind (see below) + * On failure the function writes an owned, human-readable message into + * `*out_error` (when `out_error` is non-null). The caller must release that + * message with `free_string`. On success `*out_error` is set to NULL. + * + * Output buffers + * -------------- + * The native layer owns every output allocation. `generate_*` write a pointer + * to a freshly allocated float buffer into `*out_data`; the caller copies it + * and then releases it with `free_floats(ptr, len)`. String outputs are + * released with `free_string`. + * + * Stable error codes (Model2VecErrorKind): + * 1 not initialized 2 model load failed 3 init from bytes failed + * 4 lock poisoned 5 null argument 6 tokenization failed + * 7 empty result 8 panic + */ int init_embedder_advanced( const char* model_path, const char* hf_token, const char* cache_dir, int normalize, - const char* subfolder + const char* subfolder, + char** out_error ); int init_embedder_from_bytes( const unsigned char* tokenizer_ptr, - unsigned long tokenizer_len, + size_t tokenizer_len, const unsigned char* model_ptr, - unsigned long model_len, + size_t model_len, const unsigned char* config_ptr, - unsigned long config_len + size_t config_len, + char** out_error ); -int get_embedding_dimension(); +int get_embedding_dimension(int* out_value, char** out_error); -int get_vocabulary_size(); +int get_vocabulary_size(int* out_value, char** out_error); -int is_normalized(); +int is_normalized(int* out_value, char** out_error); -int get_median_token_length(); +int get_median_token_length(int* out_value, char** out_error); -char* tokenize(const char* text); +int tokenize(const char* text, char** out_json, char** out_error); int generate_embedding( const char* text, - float* out_vector, - unsigned long max_len -); - -int generate_batch_embeddings( - const char** texts_ptr, - unsigned long count, - float* out_vectors + size_t max_length, + float** out_data, + size_t* out_dim, + char** out_error ); int generate_batch_embeddings_advanced( const char** texts_ptr, - unsigned long count, - float* out_vectors, - unsigned long max_length, - unsigned long batch_size + size_t count, + size_t max_length, + size_t batch_size, + float** out_data, + size_t* out_dim, + size_t* out_count, + char** out_error ); void free_string(char* s); + +void free_floats(float* ptr, size_t len); diff --git a/native/src/lib.rs b/native/src/lib.rs index 2b80909..89bca8f 100644 --- a/native/src/lib.rs +++ b/native/src/lib.rs @@ -1,236 +1,359 @@ +use anyhow::Result; +use std::any::Any; use std::ffi::{CStr, CString}; use std::os::raw::c_char; -use std::ptr; -use std::sync::RwLock; -use anyhow::Result; +use std::panic::{catch_unwind, AssertUnwindSafe}; use std::path::Path; +use std::sync::RwLock; mod model; use model::StaticModel; static MODEL: RwLock> = RwLock::new(None); -#[no_mangle] -pub extern "C" fn init_embedder(model_path: *const c_char) -> i32 { - init_embedder_advanced(model_path, ptr::null(), ptr::null(), -1, ptr::null()) +// Stable error codes. Keep in sync with Model2VecErrorKind on the Dart side. +const CODE_NOT_INITIALIZED: i32 = 1; +const CODE_MODEL_LOAD_FAILED: i32 = 2; +const CODE_INIT_FROM_BYTES_FAILED: i32 = 3; +const CODE_LOCK_POISONED: i32 = 4; +const CODE_NULL_ARGUMENT: i32 = 5; +const CODE_TOKENIZATION_FAILED: i32 = 6; +const CODE_EMPTY_RESULT: i32 = 7; +const CODE_PANIC: i32 = 8; + +/// A failure that crosses the FFI boundary: a stable code plus an owned message. +struct FfiError { + code: i32, + message: String, } -#[no_mangle] -pub extern "C" fn init_embedder_advanced( - model_path: *const c_char, - hf_token: *const c_char, - cache_dir: *const c_char, - normalize: i32, // -1: default, 0: false, 1: true - subfolder: *const c_char, -) -> i32 { - let path = unsafe { - if model_path.is_null() { return -1; } - CStr::from_ptr(model_path).to_string_lossy().into_owned() - }; +impl FfiError { + fn new(code: i32, message: impl Into) -> Self { + FfiError { code, message: message.into() } + } +} - let token = unsafe { - if hf_token.is_null() { None } - else { Some(CStr::from_ptr(hf_token).to_string_lossy().into_owned()) } - }; +/// Runs `f`, converting any error or panic into a status code and writing the +/// message into `*out_error`. `*out_error` is set to null on success. +fn run_ffi(out_error: *mut *mut c_char, f: F) -> i32 +where + F: FnOnce() -> Result<(), FfiError>, +{ + write_error(out_error, None); + match catch_unwind(AssertUnwindSafe(f)) { + Ok(Ok(())) => 0, + Ok(Err(e)) => { + write_error(out_error, Some(&e.message)); + e.code + } + Err(panic) => { + write_error(out_error, Some(&panic_message(&panic))); + CODE_PANIC + } + } +} - let cache = unsafe { - if cache_dir.is_null() { None } - else { Some(CStr::from_ptr(cache_dir).to_string_lossy().into_owned()) } - }; +fn panic_message(panic: &Box) -> String { + if let Some(s) = panic.downcast_ref::<&str>() { + format!("panic in native code: {s}") + } else if let Some(s) = panic.downcast_ref::() { + format!("panic in native code: {s}") + } else { + "panic in native code".to_string() + } +} - let norm = match normalize { - 0 => Some(false), - 1 => Some(true), - _ => None, +/// Writes an owned copy of `msg` into `*out_error`, or null when `msg` is None. +fn write_error(out_error: *mut *mut c_char, msg: Option<&str>) { + if out_error.is_null() { + return; + } + let value = match msg { + None => std::ptr::null_mut(), + Some(m) => match CString::new(m.replace('\0', " ")) { + Ok(c) => c.into_raw(), + Err(_) => std::ptr::null_mut(), + }, }; + unsafe { *out_error = value }; +} - let sub = unsafe { - if subfolder.is_null() { None } - else { Some(CStr::from_ptr(subfolder).to_string_lossy().into_owned()) } - }; +/// Acquires a read lock and hands the active model to `f`. +fn with_model(f: F) -> Result<(), FfiError> +where + F: FnOnce(&StaticModel) -> Result<(), FfiError>, +{ + let lock = MODEL + .read() + .map_err(|_| FfiError::new(CODE_LOCK_POISONED, "model lock poisoned"))?; + let model = lock + .as_ref() + .ok_or_else(|| FfiError::new(CODE_NOT_INITIALIZED, "no model initialized; call initEmbedder first"))?; + f(model) +} - match load_model_advanced(&path, token.as_deref(), cache.as_deref(), norm, sub.as_deref()) { - Ok(_) => 0, - Err(_) => -2, +fn opt_cstr(p: *const c_char) -> Option { + if p.is_null() { + None + } else { + Some(unsafe { CStr::from_ptr(p) }.to_string_lossy().into_owned()) } } -#[no_mangle] -pub extern "C" fn init_embedder_from_bytes( - tokenizer_ptr: *const u8, - tokenizer_len: usize, - model_ptr: *const u8, - model_len: usize, - config_ptr: *const u8, - config_len: usize, -) -> i32 { - if tokenizer_ptr.is_null() || model_ptr.is_null() || config_ptr.is_null() { - return -1; +fn write_i32(out: *mut i32, value: i32) { + if !out.is_null() { + unsafe { *out = value }; } +} - let tokenizer_bytes = unsafe { std::slice::from_raw_parts(tokenizer_ptr, tokenizer_len) }; - let model_bytes = unsafe { std::slice::from_raw_parts(model_ptr, model_len) }; - let config_bytes = unsafe { std::slice::from_raw_parts(config_ptr, config_len) }; - - match StaticModel::from_bytes(tokenizer_bytes, model_bytes, config_bytes, None) { - Ok(model) => { - let mut model_lock = match MODEL.write() { - Ok(lock) => lock, - Err(_) => return -4, // Poisoned lock - }; - *model_lock = Some(model); - 0 - } - Err(_) => -3, +fn write_usize(out: *mut usize, value: usize) { + if !out.is_null() { + unsafe { *out = value }; } } +/// Moves `v` into a caller-owned allocation and writes its pointer into +/// `*out_data`. Released by `free_floats(ptr, v.len())`. +fn alloc_floats(out_data: *mut *mut f32, v: Vec) { + if out_data.is_null() { + return; // v is dropped + } + let mut boxed = v.into_boxed_slice(); + let ptr = boxed.as_mut_ptr(); + std::mem::forget(boxed); + unsafe { *out_data = ptr }; +} - -fn load_model_advanced( - path: &str, - token: Option<&str>, - cache_dir: Option<&str>, - normalize: Option, - subfolder: Option<&str> -) -> Result<()> { - let cache_path = cache_dir.map(Path::new); - let model = StaticModel::from_pretrained(path, token, cache_path, normalize, subfolder)?; - - let mut model_lock = MODEL.write().map_err(|_| anyhow::anyhow!("Poisoned lock"))?; - *model_lock = Some(model); +fn write_string(out: *mut *mut c_char, s: String) -> Result<(), FfiError> { + if out.is_null() { + return Ok(()); + } + let c = CString::new(s) + .map_err(|_| FfiError::new(CODE_TOKENIZATION_FAILED, "result contained an interior NUL byte"))?; + unsafe { *out = c.into_raw() }; Ok(()) } #[no_mangle] -pub extern "C" fn get_embedding_dimension() -> i32 { - MODEL.read().ok() - .and_then(|lock| lock.as_ref().map(|m| m.dim() as i32)) - .unwrap_or(-1) +pub extern "C" fn init_embedder_advanced( + model_path: *const c_char, + hf_token: *const c_char, + cache_dir: *const c_char, + normalize: i32, // -1: default, 0: false, 1: true + subfolder: *const c_char, + out_error: *mut *mut c_char, +) -> i32 { + run_ffi(out_error, || { + if model_path.is_null() { + return Err(FfiError::new(CODE_NULL_ARGUMENT, "model_path is null")); + } + let path = unsafe { CStr::from_ptr(model_path) }.to_string_lossy().into_owned(); + let token = opt_cstr(hf_token); + let cache = opt_cstr(cache_dir); + let sub = opt_cstr(subfolder); + let norm = match normalize { + 0 => Some(false), + 1 => Some(true), + _ => None, + }; + + let model = StaticModel::from_pretrained( + &path, + token.as_deref(), + cache.as_deref().map(Path::new), + norm, + sub.as_deref(), + ) + .map_err(|e| FfiError::new(CODE_MODEL_LOAD_FAILED, format!("failed to load model '{path}': {e}")))?; + + let mut lock = MODEL + .write() + .map_err(|_| FfiError::new(CODE_LOCK_POISONED, "model lock poisoned"))?; + *lock = Some(model); + Ok(()) + }) } #[no_mangle] -pub extern "C" fn get_vocabulary_size() -> i32 { - MODEL.read().ok() - .and_then(|lock| lock.as_ref().map(|m| m.vocabulary_size() as i32)) - .unwrap_or(-1) +pub extern "C" fn init_embedder_from_bytes( + tokenizer_ptr: *const u8, + tokenizer_len: usize, + model_ptr: *const u8, + model_len: usize, + config_ptr: *const u8, + config_len: usize, + out_error: *mut *mut c_char, +) -> i32 { + run_ffi(out_error, || { + if tokenizer_ptr.is_null() || model_ptr.is_null() || config_ptr.is_null() { + return Err(FfiError::new(CODE_NULL_ARGUMENT, "one or more byte pointers are null")); + } + let tokenizer_bytes = unsafe { std::slice::from_raw_parts(tokenizer_ptr, tokenizer_len) }; + let model_bytes = unsafe { std::slice::from_raw_parts(model_ptr, model_len) }; + let config_bytes = unsafe { std::slice::from_raw_parts(config_ptr, config_len) }; + + let model = StaticModel::from_bytes(tokenizer_bytes, model_bytes, config_bytes, None) + .map_err(|e| FfiError::new(CODE_INIT_FROM_BYTES_FAILED, format!("init from bytes failed: {e}")))?; + + let mut lock = MODEL + .write() + .map_err(|_| FfiError::new(CODE_LOCK_POISONED, "model lock poisoned"))?; + *lock = Some(model); + Ok(()) + }) } #[no_mangle] -pub extern "C" fn is_normalized() -> i32 { - MODEL.read().ok() - .and_then(|lock| lock.as_ref().map(|m| if m.is_normalized() { 1 } else { 0 })) - .unwrap_or(-1) +pub extern "C" fn get_embedding_dimension(out_value: *mut i32, out_error: *mut *mut c_char) -> i32 { + run_ffi(out_error, || with_model(|m| { + write_i32(out_value, m.dim() as i32); + Ok(()) + })) } #[no_mangle] -pub extern "C" fn get_median_token_length() -> i32 { - MODEL.read().ok() - .and_then(|lock| lock.as_ref().map(|m| m.median_token_length() as i32)) - .unwrap_or(-1) +pub extern "C" fn get_vocabulary_size(out_value: *mut i32, out_error: *mut *mut c_char) -> i32 { + run_ffi(out_error, || with_model(|m| { + write_i32(out_value, m.vocabulary_size() as i32); + Ok(()) + })) } #[no_mangle] -pub extern "C" fn tokenize(text: *const c_char) -> *mut c_char { - let model_lock = match MODEL.read() { - Ok(lock) => lock, - Err(_) => return ptr::null_mut(), - }; - let model = match model_lock.as_ref() { - Some(m) => m, - None => return ptr::null_mut(), - }; - - let input_text = unsafe { - if text.is_null() { return ptr::null_mut(); } - CStr::from_ptr(text).to_string_lossy() - }; - - let encoding = model.tokenizer.encode(input_text.to_string(), false).unwrap(); - let tokens = encoding.get_tokens(); - let json = serde_json::to_string(&tokens).unwrap_or_default(); - CString::new(json).unwrap().into_raw() +pub extern "C" fn is_normalized(out_value: *mut i32, out_error: *mut *mut c_char) -> i32 { + run_ffi(out_error, || with_model(|m| { + write_i32(out_value, if m.is_normalized() { 1 } else { 0 }); + Ok(()) + })) } #[no_mangle] -pub extern "C" fn generate_embedding(text: *const c_char, out_vector: *mut f32, max_length: usize) -> i32 { - let model_lock = match MODEL.read() { - Ok(lock) => lock, - Err(_) => return -4, - }; - let model = match model_lock.as_ref() { - Some(m) => m, - None => return -1, - }; - - if text.is_null() || out_vector.is_null() { return -2; } +pub extern "C" fn get_median_token_length(out_value: *mut i32, out_error: *mut *mut c_char) -> i32 { + run_ffi(out_error, || with_model(|m| { + write_i32(out_value, m.median_token_length() as i32); + Ok(()) + })) +} - let text_str = unsafe { CStr::from_ptr(text).to_string_lossy().into_owned() }; - let results = model.encode_with_args(&[text_str], Some(max_length), 1); - - if let Some(embedding) = results.first() { - unsafe { - ptr::copy_nonoverlapping(embedding.as_ptr(), out_vector, model.dim()); +#[no_mangle] +pub extern "C" fn tokenize( + text: *const c_char, + out_json: *mut *mut c_char, + out_error: *mut *mut c_char, +) -> i32 { + run_ffi(out_error, || { + if text.is_null() { + return Err(FfiError::new(CODE_NULL_ARGUMENT, "text is null")); } - 0 - } else { - -5 // Embedding generation resulted in empty array - } + let input = unsafe { CStr::from_ptr(text) }.to_string_lossy().into_owned(); + with_model(|model| { + let encoding = model + .tokenizer + .encode(input, false) + .map_err(|e| FfiError::new(CODE_TOKENIZATION_FAILED, format!("tokenization failed: {e}")))?; + let json = serde_json::to_string(encoding.get_tokens()) + .map_err(|e| FfiError::new(CODE_TOKENIZATION_FAILED, format!("failed to serialize tokens: {e}")))?; + write_string(out_json, json) + }) + }) } #[no_mangle] -pub extern "C" fn generate_batch_embeddings(texts_ptr: *const *const c_char, count: usize, out_vectors: *mut f32) -> i32 { - generate_batch_embeddings_advanced(texts_ptr, count, out_vectors, 512, 1024) +pub extern "C" fn generate_embedding( + text: *const c_char, + max_length: usize, + out_data: *mut *mut f32, + out_dim: *mut usize, + out_error: *mut *mut c_char, +) -> i32 { + run_ffi(out_error, || { + if text.is_null() { + return Err(FfiError::new(CODE_NULL_ARGUMENT, "text is null")); + } + let text_str = unsafe { CStr::from_ptr(text) }.to_string_lossy().into_owned(); + with_model(|model| { + let results = model + .encode_with_args(&[text_str], Some(max_length), 1) + .map_err(|e| FfiError::new(CODE_TOKENIZATION_FAILED, format!("{e}")))?; + let embedding = results + .into_iter() + .next() + .ok_or_else(|| FfiError::new(CODE_EMPTY_RESULT, "embedding result was empty"))?; + write_usize(out_dim, embedding.len()); + alloc_floats(out_data, embedding); + Ok(()) + }) + }) } #[no_mangle] pub extern "C" fn generate_batch_embeddings_advanced( texts_ptr: *const *const c_char, count: usize, - out_vectors: *mut f32, max_length: usize, batch_size: usize, + out_data: *mut *mut f32, + out_dim: *mut usize, + out_count: *mut usize, + out_error: *mut *mut c_char, ) -> i32 { - let model_lock = match MODEL.read() { - Ok(lock) => lock, - Err(_) => return -4, // Poisoned lock - }; - let model = match model_lock.as_ref() { - Some(m) => m, - None => return -1, // Not initialized - }; - - if texts_ptr.is_null() || out_vectors.is_null() { return -2; } - - let mut texts = Vec::with_capacity(count); - for i in 0..count { - unsafe { - let ptr = *texts_ptr.add(i); - if ptr.is_null() { return -3; } - texts.push(CStr::from_ptr(ptr).to_string_lossy().into_owned()); + run_ffi(out_error, || { + if texts_ptr.is_null() { + return Err(FfiError::new(CODE_NULL_ARGUMENT, "texts pointer is null")); + } + let mut texts = Vec::with_capacity(count); + for i in 0..count { + let ptr = unsafe { *texts_ptr.add(i) }; + if ptr.is_null() { + return Err(FfiError::new(CODE_NULL_ARGUMENT, format!("null text at index {i}"))); + } + texts.push(unsafe { CStr::from_ptr(ptr) }.to_string_lossy().into_owned()); } - } - let results = model.encode_with_args(&texts, Some(max_length), batch_size); - let dim = model.dim(); + with_model(|model| { + let results = model + .encode_with_args(&texts, Some(max_length), batch_size) + .map_err(|e| FfiError::new(CODE_TOKENIZATION_FAILED, format!("{e}")))?; + if results.len() != count { + return Err(FfiError::new( + CODE_EMPTY_RESULT, + format!("expected {count} embeddings, got {}", results.len()), + )); + } + // dim + data are produced under this single read lock, so the + // dimension the caller frees against always matches the data. + let dim = model.dim(); + let mut flat = Vec::with_capacity(count * dim); + for embedding in &results { + if embedding.len() != dim { + return Err(FfiError::new(CODE_EMPTY_RESULT, "embedding dimension mismatch")); + } + flat.extend_from_slice(embedding); + } + write_usize(out_dim, dim); + write_usize(out_count, count); + alloc_floats(out_data, flat); + Ok(()) + }) + }) +} - if results.len() != count { - return -5; // Generation failed or mismatched count +#[no_mangle] +pub extern "C" fn free_string(s: *mut c_char) { + if s.is_null() { + return; } - - for (i, embedding) in results.iter().enumerate() { - unsafe { - let target_ptr = out_vectors.add(i * dim); - ptr::copy_nonoverlapping(embedding.as_ptr(), target_ptr, dim); - } + unsafe { + let _ = CString::from_raw(s); } - 0 } #[no_mangle] -pub extern "C" fn free_string(s: *mut c_char) { +pub extern "C" fn free_floats(ptr: *mut f32, len: usize) { + if ptr.is_null() { + return; + } unsafe { - if s.is_null() { return; } - let _ = CString::from_raw(s); + let _ = Vec::from_raw_parts(ptr, len, len); } } diff --git a/native/src/model.rs b/native/src/model.rs index 110017f..181705b 100644 --- a/native/src/model.rs +++ b/native/src/model.rs @@ -225,13 +225,15 @@ impl StaticModel { sentences: &[String], max_length: Option, batch_size: usize, - ) -> Vec> { + ) -> Result>> { let mut embeddings = Vec::with_capacity(sentences.len()); for batch in sentences.chunks(batch_size) { let truncated: Vec<&str> = batch.iter().map(|text| { max_length.map(|max_tok| Self::truncate_str(text, max_tok, self.median_token_length)).unwrap_or(text.as_str()) }).collect(); - let encodings = self.tokenizer.encode_batch_fast::(truncated.into_iter().map(Into::into).collect(), false).expect("tokenization failed"); + let encodings = self.tokenizer + .encode_batch_fast::(truncated.into_iter().map(Into::into).collect(), false) + .map_err(|e| anyhow!("tokenization failed: {e}"))?; for encoding in encodings { let mut token_ids = encoding.get_ids().to_vec(); if let Some(unk_id) = self.unk_token_id { @@ -243,17 +245,20 @@ impl StaticModel { embeddings.push(self.pool_ids(token_ids)); } } - embeddings + Ok(embeddings) } #[allow(dead_code)] - pub fn encode(&self, sentences: &[String]) -> Vec> { + pub fn encode(&self, sentences: &[String]) -> Result>> { self.encode_with_args(sentences, Some(512), 1024) } #[allow(dead_code)] pub fn encode_single(&self, sentence: &str) -> Vec { - self.encode(&[sentence.to_string()]).into_iter().next().unwrap_or_default() + self.encode(&[sentence.to_string()]) + .ok() + .and_then(|v| v.into_iter().next()) + .unwrap_or_default() } fn pool_ids(&self, ids: Vec) -> Vec { diff --git a/pubspec.yaml b/pubspec.yaml index 1203d77..b078b7a 100644 --- a/pubspec.yaml +++ b/pubspec.yaml @@ -1,6 +1,6 @@ name: model2vec description: A high-performance Dart wrapper for model2vec-rs using Rust FFI. Generate fast, local, and static text embeddings with minimal memory footprint using Native Assets. -version: 1.2.0 +version: 2.0.0 repository: https://github.com/pro100andrey/model2vec homepage: https://github.com/pro100andrey/model2vec issue_tracker: https://github.com/pro100andrey/model2vec/issues @@ -31,9 +31,13 @@ ffigen: name: Model2VecBindings description: Bindings for model2vec-rs. output: "lib/src/model2vec_bindings.g.dart" + ffi-native: + asset-id: "package:model2vec/model2vec.so" headers: entry-points: - "native/model2vec.h" + include-directives: + - "**/model2vec.h" preamble: | // ignore_for_file: always_specify_types // ignore_for_file: camel_case_types diff --git a/test/batcher_test.dart b/test/batcher_test.dart new file mode 100644 index 0000000..f3aea90 --- /dev/null +++ b/test/batcher_test.dart @@ -0,0 +1,53 @@ +import 'package:model2vec/src/batcher.dart'; +import 'package:test/test.dart'; + +void main() { + group('batched', () { + test('groups into full batches with a smaller final remainder', () async { + final result = await batched(Stream.fromIterable([1, 2, 3, 4, 5]), 2) + .toList(); + expect(result, [ + [1, 2], + [3, 4], + [5], + ]); + }); + + test('an exact multiple leaves no remainder', () async { + final result = await batched(Stream.fromIterable([1, 2, 3, 4]), 2) + .toList(); + expect(result, [ + [1, 2], + [3, 4], + ]); + }); + + test('size larger than the input yields a single batch', () async { + final result = await batched(Stream.fromIterable([1, 2]), 10).toList(); + expect(result, [ + [1, 2], + ]); + }); + + test('size 1 yields singletons', () async { + final result = await batched(Stream.fromIterable([1, 2, 3]), 1).toList(); + expect(result, [ + [1], + [2], + [3], + ]); + }); + + test('an empty stream yields nothing', () async { + final result = await batched(const Stream.empty(), 3).toList(); + expect(result, isEmpty); + }); + + test('rejects a size below 1', () { + expect( + batched(Stream.fromIterable([1]), 0).toList(), + throwsArgumentError, + ); + }); + }); +} diff --git a/test/embedding_worker_test.dart b/test/embedding_worker_test.dart new file mode 100644 index 0000000..f051082 --- /dev/null +++ b/test/embedding_worker_test.dart @@ -0,0 +1,123 @@ +import 'dart:isolate'; +import 'dart:typed_data'; + +import 'package:model2vec/src/embedding_worker.dart'; +import 'package:model2vec/src/exception.dart'; +import 'package:model2vec/src/worker_protocol.dart'; +import 'package:test/test.dart'; + +/// Fake worker entry point: replies with one vector per input text, whose only +/// element is the text's length. No model, no native calls. +void echoEntryPoint(SendPort mainSendPort) { + final receivePort = ReceivePort(); + mainSendPort.send(receivePort.sendPort); + receivePort.listen((message) { + if (message == 'close') { + receivePort.close(); + return; + } + final request = message as EmbedRequest; + final results = [ + for (final text in request.batch) + Float32List.fromList([text.length.toDouble()]), + ]; + mainSendPort.send(results); + }); +} + +/// Fake worker entry point that always replies with a typed error. +void errorEntryPoint(SendPort mainSendPort) { + final receivePort = ReceivePort(); + mainSendPort.send(receivePort.sendPort); + receivePort.listen((message) { + if (message == 'close') { + receivePort.close(); + return; + } + mainSendPort.send( + const Model2VecException( + Model2VecErrorKind.tokenizationFailed, + 'boom', + 6, + ), + ); + }); +} + +/// Fake entry point that dies (closes its port, exiting the isolate) on the +/// first request without ever replying. +void dyingEntryPoint(SendPort mainSendPort) { + final receivePort = ReceivePort(); + mainSendPort.send(receivePort.sendPort); + receivePort.listen((message) { + if (message == 'close') { + receivePort.close(); + return; + } + receivePort.close(); // exit without replying + }); +} + +/// Fake entry point that exits immediately, never completing the handshake. +void noHandshakeEntryPoint(SendPort mainSendPort) { + // Returns at once; the isolate exits before sending its SendPort. +} + +void main() { + test('embedBatch returns results from the worker isolate', () async { + final worker = await EmbeddingWorker.start(entryPoint: echoEntryPoint); + final result = await worker.embedBatch(['a', 'bb', 'ccc']); + expect(result.map((v) => v.first).toList(), [1.0, 2.0, 3.0]); + await worker.close(); + }); + + test('multiple batches preserve order over the real isolate', () async { + final worker = await EmbeddingWorker.start(entryPoint: echoEntryPoint); + final r1 = worker.embedBatch(['x']); + final r2 = worker.embedBatch(['yy']); + expect((await r1).first.first, 1.0); + expect((await r2).first.first, 2.0); + await worker.close(); + }); + + test('a typed error survives the isolate boundary', () async { + final worker = await EmbeddingWorker.start(entryPoint: errorEntryPoint); + await expectLater( + worker.embedBatch(['a']), + throwsA( + isA().having( + (e) => e.kind, + 'kind', + Model2VecErrorKind.tokenizationFailed, + ), + ), + ); + await worker.close(); + }); + + test('close tears the worker down and later calls fail fast', () async { + final worker = await EmbeddingWorker.start(entryPoint: echoEntryPoint); + await worker.close(); + await expectLater( + worker.embedBatch(['a']), + throwsA(isA()), + ); + }); + + test('a request fails (does not hang) if the worker dies mid-request', + () async { + final worker = await EmbeddingWorker.start(entryPoint: dyingEntryPoint); + await expectLater( + worker.embedBatch(['a']), + throwsA(isA()), + ); + await worker.close(); + }); + + test('start fails if the worker never completes the handshake', () async { + await expectLater( + EmbeddingWorker.start(entryPoint: noHandshakeEntryPoint), + throwsA(isA()), + ); + }); +} diff --git a/test/error_path_test.dart b/test/error_path_test.dart new file mode 100644 index 0000000..f44b782 --- /dev/null +++ b/test/error_path_test.dart @@ -0,0 +1,65 @@ +import 'dart:typed_data'; + +import 'package:model2vec/model2vec.dart'; +import 'package:test/test.dart'; + +void main() { + group('Model2VecException.fromNative', () { + // Mirrors the CODE_* constants in native/src/lib.rs. + const cases = { + 1: Model2VecErrorKind.notInitialized, + 2: Model2VecErrorKind.modelLoadFailed, + 3: Model2VecErrorKind.initFromBytesFailed, + 4: Model2VecErrorKind.lockPoisoned, + 5: Model2VecErrorKind.nullArgument, + 6: Model2VecErrorKind.tokenizationFailed, + 7: Model2VecErrorKind.emptyResult, + 8: Model2VecErrorKind.panic, + }; + + for (final entry in cases.entries) { + test('code ${entry.key} maps to ${entry.value}', () { + expect( + Model2VecException.fromNative(entry.key, 'msg').kind, + entry.value, + ); + }); + } + + test('an unrecognized code maps to unknown', () { + expect( + Model2VecException.fromNative(99, 'msg').kind, + Model2VecErrorKind.unknown, + ); + }); + + test('carries the native message and code', () { + final e = Model2VecException.fromNative(2, 'boom'); + expect(e.message, 'boom'); + expect(e.code, 2); + expect(e.toString(), contains('boom')); + }); + }); + + group('native error path', () { + test('initEmbedderFromBytes with garbage throws a typed exception', () { + final garbage = Uint8List.fromList([0, 1, 2, 3]); + expect( + () => Model2Vec.initEmbedderFromBytes( + tokenizerBytes: garbage, + modelBytes: garbage, + configBytes: garbage, + ), + throwsA( + isA() + .having( + (e) => e.kind, + 'kind', + Model2VecErrorKind.initFromBytesFailed, + ) + .having((e) => e.message, 'message', isNotEmpty), + ), + ); + }); + }); +} diff --git a/test/model2vec_test.dart b/test/model2vec_test.dart index 5682dc6..2b0b1e3 100644 --- a/test/model2vec_test.dart +++ b/test/model2vec_test.dart @@ -4,39 +4,33 @@ import 'package:model2vec/model2vec.dart'; import 'package:test/test.dart'; void main() { - late Model2Vec m2v; - - setUpAll(() { - m2v = Model2Vec.instance; - }); - group('Model2Vec Production Tests', () { test('Successful initialization (online)', () { - // This MUST be the first init call in the process due to Rust's OnceCell + // Must be the first init call so later tests have a loaded model. expect( - () => m2v.initEmbedder('minishlab/potion-base-2M'), + () => Model2Vec.initEmbedder('minishlab/potion-base-2M'), returnsNormally, ); }); group('Model Metadata', () { test('returns valid dimension and metadata', () { - expect(m2v.embeddingDimension, equals(64)); - expect(m2v.vocabularySize, greaterThan(20000)); - expect(m2v.isNormalized, isTrue); - expect(m2v.medianTokenLength, isPositive); + expect(Model2Vec.embeddingDimension, equals(64)); + expect(Model2Vec.vocabularySize, greaterThan(20000)); + expect(Model2Vec.isNormalized, isTrue); + expect(Model2Vec.medianTokenLength, isPositive); }); }); group('Tokenization', () { test('breaks text into tokens correctly', () { - final tokens = m2v.tokenize('Dart FFI is powerful'); + final tokens = Model2Vec.tokenize('Dart FFI is powerful'); expect(tokens, isNotEmpty); expect(tokens, contains('dart')); }); test('handles empty string tokenization', () { - final tokens = m2v.tokenize(''); + final tokens = Model2Vec.tokenize(''); expect(tokens, isEmpty); }); }); @@ -45,8 +39,8 @@ void main() { test( 'generates vector of correct length and exact values (snapshot)', () { - final vector = m2v.generateEmbedding('Hello world'); - expect(vector.length, equals(m2v.embeddingDimension)); + final vector = Model2Vec.generateEmbedding('Hello world'); + expect(vector.length, equals(Model2Vec.embeddingDimension)); // Exact values for Potion Base 2M on "Hello world" // Using closeTo to handle minor floating point precision differences @@ -58,8 +52,8 @@ void main() { test('batch embedding is consistent with single embedding', () { const text = 'Consistency check'; - final single = m2v.generateEmbedding(text); - final batch = m2v.generateBatchEmbeddings([text]); + final single = Model2Vec.generateEmbedding(text); + final batch = Model2Vec.generateBatchEmbeddings([text]); expect(batch.length, 1); expect(batch[0], orderedEquals(single)); @@ -72,31 +66,31 @@ void main() { '12345 !? @#', 'Теж має працювати', ]; - final results = m2v.generateBatchEmbeddings(texts); + final results = Model2Vec.generateBatchEmbeddings(texts); expect(results.length, equals(texts.length)); for (final v in results) { - expect(v.length, equals(m2v.embeddingDimension)); + expect(v.length, equals(Model2Vec.embeddingDimension)); } }); test('batch with empty list returns empty list', () { - expect(m2v.generateBatchEmbeddings([]), isEmpty); + expect(Model2Vec.generateBatchEmbeddings([]), isEmpty); }); }); group('Asynchronous API', () { test('generates embedding asynchronously in isolate', () async { - final vector = await m2v.generateEmbeddingAsync('Async test'); - expect(vector.length, equals(m2v.embeddingDimension)); + final vector = await Model2Vec.generateEmbeddingAsync('Async test'); + expect(vector.length, equals(Model2Vec.embeddingDimension)); }); test('generates batch embeddings asynchronously in isolate', () async { - final results = await m2v.generateBatchEmbeddingsAsync([ + final results = await Model2Vec.generateBatchEmbeddingsAsync([ 'Async 1', 'Async 2', ]); expect(results.length, 2); - expect(results[0].length, m2v.embeddingDimension); + expect(results[0].length, Model2Vec.embeddingDimension); }); }); @@ -108,7 +102,7 @@ void main() { Iterable.generate(totalItems, (i) => 'This is test string number $i'), ); - final resultStream = m2v.generateEmbeddingStream( + final resultStream = Model2Vec.generateEmbeddingStream( stream, batchSize: 1000, ); @@ -116,18 +110,16 @@ void main() { final results = await resultStream.toList(); expect(results.length, equals(totalItems)); - expect(results.first.length, equals(m2v.embeddingDimension)); + expect(results.first.length, equals(Model2Vec.embeddingDimension)); }); test('works correctly with useIsolate: false', () async { final stream = Stream.fromIterable(['Text 1', 'Text 2', 'Text 3']); - final results = await m2v - .generateEmbeddingStream( - stream, - batchSize: 2, - useIsolate: false, - ) - .toList(); + final results = await Model2Vec.generateEmbeddingStream( + stream, + batchSize: 2, + useIsolate: false, + ).toList(); expect(results.length, equals(3)); }); @@ -139,7 +131,10 @@ void main() { 'Text 3', 'Text 4', ]); - final resultStream = m2v.generateEmbeddingStream(stream, batchSize: 2); + final resultStream = Model2Vec.generateEmbeddingStream( + stream, + batchSize: 2, + ); // Take only 2 elements, which will cancel the subscription early final firstTwo = await resultStream.take(2).toList(); @@ -150,12 +145,12 @@ void main() { group('Model Switching', () { test('can switch between different models successfully', () { // Switch to 8M model (dimension 256) - m2v.initEmbedder('minishlab/potion-base-8M'); - expect(m2v.embeddingDimension, equals(256)); + Model2Vec.initEmbedder('minishlab/potion-base-8M'); + expect(Model2Vec.embeddingDimension, equals(256)); // Switch back to 2M model (dimension 64) - m2v.initEmbedder('minishlab/potion-base-2M'); - expect(m2v.embeddingDimension, equals(64)); + Model2Vec.initEmbedder('minishlab/potion-base-2M'); + expect(Model2Vec.embeddingDimension, equals(64)); }); }); @@ -165,12 +160,12 @@ void main() { 'A very long sentence to test truncation with maxLength parameter'; // With very short maxLength, the embedding should be different from // default - final defaultEmbedding = m2v.generateEmbedding(text); - final shortEmbedding = m2v.generateEmbedding(text, maxLength: 2); + final defaultEmbedding = Model2Vec.generateEmbedding(text); + final shortEmbedding = Model2Vec.generateEmbedding(text, maxLength: 2); expect(defaultEmbedding, isNot(orderedEquals(shortEmbedding))); - final batchEmbedding = m2v.generateBatchEmbeddings( + final batchEmbedding = Model2Vec.generateBatchEmbeddings( [text, text], maxLength: 2, batchSize: 1, @@ -185,7 +180,7 @@ void main() { try { // Now that switching is enabled, this should return normally expect( - () => m2v.initEmbedderAdvanced( + () => Model2Vec.initEmbedderAdvanced( modelPath: 'minishlab/potion-base-2M', cacheDirectory: tempDir.path, ), @@ -198,10 +193,12 @@ void main() { }); group('Recommended Models', () { - test('returns a non-empty list of maps', () { - final models = m2v.getRecommendedModels(); - expect(models, isNotEmpty); - expect(models.first, contains('id')); + test('exposes a non-empty typed catalog', () { + expect(Model2Vec.recommendedModels, isNotEmpty); + expect( + Model2Vec.recommendedModels.first.id, + startsWith('minishlab/'), + ); }); }); }); diff --git a/test/utils_test.dart b/test/utils_test.dart index df98cf2..426d393 100644 --- a/test/utils_test.dart +++ b/test/utils_test.dart @@ -164,16 +164,14 @@ void main() { }); group('Model2VecUtils - Real World', () { - late Model2Vec m2v; - setUpAll(() { - m2v = Model2Vec.instance..initEmbedder('minishlab/potion-base-2M'); + Model2Vec.initEmbedder('minishlab/potion-base-2M'); }); test('semantic similarity makes sense', () { - final vCat = m2v.generateEmbedding('A small cute cat'); - final vKitten = m2v.generateEmbedding('A young little kitten'); - final vSpace = m2v.generateEmbedding( + final vCat = Model2Vec.generateEmbedding('A small cute cat'); + final vKitten = Model2Vec.generateEmbedding('A young little kitten'); + final vSpace = Model2Vec.generateEmbedding( 'The exploration of outer space and planets', ); @@ -197,11 +195,15 @@ void main() { }); test('task similarity example', () { - final vTask1 = m2v.generateEmbedding('Fix the login bug in production'); - final vTask2 = m2v.generateEmbedding( + final vTask1 = Model2Vec.generateEmbedding( + 'Fix the login bug in production', + ); + final vTask2 = Model2Vec.generateEmbedding( 'Resolve authentication issue on server', ); - final vTask3 = m2v.generateEmbedding('Order pizza for the team lunch'); + final vTask3 = Model2Vec.generateEmbedding( + 'Order pizza for the team lunch', + ); final sim12 = Model2VecUtils.cosineSimilarity(vTask1, vTask2); final sim13 = Model2VecUtils.cosineSimilarity(vTask1, vTask3); diff --git a/test/worker_protocol_test.dart b/test/worker_protocol_test.dart new file mode 100644 index 0000000..8f77902 --- /dev/null +++ b/test/worker_protocol_test.dart @@ -0,0 +1,138 @@ +import 'dart:async'; +import 'dart:typed_data'; + +import 'package:model2vec/src/channel.dart'; +import 'package:model2vec/src/exception.dart'; +import 'package:model2vec/src/worker_protocol.dart'; +import 'package:test/test.dart'; + +/// An in-memory [Channel] the test drives directly — no isolate, no model. +class FakeChannel implements Channel { + final sent = []; + final _incoming = StreamController(); + // ignore: omit_obvious_property_types public field needs the annotation + bool closed = false; + + @override + void send(Object? message) => sent.add(message); + + @override + Stream get incoming => _incoming.stream; + + @override + Future close() async { + closed = true; + await _incoming.close(); + } + + /// Simulates a reply arriving from the worker end. + void deliver(Object? message) => _incoming.add(message); +} + +Float32List _vec(List xs) => Float32List.fromList(xs); + +void main() { + late FakeChannel channel; + late WorkerProtocol protocol; + + setUp(() { + channel = FakeChannel(); + protocol = WorkerProtocol(channel); + }); + + test('sends the request and resolves with the delivered result', () async { + final future = protocol.embedBatch(['a'], maxLength: 8); + expect(channel.sent, hasLength(1)); // sent synchronously + + final result = [ + _vec([1, 2, 3]), + ]; + channel.deliver(result); + expect(await future, same(result)); + }); + + test('serializes: the second request waits for the first reply', () async { + final f1 = protocol.embedBatch(['a'], maxLength: 8); + final f2 = protocol.embedBatch(['b'], maxLength: 8); + + // Only the first request is in flight. + expect(channel.sent, hasLength(1)); + + channel.deliver([ + _vec([1]), + ]); + await f1; + + // Now the second is sent. + expect(channel.sent, hasLength(2)); + channel.deliver([ + _vec([2]), + ]); + await f2; + }); + + test('propagates a typed Model2VecException from the worker', () async { + final future = protocol.embedBatch(['a'], maxLength: 8); + channel.deliver( + const Model2VecException( + Model2VecErrorKind.notInitialized, + 'no model', + 1, + ), + ); + + await expectLater( + future, + throwsA( + isA().having( + (e) => e.kind, + 'kind', + Model2VecErrorKind.notInitialized, + ), + ), + ); + }); + + test('an unexpected reply shape becomes a StateError', () async { + final future = protocol.embedBatch(['a'], maxLength: 8); + channel.deliver('garbage'); + await expectLater(future, throwsStateError); + }); + + test('close fails queued and in-flight requests and closes the channel', + () async { + final f1 = protocol.embedBatch(['a'], maxLength: 8); + final f2 = protocol.embedBatch(['b'], maxLength: 8); + + // Attach listeners before close so the failed futures aren't reported as + // unhandled (real callers await the future they submitted). + final expect1 = expectLater(f1, throwsA(isA())); + final expect2 = expectLater(f2, throwsA(isA())); + + await protocol.close(); + + await expect1; + await expect2; + expect(channel.closed, isTrue); + }); + + test('embedBatch after close fails fast', () async { + await protocol.close(); + await expectLater( + protocol.embedBatch(['a'], maxLength: 8), + throwsA(isA()), + ); + }); + + test('fails a pending request when the channel closes on its own', () async { + final future = protocol.embedBatch(['a'], maxLength: 8); + final expectation = expectLater( + future, + throwsA(isA()), + ); + // Simulates the worker dying: the channel closes while a request is in + // flight, without protocol.close() first, so onDone reaches the protocol. + await channel.close(); + await expectation; + }); +} From dd2d0c1daedaff1c2f1f37fb9ecbe5298c4d01b9 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Mon, 6 Jul 2026 18:39:49 +0300 Subject: [PATCH 02/20] feat: local vector index, RAG helpers, worker pool, model lifecycle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the retrieval/RAG layer on top of the embedding engine, plus lifecycle and parallelism helpers. - EmbeddingIndex: in-memory vector store (add/search/remove) with cosine ranking, optional int8-quantized storage (~4x less memory), and binary toBytes/fromBytes persistence — a local retrieval engine for RAG. - RAG helpers: chunkText (overlapping character chunker, maxChars-bounded), Model2VecUtils.similaritySearchWithScores (index + score), and maximalMarginalRelevance (MMR reranking, O(k*n*dim)). - Lifecycle & DX: Model2Vec.isInitialized (non-throwing), unloadModel() (frees native memory via a new free_embedder FFI), modelInfo (metadata in one ModelInfo), Model2VecUtils.dequantizeInt8. - EmbeddingPool: N worker isolates embedding concurrently (the Rust RwLock allows concurrent readers), with order-preserving embedBatches and leak-safe startup. - Native: is_model_loaded + free_embedder FFI functions; bindings regenerated. 92 tests, including no-model/no-network suites for the index (persistence, quantization, corrupt-blob handling), the chunker, MMR, and the pool (fake entry point). README recipes and CHANGELOG updated. --- CHANGELOG.md | 16 ++ README.md | 50 ++++++ lib/model2vec.dart | 4 + lib/src/chunker.dart | 73 ++++++++ lib/src/embedding_index.dart | 278 ++++++++++++++++++++++++++++++ lib/src/embedding_pool.dart | 95 ++++++++++ lib/src/model2vec_base.dart | 25 +++ lib/src/model2vec_bindings.g.dart | 8 + lib/src/model_info.dart | 39 +++++ lib/src/utils.dart | 80 +++++++++ native/model2vec.h | 4 + native/src/lib.rs | 23 +++ test/chunker_test.dart | 57 ++++++ test/embedding_index_test.dart | 151 ++++++++++++++++ test/embedding_pool_test.dart | 70 ++++++++ test/lifecycle_test.dart | 37 ++++ test/utils_test.dart | 94 ++++++++++ 17 files changed, 1104 insertions(+) create mode 100644 lib/src/chunker.dart create mode 100644 lib/src/embedding_index.dart create mode 100644 lib/src/embedding_pool.dart create mode 100644 lib/src/model_info.dart create mode 100644 test/chunker_test.dart create mode 100644 test/embedding_index_test.dart create mode 100644 test/embedding_pool_test.dart create mode 100644 test/lifecycle_test.dart diff --git a/CHANGELOG.md b/CHANGELOG.md index a362b5d..78ec0e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -34,6 +34,22 @@ correctness. **This release is breaking** — see migration below. worker isolate. Worker errors cross the isolate boundary as typed `Model2VecException`s (kind + code preserved) rather than stringified errors. +**New capabilities:** + +- **Local vector index.** `EmbeddingIndex` — store embeddings by id, then + `search` the nearest by cosine similarity. Optional int8-quantized storage + (~4x less memory) and binary `toBytes`/`fromBytes` persistence. Turns the + package into a local retrieval engine for RAG. +- **RAG pipeline helpers.** `chunkText` (overlapping character chunker), + `Model2VecUtils.similaritySearchWithScores` (index + score), and + `Model2VecUtils.maximalMarginalRelevance` (MMR reranking for diverse results). +- **Lifecycle & DX.** `Model2Vec.isInitialized` (non-throwing check), + `Model2Vec.unloadModel()` (free the native model), `Model2Vec.modelInfo` + (all metadata in one `ModelInfo`), and `Model2VecUtils.dequantizeInt8` + (the inverse of `quantizeToInt8`). +- **Parallel worker pool.** `EmbeddingPool` fans batches across N worker + isolates to embed concurrently across CPU cores. + **Migration:** | 1.x | 2.0.0 | diff --git a/README.md b/README.md index ff5ba42..c0eb165 100644 --- a/README.md +++ b/README.md @@ -172,6 +172,53 @@ final sentenceVector = Model2VecUtils.meanPooling(candidates); final base64String = Model2VecUtils.toBase64(query); ``` +### 5. Local Retrieval (RAG) with `EmbeddingIndex` + +Build a searchable, persistable index of your documents entirely on-device — chunk, embed, store, and query without a server. + +```dart +// 1. Split long documents into overlapping passages +final passages = chunkText(document, maxChars: 800, overlap: 100); + +// 2. Embed and index them (int8 storage cuts memory ~4x) +final index = EmbeddingIndex(quantized: true); +for (var i = 0; i < passages.length; i++) { + index.add('passage-$i', Model2Vec.generateEmbedding(passages[i])); +} + +// 3. Query — returns SearchResult(id, score), most similar first +final query = Model2Vec.generateEmbedding('How do I reset my password?'); +final hits = index.search(query, topK: 5); +for (final hit in hits) { + print('${hit.id}: ${hit.score.toStringAsFixed(3)}'); +} + +// 4. Persist to disk and reload later +final bytes = index.toBytes(); +final reloaded = EmbeddingIndex.fromBytes(bytes); + +// (Optional) Rerank a candidate vector list for diverse results (MMR): +final diverse = Model2VecUtils.maximalMarginalRelevance( + query, candidateVectors, topK: 5, lambda: 0.5, +); +``` + +### 6. Parallel Embedding & Lifecycle + +```dart +// Embed across all CPU cores with a pool of worker isolates +final pool = await EmbeddingPool.start(); // defaults to core count +final results = await pool.embedBatches(listOfBatches); +await pool.close(); + +// Non-throwing state check + free native memory when done +if (Model2Vec.isInitialized) { + final info = Model2Vec.modelInfo; // dim, vocab, normalized, median + print('Loaded model with dimension ${info.dimension}'); +} +Model2Vec.unloadModel(); // releases the native model +``` + ## API Reference ### Core Methods (`Model2Vec` class) @@ -187,6 +234,9 @@ final base64String = Model2VecUtils.toBase64(query); | `generateBatchEmbeddings(texts)` | Synchronously generates embeddings for a `List` using Rust SIMD. | | `generateEmbeddingAsync(text)` | Asynchronously generates an embedding in a background `Isolate`. | | `generateEmbeddingStream(stream)` | Processes a huge `Stream` into a `Stream` in batches. | +| `isInitialized` | Non-throwing check for whether a model is currently loaded. | +| `modelInfo` | All model metadata in one `ModelInfo` (dimension, vocabulary, normalized, median). | +| `unloadModel()` | Unloads the active model and frees its native memory. | | `embeddingDimension` | Property returning the vector size (e.g., 256, 384, 512). | | `vocabularySize` | Property returning the number of tokens in the model's vocabulary. | diff --git a/lib/model2vec.dart b/lib/model2vec.dart index 6659d6e..fa410c7 100644 --- a/lib/model2vec.dart +++ b/lib/model2vec.dart @@ -2,7 +2,11 @@ /// library for generating text embeddings. library; +export 'src/chunker.dart' show chunkText; +export 'src/embedding_index.dart' show EmbeddingIndex, SearchResult; +export 'src/embedding_pool.dart' show EmbeddingPool; export 'src/exception.dart' show Model2VecErrorKind, Model2VecException; export 'src/model2vec_base.dart' show Model2Vec; +export 'src/model_info.dart' show ModelInfo; export 'src/recommended_model.dart' show RecommendedModel; export 'src/utils.dart' show Model2VecUtils; diff --git a/lib/src/chunker.dart b/lib/src/chunker.dart new file mode 100644 index 0000000..5368680 --- /dev/null +++ b/lib/src/chunker.dart @@ -0,0 +1,73 @@ +/// Splits [text] into overlapping chunks suitable for embedding and retrieval. +/// +/// Chunks hold at most [maxChars] characters and carry [overlap] characters of +/// trailing context into the next chunk, so a passage that straddles a +/// boundary is still retrievable. Splitting happens on whitespace, so words are +/// never cut mid-way — a single word longer than [maxChars] becomes its own +/// (oversized) chunk rather than being broken. +/// +/// To budget by tokens instead of characters, pass +/// `maxChars: maxTokens * Model2Vec.medianTokenLength`, mirroring how the model +/// itself truncates input. +/// +/// Returns an empty list for blank input. Throws [ArgumentError] if [maxChars] +/// is below 1 or [overlap] is outside `[0, maxChars)`. +List chunkText(String text, {int maxChars = 1000, int overlap = 100}) { + if (maxChars < 1) { + throw ArgumentError.value(maxChars, 'maxChars', 'must be >= 1'); + } + if (overlap < 0 || overlap >= maxChars) { + throw ArgumentError.value(overlap, 'overlap', 'must be in [0, maxChars)'); + } + + final words = text.trim().split(RegExp(r'\s+')).where((w) => w.isNotEmpty); + if (words.isEmpty) { + return []; + } + + final chunks = []; + var current = StringBuffer(); + + for (final word in words) { + if (current.isEmpty) { + current.write(word); + } else if (current.length + 1 + word.length <= maxChars) { + current + ..write(' ') + ..write(word); + } else { + final chunk = current.toString(); + chunks.add(chunk); + current = StringBuffer(); + final tail = _overlapTail(chunk, overlap); + // Only carry the overlap when it still leaves room for this word, so a + // chunk never exceeds maxChars (except a single oversized word). + if (tail.isNotEmpty && tail.length + 1 + word.length <= maxChars) { + current + ..write(tail) + ..write(' '); + } + current.write(word); + } + } + + if (current.isNotEmpty) { + chunks.add(current.toString()); + } + return chunks; +} + +/// The last [overlap] characters of [chunk], advanced to the next word +/// boundary so the tail starts on a whole word. +String _overlapTail(String chunk, int overlap) { + if (overlap <= 0 || chunk.length <= overlap) { + return ''; + } + var start = chunk.length - overlap; + final space = chunk.indexOf(' ', start); + if (space == -1) { + return ''; // the tail would be a single partial word; skip it + } + start = space + 1; + return chunk.substring(start); +} diff --git a/lib/src/embedding_index.dart b/lib/src/embedding_index.dart new file mode 100644 index 0000000..59d883c --- /dev/null +++ b/lib/src/embedding_index.dart @@ -0,0 +1,278 @@ +import 'dart:convert'; +import 'dart:typed_data'; + +import 'utils.dart'; + +/// A single hit from [EmbeddingIndex.search]: the stored entry's id and its +/// similarity to the query. +final class SearchResult { + /// Creates a search hit. + const SearchResult(this.id, this.score); + + /// The id the vector was stored under. + final String id; + + /// Cosine similarity to the query, in `[-1.0, 1.0]`. + final double score; + + @override + String toString() => 'SearchResult($id, ${score.toStringAsFixed(4)})'; +} + +/// An in-memory index of embeddings for local semantic search / retrieval. +/// +/// Store vectors by id with [add], then query the nearest ones with [search]. +/// Set `quantized: true` to keep vectors int8-quantized (~4x less memory, at a +/// small accuracy cost). Persist the whole index with [toBytes] and restore it +/// with [EmbeddingIndex.fromBytes]. +/// +/// This is a pure data structure — it never touches the native model, so it is +/// fully testable with hand-made vectors. +class EmbeddingIndex { + /// Creates an empty index. When [quantized], stored vectors are int8. + EmbeddingIndex({bool quantized = false}) : _quantized = quantized; + + final bool _quantized; + + // id -> stored vector (Float32List, or Int8List when quantized). Insertion + // order is preserved, which keeps [toBytes] output stable. + final _store = {}; + int? _dimension; + + /// Number of stored vectors. + int get length => _store.length; + + /// Whether the index is empty. + bool get isEmpty => _store.isEmpty; + + /// Whether vectors are stored int8-quantized. + bool get isQuantized => _quantized; + + /// Dimension of the stored vectors, or `null` while the index is empty. + int? get dimension => _dimension; + + /// The ids currently stored, in insertion order. + Iterable get ids => _store.keys; + + /// Whether [id] is present. + bool contains(String id) => _store.containsKey(id); + + /// Stores [vector] under [id], replacing any existing entry for that id. + /// + /// Throws [ArgumentError] if [vector]'s length differs from the vectors + /// already in the index. + void add(String id, Float32List vector) { + _checkDimension(vector.length); + _dimension = vector.length; + _store[id] = _quantized + ? Model2VecUtils.quantizeToInt8(vector) + : Float32List.fromList(vector); + } + + /// Stores every entry in [entries]. See [add]. + void addAll(Map entries) { + entries.forEach(add); + } + + /// Removes the entry for [id]. Returns `true` if it was present. + bool remove(String id) { + final removed = _store.remove(id) != null; + if (_store.isEmpty) { + _dimension = null; + } + return removed; + } + + /// Removes all entries. + void clear() { + _store.clear(); + _dimension = null; + } + + /// Returns the [topK] entries most similar to [query], most similar first. + /// + /// Throws [ArgumentError] if [query]'s length differs from the index's + /// [dimension]. Returns an empty list for an empty index. + List search(Float32List query, {int topK = 5}) { + final scored = _scoreAll(query) + ..sort((a, b) => b.score.compareTo(a.score)); + if (topK <= 0) { + return []; + } + if (topK >= scored.length) { + return scored; + } + return scored.sublist(0, topK); + } + + /// Returns every entry whose similarity to [query] is `>= threshold`, most + /// similar first. + List searchWithThreshold( + Float32List query, { + required double threshold, + }) { + final scored = _scoreAll(query) + .where((r) => r.score >= threshold) + .toList() + ..sort((a, b) => b.score.compareTo(a.score)); + return scored; + } + + List _scoreAll(Float32List query) { + if (_store.isEmpty) { + return []; + } + if (query.length != _dimension) { + throw ArgumentError( + 'query length ${query.length} != index dimension $_dimension', + ); + } + final results = []; + for (final entry in _store.entries) { + final vector = _asFloat32(entry.value); + results.add( + SearchResult(entry.key, Model2VecUtils.cosineSimilarity(query, vector)), + ); + } + return results; + } + + Float32List _asFloat32(Object stored) => _quantized + ? Model2VecUtils.dequantizeInt8(stored as Int8List) + : stored as Float32List; + + void _checkDimension(int length) { + if (_dimension != null && length != _dimension) { + throw ArgumentError( + 'vector length $length != index dimension $_dimension', + ); + } + } + + // --- persistence --------------------------------------------------------- + + static const _magic = 'M2VI'; + static const _version = 1; + + /// Serializes the whole index to a compact binary blob. Restore it with + /// [EmbeddingIndex.fromBytes]. + Uint8List toBytes() { + final dim = _dimension ?? 0; + final elemBytes = _quantized ? 1 : 4; + + final idByteList = _store.keys.map(utf8.encode).toList(growable: false); + var size = 4 + 1 + 1 + 4 + 4; // magic + version + flags + dim + count + for (final idBytes in idByteList) { + size += 4 + idBytes.length + dim * elemBytes; // u32 idLen + id + vector + } + + final bytes = Uint8List(size); + final data = ByteData.sublistView(bytes); + var o = 0; + + bytes.setRange(0, 4, ascii.encode(_magic)); + o = 4; + data.setUint8(o, _version); + o += 1; + data.setUint8(o, _quantized ? 1 : 0); + o += 1; + data.setUint32(o, dim, Endian.little); + o += 4; + data.setUint32(o, _store.length, Endian.little); + o += 4; + + var e = 0; + for (final stored in _store.values) { + final idBytes = idByteList[e++]; + data.setUint32(o, idBytes.length, Endian.little); + o += 4; + bytes.setRange(o, o + idBytes.length, idBytes); + o += idBytes.length; + if (_quantized) { + final q = stored as Int8List; + for (var i = 0; i < dim; i++) { + data.setInt8(o, q[i]); + o += 1; + } + } else { + final f = stored as Float32List; + for (var i = 0; i < dim; i++) { + data.setFloat32(o, f[i], Endian.little); + o += 4; + } + } + } + return bytes; + } + + /// Reconstructs an index from a blob produced by [toBytes]. + /// + /// Throws [ArgumentError] if the blob is not a recognized index. + static EmbeddingIndex fromBytes(Uint8List bytes) { + if (bytes.length < 14 || !_hasMagic(bytes)) { + throw ArgumentError('not a valid EmbeddingIndex blob'); + } + try { + return _decode(bytes); + } on FormatException catch (e) { + throw ArgumentError('corrupt EmbeddingIndex blob: $e'); + // A short/corrupt buffer surfaces as a RangeError from the typed reads; + // convert it to the documented ArgumentError for callers. + // ignore: avoid_catching_errors + } on RangeError catch (e) { + throw ArgumentError('corrupt EmbeddingIndex blob: $e'); + } + } + + static bool _hasMagic(Uint8List bytes) { + final magic = ascii.encode(_magic); + for (var i = 0; i < magic.length; i++) { + if (bytes[i] != magic[i]) { + return false; + } + } + return true; + } + + static EmbeddingIndex _decode(Uint8List bytes) { + final data = ByteData.sublistView(bytes); + var o = 4; // past the magic + + final version = data.getUint8(o); + o += 1; + if (version != _version) { + throw ArgumentError('unsupported EmbeddingIndex version $version'); + } + final quantized = data.getUint8(o) == 1; + o += 1; + final dim = data.getUint32(o, Endian.little); + o += 4; + final count = data.getUint32(o, Endian.little); + o += 4; + + final index = EmbeddingIndex(quantized: quantized); + for (var e = 0; e < count; e++) { + final idLen = data.getUint32(o, Endian.little); + o += 4; + final id = utf8.decode(bytes.sublist(o, o + idLen)); + o += idLen; + if (quantized) { + final q = Int8List(dim); + for (var i = 0; i < dim; i++) { + q[i] = data.getInt8(o); + o += 1; + } + index._store[id] = q; + } else { + final f = Float32List(dim); + for (var i = 0; i < dim; i++) { + f[i] = data.getFloat32(o, Endian.little); + o += 4; + } + index._store[id] = f; + } + } + index._dimension = count > 0 ? dim : null; + return index; + } +} diff --git a/lib/src/embedding_pool.dart b/lib/src/embedding_pool.dart new file mode 100644 index 0000000..bf07ea4 --- /dev/null +++ b/lib/src/embedding_pool.dart @@ -0,0 +1,95 @@ +import 'dart:async'; +import 'dart:io'; +import 'dart:isolate'; +import 'dart:typed_data'; + +import 'embedding_worker.dart'; + +/// A pool of worker isolates that embed batches concurrently. +/// +/// The native model is a single process-global behind a read/write lock, and +/// the Rust engine allows concurrent readers — so N worker isolates genuinely +/// parallelize embedding across CPU cores. Use this for large workloads where a +/// single `Model2Vec.generateEmbeddingStream` worker is the bottleneck. +/// +/// A model must be initialized (via `Model2Vec.initEmbedder`) before use, and +/// must not be switched while the pool is working (the one-model-per-run +/// contract, same as the single worker). +class EmbeddingPool { + EmbeddingPool._(this._workers) + : _inFlight = List.filled(_workers.length, 0); + + final List _workers; + final List _inFlight; + + /// Number of worker isolates in the pool. + int get size => _workers.length; + + /// Spawns a pool of [size] workers (defaulting to the number of CPU cores). + /// + /// [entryPoint] defaults to the model-backed worker; tests inject a fake one. + static Future start({ + int? size, + void Function(SendPort) entryPoint = embeddingWorkerEntryPoint, + }) async { + final requested = size ?? Platform.numberOfProcessors; + final count = requested < 1 ? 1 : requested; + // Start all in parallel; if any worker fails to start, close the ones that + // did so we never leak orphaned isolates. + final outcomes = await Future.wait( + List.generate(count, (_) async { + try { + return await EmbeddingWorker.start(entryPoint: entryPoint); + } on Object { + return null; + } + }), + ); + final workers = outcomes.whereType().toList(); + if (workers.length != count) { + await Future.wait(workers.map((worker) => worker.close())); + throw StateError('failed to start $count embedding workers'); + } + return EmbeddingPool._(workers); + } + + /// Embeds [batch] on the least-busy worker. Safe to call concurrently — each + /// call runs on a (potentially) different worker, in parallel. + Future> embedBatch( + List batch, { + int maxLength = 512, + }) { + final worker = _leastBusyIndex(); + _inFlight[worker]++; + return _workers[worker] + .embedBatch(batch, maxLength: maxLength) + .whenComplete(() => _inFlight[worker]--); + } + + /// Embeds every batch in [batches] concurrently across the pool and returns + /// the results in the same order as the input. + /// + /// All batches are dispatched at once; for very large inputs, feed them in + /// windows to bound memory. + Future>> embedBatches( + List> batches, { + int maxLength = 512, + }) => Future.wait( + batches.map((batch) => embedBatch(batch, maxLength: maxLength)), + ); + + /// Tears down every worker in the pool. + Future close() async { + await Future.wait(_workers.map((worker) => worker.close())); + } + + int _leastBusyIndex() { + var best = 0; + for (var i = 1; i < _inFlight.length; i++) { + if (_inFlight[i] < _inFlight[best]) { + best = i; + } + } + return best; + } +} diff --git a/lib/src/model2vec_base.dart b/lib/src/model2vec_base.dart index 0b82667..5035040 100644 --- a/lib/src/model2vec_base.dart +++ b/lib/src/model2vec_base.dart @@ -9,6 +9,7 @@ import 'batcher.dart'; import 'embedding_worker.dart'; import 'exception.dart'; import 'model2vec_bindings.g.dart' as native; +import 'model_info.dart'; import 'recommended_model.dart'; /// The main entry point for the Model2Vec library. @@ -19,6 +20,10 @@ import 'recommended_model.dart'; /// SDK's code-asset resolver — there is nothing to boot or inject. // ignore: avoid_classes_with_only_static_members abstract final class Model2Vec { + /// Whether a model is currently loaded. Unlike the metadata getters, this + /// never throws — use it to guard calls before a model is initialized. + static bool get isInitialized => native.is_model_loaded() == 1; + /// Returns the embedding dimension of the currently loaded model. /// /// Throws a [Model2VecException] if no model has been initialized yet. @@ -41,6 +46,17 @@ abstract final class Model2Vec { static int get medianTokenLength => _readInt(native.get_median_token_length); + /// Reads the current model's metadata into a [ModelInfo]. + /// + /// Reflects the model loaded at call time; do not switch or unload the model + /// concurrently. Throws a [Model2VecException] if no model is loaded. + static ModelInfo get modelInfo => ModelInfo( + dimension: embeddingDimension, + vocabularySize: vocabularySize, + isNormalized: isNormalized, + medianTokenLength: medianTokenLength, + ); + static int _readInt( int Function(Pointer, Pointer>) fn, ) => using((arena) { @@ -150,6 +166,15 @@ abstract final class Model2Vec { ); }); + /// Unloads the active model and frees its native memory. + /// + /// Safe to call when no model is loaded. After this, [isInitialized] is + /// `false` and the metadata getters throw until a model is re-initialized. + static void unloadModel() => using((arena) { + final outError = arena>(); + _check(native.free_embedder(outError), outError); + }); + /// Officially recommended Potion models to start from. /// /// A curated, offline catalog — not fetched from Hugging Face (see diff --git a/lib/src/model2vec_bindings.g.dart b/lib/src/model2vec_bindings.g.dart index 86abc50..ad221d3 100644 --- a/lib/src/model2vec_bindings.g.dart +++ b/lib/src/model2vec_bindings.g.dart @@ -146,3 +146,11 @@ external void free_floats( ffi.Pointer ptr, int len, ); + +@ffi.Native() +external int is_model_loaded(); + +@ffi.Native>)>() +external int free_embedder( + ffi.Pointer> out_error, +); diff --git a/lib/src/model_info.dart b/lib/src/model_info.dart new file mode 100644 index 0000000..5b98f8a --- /dev/null +++ b/lib/src/model_info.dart @@ -0,0 +1,39 @@ +/// A snapshot of the active model's metadata, read in one call. +final class ModelInfo { + /// Creates a metadata snapshot. + const ModelInfo({ + required this.dimension, + required this.vocabularySize, + required this.isNormalized, + required this.medianTokenLength, + }); + + /// Length of the embedding vectors the model produces. + final int dimension; + + /// Number of unique tokens in the model's vocabulary. + final int vocabularySize; + + /// Whether the model L2-normalizes its output embeddings. + final bool isNormalized; + + /// Median token length, in characters, across the vocabulary. + final int medianTokenLength; + + @override + bool operator ==(Object other) => + other is ModelInfo && + other.dimension == dimension && + other.vocabularySize == vocabularySize && + other.isNormalized == isNormalized && + other.medianTokenLength == medianTokenLength; + + @override + int get hashCode => + Object.hash(dimension, vocabularySize, isNormalized, medianTokenLength); + + @override + String toString() => + 'ModelInfo(dimension: $dimension, vocabularySize: $vocabularySize, ' + 'isNormalized: $isNormalized, medianTokenLength: $medianTokenLength)'; +} diff --git a/lib/src/utils.dart b/lib/src/utils.dart index db03415..f43a1e5 100644 --- a/lib/src/utils.dart +++ b/lib/src/utils.dart @@ -126,6 +126,75 @@ final class Model2VecUtils { return similarities.map((s) => s.i).toList(growable: false); } + /// Like [similaritySearch] but returns each candidate index paired with its + /// cosine similarity to [query], most similar first. + static List<({int index, double score})> similaritySearchWithScores( + Float32List query, + List candidates, { + int topK = 5, + }) { + if (candidates.isEmpty || topK <= 0) { + return []; + } + final scored = <({int index, double score})>[ + for (var i = 0; i < candidates.length; i++) + (index: i, score: cosineSimilarity(query, candidates[i])), + ]..sort((a, b) => b.score.compareTo(a.score)); + return scored.take(math.min(topK, scored.length)).toList(growable: false); + } + + /// Reranks [candidates] against [query] with Maximal Marginal Relevance, + /// trading off relevance (similarity to the query) against diversity + /// (dissimilarity to already-picked results). + /// + /// [lambda] in `[0.0, 1.0]`: 1.0 is pure relevance, 0.0 is pure diversity. + /// Returns the selected candidate indices, best first. + static List maximalMarginalRelevance( + Float32List query, + List candidates, { + int topK = 5, + double lambda = 0.5, + }) { + if (lambda < 0.0 || lambda > 1.0) { + throw ArgumentError.value(lambda, 'lambda', 'must be in [0.0, 1.0]'); + } + if (candidates.isEmpty) { + return []; + } + final k = math.min(topK, candidates.length); + final relevance = [for (final c in candidates) cosineSimilarity(query, c)]; + // Running max similarity of each candidate to any already-selected item, + // updated once per selection — keeps the whole reranking O(k * n * dim). + final maxSimToSelected = List.filled( + candidates.length, + double.negativeInfinity, + ); + final selected = []; + final remaining = List.generate(candidates.length, (i) => i); + + while (selected.length < k && remaining.isNotEmpty) { + var bestIdx = remaining.first; + var bestScore = double.negativeInfinity; + for (final i in remaining) { + final penalty = selected.isEmpty ? 0.0 : maxSimToSelected[i]; + final mmr = lambda * relevance[i] - (1 - lambda) * penalty; + if (mmr > bestScore) { + bestScore = mmr; + bestIdx = i; + } + } + selected.add(bestIdx); + remaining.remove(bestIdx); + for (final i in remaining) { + final sim = cosineSimilarity(candidates[i], candidates[bestIdx]); + if (sim > maxSimToSelected[i]) { + maxSimToSelected[i] = sim; + } + } + } + return selected; + } + /// Calculates the cosine distance between two vectors. /// Distance is `1.0 - cosineSimilarity`. Range is 0.0 to 2.0. static double cosineDistance(Float32List a, Float32List b) => @@ -191,6 +260,17 @@ final class Model2VecUtils { return result; } + /// Reconstructs an approximate [Float32List] from an [Int8List] produced by + /// [quantizeToInt8] — the inverse scaling (value / 127). Quantization is + /// lossy, so the result is close to, but not exactly, the original vector. + static Float32List dequantizeInt8(Int8List quantized) { + final result = Float32List(quantized.length); + for (var i = 0; i < quantized.length; i++) { + result[i] = quantized[i] / 127.0; + } + return result; + } + /// Serializes a [Float32List] to a Base64 string for easy storage. static String toBase64(Float32List vector) => base64Encode(vector.buffer.asUint8List()); diff --git a/native/model2vec.h b/native/model2vec.h index 765aec5..91b05cb 100644 --- a/native/model2vec.h +++ b/native/model2vec.h @@ -74,3 +74,7 @@ int generate_batch_embeddings_advanced( void free_string(char* s); void free_floats(float* ptr, size_t len); + +int is_model_loaded(); + +int free_embedder(char** out_error); diff --git a/native/src/lib.rs b/native/src/lib.rs index 89bca8f..38fb4df 100644 --- a/native/src/lib.rs +++ b/native/src/lib.rs @@ -357,3 +357,26 @@ pub extern "C" fn free_floats(ptr: *mut f32, len: usize) { let _ = Vec::from_raw_parts(ptr, len, len); } } + +/// Returns 1 if a model is currently loaded, 0 otherwise. Never fails: a +/// poisoned lock is reported as "not loaded". +#[no_mangle] +pub extern "C" fn is_model_loaded() -> i32 { + match MODEL.read() { + Ok(lock) => i32::from(lock.is_some()), + Err(_) => 0, + } +} + +/// Unloads the active model, freeing its memory. Idempotent (no-op when nothing +/// is loaded). +#[no_mangle] +pub extern "C" fn free_embedder(out_error: *mut *mut c_char) -> i32 { + run_ffi(out_error, || { + let mut lock = MODEL + .write() + .map_err(|_| FfiError::new(CODE_LOCK_POISONED, "model lock poisoned"))?; + *lock = None; + Ok(()) + }) +} diff --git a/test/chunker_test.dart b/test/chunker_test.dart new file mode 100644 index 0000000..b79db30 --- /dev/null +++ b/test/chunker_test.dart @@ -0,0 +1,57 @@ +import 'package:model2vec/model2vec.dart'; +import 'package:test/test.dart'; + +void main() { + group('chunkText', () { + test('short text is a single chunk', () { + expect(chunkText('hello world'), ['hello world']); + }); + + test('blank input yields no chunks', () { + expect(chunkText(' \n\t '), isEmpty); + }); + + test('splits long text and covers every word exactly once (no overlap)', + () { + final words = List.generate(40, (i) => 'token$i'); + final chunks = chunkText(words.join(' '), maxChars: 40, overlap: 0); + + expect(chunks.length, greaterThan(1)); + final emitted = chunks.expand((c) => c.split(' ')).toList(); + expect(emitted, equals(words)); // order preserved, no duplication + }); + + test('overlap re-includes boundary words', () { + final words = List.generate(40, (i) => 'token$i').join(' '); + final withOverlap = chunkText(words, maxChars: 40, overlap: 14); + final without = chunkText(words, maxChars: 40, overlap: 0); + + expect(withOverlap.length, greaterThan(1)); + final totalWith = withOverlap.expand((c) => c.split(' ')).length; + final totalWithout = without.expand((c) => c.split(' ')).length; + expect(totalWith, greaterThan(totalWithout)); + }); + + test('chunks stay within maxChars except a single oversized word', () { + const text = 'aaaa bbbbbbbbbb cccccccccccc dddddddddd eeeeeeee ffffff'; + final chunks = chunkText(text, maxChars: 30, overlap: 25); + for (final chunk in chunks) { + if (chunk.split(' ').length > 1) { + expect(chunk.length, lessThanOrEqualTo(30)); + } + } + }); + + test('validates arguments', () { + expect(() => chunkText('a', maxChars: 0), throwsArgumentError); + expect( + () => chunkText('a', maxChars: 10, overlap: 10), + throwsArgumentError, + ); + expect( + () => chunkText('a', maxChars: 10, overlap: -1), + throwsArgumentError, + ); + }); + }); +} diff --git a/test/embedding_index_test.dart b/test/embedding_index_test.dart new file mode 100644 index 0000000..ca0dc86 --- /dev/null +++ b/test/embedding_index_test.dart @@ -0,0 +1,151 @@ +import 'dart:typed_data'; + +import 'package:model2vec/model2vec.dart'; +import 'package:test/test.dart'; + +Float32List _v(List xs) => Float32List.fromList(xs); + +void main() { + group('EmbeddingIndex', () { + test('add / length / contains / remove / clear', () { + final index = EmbeddingIndex() + ..add('a', _v([1, 0, 0])) + ..add('b', _v([0, 1, 0])); + + expect(index.length, 2); + expect(index.contains('a'), isTrue); + expect(index.contains('z'), isFalse); + expect(index.dimension, 3); + + expect(index.remove('a'), isTrue); + expect(index.remove('a'), isFalse); + expect(index.length, 1); + + index.clear(); + expect(index.isEmpty, isTrue); + expect(index.dimension, isNull); + }); + + test('search ranks by cosine similarity, most similar first', () { + final index = EmbeddingIndex() + ..add('a', _v([1, 0, 0])) + ..add('b', _v([0, 1, 0])) + ..add('near-a', _v([0.9, 0.1, 0])); + + final hits = index.search(_v([1, 0, 0]), topK: 2); + expect(hits.map((h) => h.id).toList(), ['a', 'near-a']); + expect(hits.first.score, closeTo(1.0, 1e-6)); + expect(hits[1].score, greaterThan(0.9)); + }); + + test('topK larger than size returns all', () { + final index = EmbeddingIndex()..add('a', _v([1, 0])); + expect(index.search(_v([1, 0]), topK: 10).length, 1); + }); + + test('searchWithThreshold filters and sorts', () { + final index = EmbeddingIndex() + ..add('a', _v([1, 0, 0])) + ..add('b', _v([0, 1, 0])) + ..add('near-a', _v([0.9, 0.1, 0])); + + final hits = index.searchWithThreshold(_v([1, 0, 0]), threshold: 0.5); + expect(hits.map((h) => h.id).toList(), ['a', 'near-a']); + }); + + test('empty index search returns empty', () { + expect(EmbeddingIndex().search(_v([1, 0])), isEmpty); + }); + + test('addAll and overwrite by id', () { + final index = EmbeddingIndex() + ..addAll({'a': _v([1, 0]), 'b': _v([0, 1])}); + expect(index.length, 2); + + index.add('a', _v([0, 1])); // overwrite + expect(index.length, 2); + expect(index.search(_v([0, 1]), topK: 1).first.id, anyOf('a', 'b')); + }); + + test('dimension mismatch throws on add and on query', () { + final index = EmbeddingIndex()..add('a', _v([1, 0, 0])); + expect(() => index.add('b', _v([1, 0])), throwsArgumentError); + expect(() => index.search(_v([1, 0])), throwsArgumentError); + }); + + test('toBytes/fromBytes round-trips exactly (float)', () { + final index = EmbeddingIndex() + ..add('doc-1', _v([0.1, 0.2, 0.3, 0.4])) + ..add('doc-2', _v([-0.5, 0.5, 0.0, 1.0])); + + final restored = EmbeddingIndex.fromBytes(index.toBytes()); + expect(restored.length, 2); + expect(restored.isQuantized, isFalse); + expect(restored.dimension, 4); + + final q = _v([0.1, 0.2, 0.3, 0.4]); + final before = index.search(q); + final after = restored.search(q); + expect(after.map((h) => h.id).toList(), before.map((h) => h.id).toList()); + for (var i = 0; i < before.length; i++) { + expect(after[i].score, closeTo(before[i].score, 1e-6)); + } + }); + + test('quantized storage round-trips and searches', () { + final index = EmbeddingIndex(quantized: true) + ..add('a', _v([1, 0, 0])) + ..add('b', _v([0, 1, 0])) + ..add('near-a', _v([0.9, 0.1, 0])); + + expect(index.isQuantized, isTrue); + final hits = index.search(_v([1, 0, 0]), topK: 2); + expect(hits.map((h) => h.id).toList(), ['a', 'near-a']); + + final restored = EmbeddingIndex.fromBytes(index.toBytes()); + expect(restored.isQuantized, isTrue); + final after = restored.search(_v([1, 0, 0]), topK: 2); + expect(after.map((h) => h.id).toList(), ['a', 'near-a']); + expect(after.first.score, closeTo(hits.first.score, 1e-6)); + }); + + test('fromBytes rejects a non-index blob', () { + expect( + () => EmbeddingIndex.fromBytes(Uint8List.fromList([1, 2, 3])), + throwsArgumentError, + ); + }); + + test('empty index round-trips', () { + final restored = EmbeddingIndex.fromBytes(EmbeddingIndex().toBytes()); + expect(restored.isEmpty, isTrue); + expect(restored.dimension, isNull); + }); + + test('a very long id (> 65535 bytes) round-trips', () { + final longId = 'x' * 70000; + final index = EmbeddingIndex()..add(longId, _v([1, 0, 0])); + final restored = EmbeddingIndex.fromBytes(index.toBytes()); + expect(restored.length, 1); + expect(restored.contains(longId), isTrue); + }); + + test('fromBytes rejects a high-byte / truncated blob with ArgumentError', + () { + expect( + () => EmbeddingIndex.fromBytes(Uint8List(14)..[0] = 0xFF), + throwsArgumentError, + ); + final full = (EmbeddingIndex()..add('a', _v([1, 2, 3, 4]))).toBytes(); + expect( + () => EmbeddingIndex.fromBytes(full.sublist(0, full.length - 2)), + throwsArgumentError, + ); + }); + + test('negative topK returns empty', () { + final index = EmbeddingIndex()..add('a', _v([1, 0])); + expect(index.search(_v([1, 0]), topK: -1), isEmpty); + }); + }); +} diff --git a/test/embedding_pool_test.dart b/test/embedding_pool_test.dart new file mode 100644 index 0000000..6f8b326 --- /dev/null +++ b/test/embedding_pool_test.dart @@ -0,0 +1,70 @@ +import 'dart:isolate'; +import 'dart:typed_data'; + +import 'package:model2vec/model2vec.dart'; +import 'package:model2vec/src/worker_protocol.dart'; +import 'package:test/test.dart'; + +/// Fake worker entry point: one vector per input text, whose only element is +/// the text's length. No model, no native calls. +void echoEntryPoint(SendPort mainSendPort) { + final receivePort = ReceivePort(); + mainSendPort.send(receivePort.sendPort); + receivePort.listen((message) { + if (message == 'close') { + receivePort.close(); + return; + } + final request = message as EmbedRequest; + mainSendPort.send([ + for (final text in request.batch) + Float32List.fromList([text.length.toDouble()]), + ]); + }); +} + +void main() { + test('starts with the requested number of workers', () async { + final pool = await EmbeddingPool.start(size: 3, entryPoint: echoEntryPoint); + expect(pool.size, 3); + await pool.close(); + }); + + test('clamps a non-positive size to at least one worker', () async { + final pool = await EmbeddingPool.start(size: 0, entryPoint: echoEntryPoint); + expect(pool.size, 1); + await pool.close(); + }); + + test('embedBatch returns results', () async { + final pool = await EmbeddingPool.start(size: 2, entryPoint: echoEntryPoint); + final result = await pool.embedBatch(['a', 'bb']); + expect(result.map((v) => v.first).toList(), [1.0, 2.0]); + await pool.close(); + }); + + test('embedBatches preserves input order under parallelism', () async { + final pool = await EmbeddingPool.start(size: 3, entryPoint: echoEntryPoint); + final results = await pool.embedBatches([ + ['a'], + ['bb'], + ['ccc'], + ['dddd'], + ['eeeee'], + ]); + expect( + results.map((r) => r.first.first).toList(), + [1.0, 2.0, 3.0, 4.0, 5.0], + ); + await pool.close(); + }); + + test('close tears down all workers', () async { + final pool = await EmbeddingPool.start(size: 2, entryPoint: echoEntryPoint); + await pool.close(); + await expectLater( + pool.embedBatch(['a']), + throwsA(isA()), + ); + }); +} diff --git a/test/lifecycle_test.dart b/test/lifecycle_test.dart new file mode 100644 index 0000000..fe133fc --- /dev/null +++ b/test/lifecycle_test.dart @@ -0,0 +1,37 @@ +import 'package:model2vec/model2vec.dart'; +import 'package:test/test.dart'; + +void main() { + group('model lifecycle', () { + test('isInitialized / modelInfo / unloadModel round-trip', () { + Model2Vec.initEmbedder('minishlab/potion-base-2M'); + expect(Model2Vec.isInitialized, isTrue); + + final info = Model2Vec.modelInfo; + expect(info.dimension, equals(64)); + expect(info.vocabularySize, greaterThan(20000)); + expect(info.isNormalized, isTrue); + expect(info.medianTokenLength, isPositive); + + Model2Vec.unloadModel(); + expect(Model2Vec.isInitialized, isFalse); + expect( + () => Model2Vec.embeddingDimension, + throwsA( + isA().having( + (e) => e.kind, + 'kind', + Model2VecErrorKind.notInitialized, + ), + ), + ); + + // Idempotent: unloading again is a no-op. + expect(Model2Vec.unloadModel, returnsNormally); + + // Re-initialize so later suites in this process find a model. + Model2Vec.initEmbedder('minishlab/potion-base-2M'); + expect(Model2Vec.isInitialized, isTrue); + }); + }); +} diff --git a/test/utils_test.dart b/test/utils_test.dart index 426d393..a406302 100644 --- a/test/utils_test.dart +++ b/test/utils_test.dart @@ -125,6 +125,100 @@ void main() { expect(quantized[3], equals(-127)); }); + test('dequantizeInt8 approximately inverts quantizeToInt8', () { + final vector = Float32List.fromList([0.5, -0.25, 1.0, -1.0, 0.0]); + final restored = Model2VecUtils.dequantizeInt8( + Model2VecUtils.quantizeToInt8(vector), + ); + + expect(restored.length, equals(vector.length)); + for (var i = 0; i < vector.length; i++) { + expect(restored[i], closeTo(vector[i], 0.01)); + } + }); + + test('similaritySearchWithScores returns sorted index+score pairs', () { + final query = Float32List.fromList([1, 0, 0]); + final candidates = [ + Float32List.fromList([1, 0.2, 0]), + Float32List.fromList([1, 0.25, 0]), + Float32List.fromList([0.6, 0, 0.8]), + ]; + + final results = Model2VecUtils.similaritySearchWithScores( + query, + candidates, + topK: 2, + ); + expect(results.map((r) => r.index).toList(), [0, 1]); + expect(results.first.score, greaterThan(results[1].score)); + expect( + Model2VecUtils.similaritySearchWithScores(query, const []), + isEmpty, + ); + }); + + test('MMR prefers a diverse result over a near-duplicate', () { + final query = Float32List.fromList([1, 0, 0]); + final candidates = [ + Float32List.fromList([1, 0.2, 0]), // 0: most relevant + Float32List.fromList([1, 0.25, 0]), // 1: near-duplicate of 0 + Float32List.fromList([0.6, 0, 0.8]), // 2: relevant-ish but diverse + ]; + + final selected = Model2VecUtils.maximalMarginalRelevance( + query, + candidates, + topK: 2, + ); + expect(selected.first, 0); // most relevant first + expect(selected[1], 2); // diversity beats the near-duplicate + }); + + test('MMR with lambda 1.0 is pure relevance order', () { + final query = Float32List.fromList([1, 0, 0]); + final candidates = [ + Float32List.fromList([0.6, 0, 0.8]), + Float32List.fromList([1, 0.2, 0]), + Float32List.fromList([1, 0.25, 0]), + ]; + + final selected = Model2VecUtils.maximalMarginalRelevance( + query, + candidates, + topK: 3, + lambda: 1, + ); + expect(selected, [1, 2, 0]); + expect( + Model2VecUtils.maximalMarginalRelevance(query, const []), + isEmpty, + ); + }); + + test('MMR rejects lambda outside [0, 1]', () { + final query = Float32List.fromList([1, 0]); + final candidates = [ + Float32List.fromList([1, 0]), + ]; + expect( + () => Model2VecUtils.maximalMarginalRelevance( + query, + candidates, + lambda: 1.5, + ), + throwsArgumentError, + ); + expect( + () => Model2VecUtils.maximalMarginalRelevance( + query, + candidates, + lambda: -0.1, + ), + throwsArgumentError, + ); + }); + test('Base64 serialization works bidirectionally', () { final vector = Float32List.fromList([0.1, 0.2, -0.3, 100.5]); final base64 = Model2VecUtils.toBase64(vector); From a869705792d535ba94ccf0fb4d01a24f93031698 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Mon, 6 Jul 2026 19:14:26 +0300 Subject: [PATCH 03/20] add: rag example --- example/rag_example.dart | 109 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 109 insertions(+) create mode 100644 example/rag_example.dart diff --git a/example/rag_example.dart b/example/rag_example.dart new file mode 100644 index 0000000..8e8af06 --- /dev/null +++ b/example/rag_example.dart @@ -0,0 +1,109 @@ +import 'dart:io'; + +import 'package:model2vec/model2vec.dart'; + +/// A tiny knowledge base: article id -> body text. +const _knowledgeBase = { + 'password': + 'To reset your password, open Settings, tap Security, and choose Reset ' + 'Password. We email you a secure link that expires in one hour. If the ' + 'email does not arrive, check your spam folder and confirm the address ' + 'on your account is correct.', + 'twoFactor': + 'Two-factor authentication adds a second step at sign-in. After your ' + 'password you enter a short code from an authenticator app or an SMS ' + 'message, so a stolen password alone is not enough to log in.', + 'refunds': + 'Our refund policy allows returns within 30 days of purchase for a full ' + 'refund. Items must be unused and in the original packaging. Refunds are ' + 'issued to the original payment method within five business days.', + 'shipping': + 'Shipping is free on orders over fifty dollars. Standard delivery ' + 'usually arrives in three to five business days; express delivery ' + 'arrives the next day for an extra fee.', + 'billing': + 'To cancel your subscription, open Billing and choose Cancel Plan before ' + 'the renewal date. You keep access until the end of the current billing ' + 'period and are not charged again.', +}; + +Future main() async { + try { + stdout.writeln('=== Model2Vec RAG example ===\n'); + + // 1. Load an embedding model. For real retrieval, potion-retrieval-32M is + // purpose-built; potion-base-8M is a lighter, still-capable choice. + Model2Vec.initEmbedder('minishlab/potion-base-8M'); + + // 2. Chunk each article and index the passages. The index stores id -> + // vector; we keep our own id -> text map to show the passage behind a + // hit (the index deliberately does not store the text). + final index = EmbeddingIndex(); + final passages = {}; + + for (final article in _knowledgeBase.entries) { + final chunks = chunkText(article.value, maxChars: 160, overlap: 32); + for (var i = 0; i < chunks.length; i++) { + final id = '${article.key}#$i'; + passages[id] = chunks[i]; + index.add(id, Model2Vec.generateEmbedding(chunks[i])); + } + } + stdout.writeln( + 'Indexed ${index.length} passages from ' + '${_knowledgeBase.length} articles.\n', + ); + + // 3. Answer questions by retrieving the closest passages. + const questions = [ + 'How do I change my login password?', + 'Can I get my money back?', + 'When will my package arrive?', + ]; + + for (final question in questions) { + final query = Model2Vec.generateEmbedding(question); + stdout.writeln('Q: $question'); + for (final hit in index.search(query, topK: 2)) { + final score = hit.score.toStringAsFixed(3); + stdout.writeln(' [$score] ${hit.id}: ${passages[hit.id]}'); + } + stdout.writeln(); + } + + // 4. Threshold search — keep only confident matches. + final twoFa = Model2Vec.generateEmbedding('two step verification code'); + final strong = index.searchWithThreshold(twoFa, threshold: 0.4); + stdout.writeln( + 'Confident matches for "two step verification code": ' + '${strong.map((h) => h.id).toList()}\n', + ); + + // 5. Diverse reranking (MMR): pull a wider candidate set, then rerank to + // avoid near-duplicate passages in the top results. + final broad = index.search(twoFa, topK: 6); + final candidateVectors = [ + for (final hit in broad) Model2Vec.generateEmbedding(passages[hit.id]!), + ]; + final order = Model2VecUtils.maximalMarginalRelevance( + twoFa, + candidateVectors, + topK: 3, + lambda: 0.6, + ); + final reranked = [for (final i in order) broad[i].id]; + stdout.writeln('MMR-reranked top 3: $reranked\n'); + + // 6. Persist the index to disk and reload it. Loading needs no model — + // only generating new query vectors does. + final file = File('${Directory.systemTemp.path}/m2v_rag_index.bin'); + await file.writeAsBytes(index.toBytes()); + final reloaded = EmbeddingIndex.fromBytes(await file.readAsBytes()); + stdout.writeln('Saved and reloaded index: ${reloaded.length} passages.'); + await file.delete(); + + stdout.writeln('\n🎉 Done.'); + } on Model2VecException catch (e) { + stdout.writeln('❌ Model2Vec error (${e.kind.name}): ${e.message}'); + } +} From 84e75ad017ea21e48ea7905fd733e2f2a7b85d8f Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Mon, 6 Jul 2026 22:07:28 +0300 Subject: [PATCH 04/20] feat: add async model loading (initEmbedderAsync) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Loading is the heaviest, most download-prone step and had no async door, while microsecond generate* calls did — inverted from cost. On Flutter the synchronous initEmbedder freezes the UI during the first model download. Add initEmbedderAsync and initEmbedderAdvancedAsync, which load on a background isolate via Isolate.run. The native model is a process-global (the worker pool already relies on this), so the model loaded off-thread is visible on every isolate, including the one that awaited the call. --- README.md | 7 ++++++ lib/src/model2vec_base.dart | 32 ++++++++++++++++++++++++++ test/init_async_test.dart | 46 +++++++++++++++++++++++++++++++++++++ 3 files changed, 85 insertions(+) create mode 100644 test/init_async_test.dart diff --git a/README.md b/README.md index c0eb165..1799e36 100644 --- a/README.md +++ b/README.md @@ -138,7 +138,12 @@ Future processHugeFile() async { Never block the main thread. If you are building a Flutter app, always use the `Async` variants to perform generation in a background `Isolate`. +Loading is the heaviest step of all — the first time a model is fetched it downloads tens to hundreds of MB. `initEmbedderAsync` runs that load on a background isolate; because the native model is a single process-global, the loaded model is then visible to every isolate. + ```dart +// Load off the main thread so the first download never freezes the UI. +await Model2Vec.initEmbedderAsync('minishlab/potion-base-2M'); + final embedding = await Model2Vec.generateEmbeddingAsync('A very long text...'); final batch = await Model2Vec.generateBatchEmbeddingsAsync(['A', 'B', 'C']); ``` @@ -228,6 +233,8 @@ Model2Vec.unloadModel(); // releases the native model | `initEmbedder(path)` | Initializes the model from a Hugging Face repo ID or local path. | | `initEmbedderAdvanced(...)` | Advanced initialization with custom `cacheDirectory`, `hfToken`, or `normalize` overrides. | | `initEmbedderFromBytes(...)` | Initializes the model directly from raw `Uint8List` bytes (`model.safetensors`, `tokenizer.json`, etc). | +| `initEmbedderAsync(path)` | Loads a model on a background isolate; await it so the first download never blocks the UI. | +| `initEmbedderAdvancedAsync(...)` | Async form of `initEmbedderAdvanced` (background isolate). | | `recommendedModels` | A typed `List` catalog of officially recommended Potion models (offline). | | `tokenize(text)` | Runs the internal BPE tokenizer and returns a `List`. | | `generateEmbedding(text)` | Synchronously generates a `Float32List` embedding vector. | diff --git a/lib/src/model2vec_base.dart b/lib/src/model2vec_base.dart index 5035040..58137c4 100644 --- a/lib/src/model2vec_base.dart +++ b/lib/src/model2vec_base.dart @@ -166,6 +166,38 @@ abstract final class Model2Vec { ); }); + /// Loads a model asynchronously on a background isolate. + /// + /// Prefer this over [initEmbedder] when the model may be downloaded from + /// Hugging Face for the first time (tens to hundreds of MB): the synchronous + /// [initEmbedder] blocks the calling isolate for the entire download, which + /// freezes a Flutter UI. The native model is a single process-global, so once + /// the background isolate has loaded it the model is visible to every isolate + /// — including the one that awaited this call. + static Future initEmbedderAsync(String modelPath) => + Isolate.run(() => initEmbedder(modelPath)); + + /// Advanced counterpart to [initEmbedderAsync]. + /// + /// Takes the same options as [initEmbedderAdvanced] and loads them on a + /// background isolate; see [initEmbedderAsync] for why loading off-thread + /// matters and why the loaded model is visible on every isolate. + static Future initEmbedderAdvancedAsync({ + required String modelPath, + String? hfToken, + String? cacheDirectory, + bool? normalize, + String? subfolder, + }) => Isolate.run( + () => initEmbedderAdvanced( + modelPath: modelPath, + hfToken: hfToken, + cacheDirectory: cacheDirectory, + normalize: normalize, + subfolder: subfolder, + ), + ); + /// Unloads the active model and frees its native memory. /// /// Safe to call when no model is loaded. After this, [isInitialized] is diff --git a/test/init_async_test.dart b/test/init_async_test.dart new file mode 100644 index 0000000..28d5650 --- /dev/null +++ b/test/init_async_test.dart @@ -0,0 +1,46 @@ +import 'dart:io'; + +import 'package:model2vec/model2vec.dart'; +import 'package:test/test.dart'; + +void main() { + group('async model loading', () { + test( + 'initEmbedderAsync loads the process-global model, visible on the ' + 'calling isolate', + () async { + await Model2Vec.initEmbedderAsync('minishlab/potion-base-2M'); + + // The load ran on a background isolate; the native model is a + // process-global, so a synchronous call on this isolate sees it. + expect(Model2Vec.isInitialized, isTrue); + expect(Model2Vec.embeddingDimension, equals(64)); + expect( + Model2Vec.generateEmbedding('async load').length, + equals(64), + ); + }, + ); + + test('initEmbedderAdvancedAsync honors a custom cache directory', () async { + final tempDir = Directory.systemTemp.createTempSync('m2v_async_cache_'); + try { + await Model2Vec.initEmbedderAdvancedAsync( + modelPath: 'minishlab/potion-base-2M', + cacheDirectory: tempDir.path, + ); + expect(Model2Vec.isInitialized, isTrue); + expect(Model2Vec.embeddingDimension, equals(64)); + } finally { + tempDir.deleteSync(recursive: true); + } + }); + + test('a failed async load surfaces a typed Model2VecException', () { + expect( + Model2Vec.initEmbedderAsync('definitely/not-a-real-model-xyz'), + throwsA(isA()), + ); + }); + }); +} From 5f35bc0688a51ac753660bd0ff033500ad7543e9 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Mon, 6 Jul 2026 22:13:22 +0300 Subject: [PATCH 05/20] feat: let EmbeddingIndex carry an optional payload MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The index stored id->vector only, so every retrieval caller kept a parallel id->text map and re-joined it on each hit (the RAG example and README recipe both did exactly this). Add an optional String payload per entry: add(id, vector, {payload}), surfaced as SearchResult.payload and persisted through toBytes. The blob format bumps to v2 (a per-entry payload section with a null sentinel distinct from an empty string); v1 blobs still decode. payload is a String (not a generic) so the index stays a pure, serializable data structure — hold JSON for structured metadata. The RAG example drops its parallel passages map. --- README.md | 13 ++-- example/rag_example.dart | 18 +++--- lib/src/embedding_index.dart | 108 ++++++++++++++++++++++++--------- test/embedding_index_test.dart | 70 +++++++++++++++++++++ 4 files changed, 170 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index 1799e36..159e5a6 100644 --- a/README.md +++ b/README.md @@ -185,17 +185,22 @@ Build a searchable, persistable index of your documents entirely on-device — c // 1. Split long documents into overlapping passages final passages = chunkText(document, maxChars: 800, overlap: 100); -// 2. Embed and index them (int8 storage cuts memory ~4x) +// 2. Embed and index them, keeping each passage as the entry's payload so a +// hit carries its text directly (int8 storage cuts memory ~4x) final index = EmbeddingIndex(quantized: true); for (var i = 0; i < passages.length; i++) { - index.add('passage-$i', Model2Vec.generateEmbedding(passages[i])); + index.add( + 'passage-$i', + Model2Vec.generateEmbedding(passages[i]), + payload: passages[i], + ); } -// 3. Query — returns SearchResult(id, score), most similar first +// 3. Query — returns SearchResult(id, score, payload), most similar first final query = Model2Vec.generateEmbedding('How do I reset my password?'); final hits = index.search(query, topK: 5); for (final hit in hits) { - print('${hit.id}: ${hit.score.toStringAsFixed(3)}'); + print('${hit.score.toStringAsFixed(3)} ${hit.payload}'); } // 4. Persist to disk and reload later diff --git a/example/rag_example.dart b/example/rag_example.dart index 8e8af06..0b908f7 100644 --- a/example/rag_example.dart +++ b/example/rag_example.dart @@ -35,18 +35,20 @@ Future main() async { // purpose-built; potion-base-8M is a lighter, still-capable choice. Model2Vec.initEmbedder('minishlab/potion-base-8M'); - // 2. Chunk each article and index the passages. The index stores id -> - // vector; we keep our own id -> text map to show the passage behind a - // hit (the index deliberately does not store the text). + // 2. Chunk each article and index the passages. We store the passage text + // as the entry's payload, so a hit carries its document directly — no + // parallel id -> text map to keep in sync. final index = EmbeddingIndex(); - final passages = {}; for (final article in _knowledgeBase.entries) { final chunks = chunkText(article.value, maxChars: 160, overlap: 32); for (var i = 0; i < chunks.length; i++) { final id = '${article.key}#$i'; - passages[id] = chunks[i]; - index.add(id, Model2Vec.generateEmbedding(chunks[i])); + index.add( + id, + Model2Vec.generateEmbedding(chunks[i]), + payload: chunks[i], + ); } } stdout.writeln( @@ -66,7 +68,7 @@ Future main() async { stdout.writeln('Q: $question'); for (final hit in index.search(query, topK: 2)) { final score = hit.score.toStringAsFixed(3); - stdout.writeln(' [$score] ${hit.id}: ${passages[hit.id]}'); + stdout.writeln(' [$score] ${hit.id}: ${hit.payload}'); } stdout.writeln(); } @@ -83,7 +85,7 @@ Future main() async { // avoid near-duplicate passages in the top results. final broad = index.search(twoFa, topK: 6); final candidateVectors = [ - for (final hit in broad) Model2Vec.generateEmbedding(passages[hit.id]!), + for (final hit in broad) Model2Vec.generateEmbedding(hit.payload!), ]; final order = Model2VecUtils.maximalMarginalRelevance( twoFa, diff --git a/lib/src/embedding_index.dart b/lib/src/embedding_index.dart index 59d883c..8a98a58 100644 --- a/lib/src/embedding_index.dart +++ b/lib/src/embedding_index.dart @@ -3,11 +3,11 @@ import 'dart:typed_data'; import 'utils.dart'; -/// A single hit from [EmbeddingIndex.search]: the stored entry's id and its -/// similarity to the query. +/// A single hit from [EmbeddingIndex.search]: the stored entry's id, its +/// similarity to the query, and the payload stored alongside it (if any). final class SearchResult { /// Creates a search hit. - const SearchResult(this.id, this.score); + const SearchResult(this.id, this.score, {this.payload}); /// The id the vector was stored under. final String id; @@ -15,16 +15,28 @@ final class SearchResult { /// Cosine similarity to the query, in `[-1.0, 1.0]`. final double score; + /// The document stored with the vector via `add(..., payload: ...)`, or + /// `null` when the entry was added without one. + final String? payload; + @override String toString() => 'SearchResult($id, ${score.toStringAsFixed(4)})'; } +/// One stored entry: the vector (a `Float32List`, or an `Int8List` when the +/// index is quantized) and its optional payload. +typedef _Entry = ({Object vector, String? payload}); + /// An in-memory index of embeddings for local semantic search / retrieval. /// /// Store vectors by id with [add], then query the nearest ones with [search]. +/// Attach an optional [String] payload to each entry — e.g. the source passage +/// behind the vector — so [search] hands back the document directly instead of +/// just its id, with no parallel id→text map to keep in sync. +/// /// Set `quantized: true` to keep vectors int8-quantized (~4x less memory, at a -/// small accuracy cost). Persist the whole index with [toBytes] and restore it -/// with [EmbeddingIndex.fromBytes]. +/// small accuracy cost). Persist the whole index (payloads included) with +/// [toBytes] and restore it with [EmbeddingIndex.fromBytes]. /// /// This is a pure data structure — it never touches the native model, so it is /// fully testable with hand-made vectors. @@ -34,9 +46,9 @@ class EmbeddingIndex { final bool _quantized; - // id -> stored vector (Float32List, or Int8List when quantized). Insertion - // order is preserved, which keeps [toBytes] output stable. - final _store = {}; + // id -> stored entry (vector + optional payload). Insertion order is + // preserved, which keeps [toBytes] output stable. + final _store = {}; int? _dimension; /// Number of stored vectors. @@ -59,14 +71,18 @@ class EmbeddingIndex { /// Stores [vector] under [id], replacing any existing entry for that id. /// - /// Throws [ArgumentError] if [vector]'s length differs from the vectors - /// already in the index. - void add(String id, Float32List vector) { + /// Optionally attach a [payload] (e.g. the source text) that [search] will + /// return with the hit. Throws [ArgumentError] if [vector]'s length differs + /// from the vectors already in the index. + void add(String id, Float32List vector, {String? payload}) { _checkDimension(vector.length); _dimension = vector.length; - _store[id] = _quantized - ? Model2VecUtils.quantizeToInt8(vector) - : Float32List.fromList(vector); + _store[id] = ( + vector: _quantized + ? Model2VecUtils.quantizeToInt8(vector) + : Float32List.fromList(vector), + payload: payload, + ); } /// Stores every entry in [entries]. See [add]. @@ -129,9 +145,13 @@ class EmbeddingIndex { } final results = []; for (final entry in _store.entries) { - final vector = _asFloat32(entry.value); + final vector = _asFloat32(entry.value.vector); results.add( - SearchResult(entry.key, Model2VecUtils.cosineSimilarity(query, vector)), + SearchResult( + entry.key, + Model2VecUtils.cosineSimilarity(query, vector), + payload: entry.value.payload, + ), ); } return results; @@ -152,18 +172,29 @@ class EmbeddingIndex { // --- persistence --------------------------------------------------------- static const _magic = 'M2VI'; - static const _version = 1; + static const _version = 2; - /// Serializes the whole index to a compact binary blob. Restore it with - /// [EmbeddingIndex.fromBytes]. + // Per-entry payload length sentinel meaning "no payload" (distinct from a + // zero-length, i.e. empty-string, payload). + static const _noPayload = 0xFFFFFFFF; + + /// Serializes the whole index — vectors and payloads — to a compact binary + /// blob. Restore it with [EmbeddingIndex.fromBytes]. Uint8List toBytes() { final dim = _dimension ?? 0; final elemBytes = _quantized ? 1 : 4; final idByteList = _store.keys.map(utf8.encode).toList(growable: false); + // Null payload is encoded via the _noPayload sentinel (no trailing bytes). + final payloadByteList = _store.values + .map((e) => e.payload == null ? null : utf8.encode(e.payload!)) + .toList(growable: false); + var size = 4 + 1 + 1 + 4 + 4; // magic + version + flags + dim + count - for (final idBytes in idByteList) { - size += 4 + idBytes.length + dim * elemBytes; // u32 idLen + id + vector + for (var i = 0; i < idByteList.length; i++) { + // u32 idLen + id + vector + u32 payloadLen + payload + size += 4 + idByteList[i].length + dim * elemBytes; + size += 4 + (payloadByteList[i]?.length ?? 0); } final bytes = Uint8List(size); @@ -183,24 +214,34 @@ class EmbeddingIndex { var e = 0; for (final stored in _store.values) { - final idBytes = idByteList[e++]; + final idBytes = idByteList[e]; data.setUint32(o, idBytes.length, Endian.little); o += 4; bytes.setRange(o, o + idBytes.length, idBytes); o += idBytes.length; if (_quantized) { - final q = stored as Int8List; + final q = stored.vector as Int8List; for (var i = 0; i < dim; i++) { data.setInt8(o, q[i]); o += 1; } } else { - final f = stored as Float32List; + final f = stored.vector as Float32List; for (var i = 0; i < dim; i++) { data.setFloat32(o, f[i], Endian.little); o += 4; } } + final payloadBytes = payloadByteList[e++]; + if (payloadBytes == null) { + data.setUint32(o, _noPayload, Endian.little); + o += 4; + } else { + data.setUint32(o, payloadBytes.length, Endian.little); + o += 4; + bytes.setRange(o, o + payloadBytes.length, payloadBytes); + o += payloadBytes.length; + } } return bytes; } @@ -240,9 +281,11 @@ class EmbeddingIndex { final version = data.getUint8(o); o += 1; - if (version != _version) { + // v1 has no payload section; v2 adds a payload per entry. Both still read. + if (version != 1 && version != 2) { throw ArgumentError('unsupported EmbeddingIndex version $version'); } + final hasPayloads = version >= 2; final quantized = data.getUint8(o) == 1; o += 1; final dim = data.getUint32(o, Endian.little); @@ -256,21 +299,32 @@ class EmbeddingIndex { o += 4; final id = utf8.decode(bytes.sublist(o, o + idLen)); o += idLen; + final Object vector; if (quantized) { final q = Int8List(dim); for (var i = 0; i < dim; i++) { q[i] = data.getInt8(o); o += 1; } - index._store[id] = q; + vector = q; } else { final f = Float32List(dim); for (var i = 0; i < dim; i++) { f[i] = data.getFloat32(o, Endian.little); o += 4; } - index._store[id] = f; + vector = f; + } + String? payload; + if (hasPayloads) { + final payloadLen = data.getUint32(o, Endian.little); + o += 4; + if (payloadLen != _noPayload) { + payload = utf8.decode(bytes.sublist(o, o + payloadLen)); + o += payloadLen; + } } + index._store[id] = (vector: vector, payload: payload); } index._dimension = count > 0 ? dim : null; return index; diff --git a/test/embedding_index_test.dart b/test/embedding_index_test.dart index ca0dc86..7f8cc0c 100644 --- a/test/embedding_index_test.dart +++ b/test/embedding_index_test.dart @@ -147,5 +147,75 @@ void main() { final index = EmbeddingIndex()..add('a', _v([1, 0])); expect(index.search(_v([1, 0]), topK: -1), isEmpty); }); + + group('payload', () { + test('search returns the payload stored with each vector', () { + final index = EmbeddingIndex() + ..add('a', _v([1, 0, 0]), payload: 'alpha') + ..add('b', _v([0, 1, 0]), payload: 'beta'); + + final hits = index.search(_v([1, 0, 0]), topK: 1); + expect(hits.single.id, 'a'); + expect(hits.single.payload, 'alpha'); + }); + + test('an entry added without a payload has a null payload', () { + final index = EmbeddingIndex()..add('a', _v([1, 0, 0])); + expect(index.search(_v([1, 0, 0])).single.payload, isNull); + }); + + test('payloads survive toBytes/fromBytes round-trip', () { + final index = EmbeddingIndex() + ..add('a', _v([1, 0, 0]), payload: 'alpha') + ..add('b', _v([0, 1, 0])); // no payload + + final restored = EmbeddingIndex.fromBytes(index.toBytes()); + final byId = { + for (final r in restored.search(_v([1, 0, 0]), topK: 2)) + r.id: r.payload, + }; + expect(byId['a'], 'alpha'); + expect(byId['b'], isNull); + }); + + test('an empty-string payload round-trips distinct from null', () { + final index = EmbeddingIndex() + ..add('empty', _v([1, 0]), payload: '') + ..add('none', _v([0, 1])); + + final restored = EmbeddingIndex.fromBytes(index.toBytes()); + final byId = { + for (final r in restored.search(_v([1, 0]), topK: 2)) r.id: r.payload, + }; + expect(byId['empty'], ''); + expect(byId['none'], isNull); + }); + + test('quantized index carries payloads too', () { + final index = EmbeddingIndex(quantized: true) + ..add('a', _v([1, 0, 0]), payload: 'alpha'); + final restored = EmbeddingIndex.fromBytes(index.toBytes()); + expect(restored.search(_v([1, 0, 0])).single.payload, 'alpha'); + }); + + test('a v1 blob (no payload section) still decodes', () { + // Golden v1 blob: one entry id='a', dim=2, float [1,0], not quantized. + final v1 = Uint8List.fromList([ + 0x4D, 0x32, 0x56, 0x49, // magic 'M2VI' + 0x01, // version 1 + 0x00, // flags: not quantized + 0x02, 0x00, 0x00, 0x00, // dim = 2 + 0x01, 0x00, 0x00, 0x00, // count = 1 + 0x01, 0x00, 0x00, 0x00, // idLen = 1 + 0x61, // 'a' + 0x00, 0x00, 0x80, 0x3F, // 1.0f little-endian + 0x00, 0x00, 0x00, 0x00, // 0.0f + ]); + final index = EmbeddingIndex.fromBytes(v1); + expect(index.length, 1); + expect(index.contains('a'), isTrue); + expect(index.search(_v([1, 0])).single.payload, isNull); + }); + }); }); } From 96e87c1c8b951dc770469f4d7e587f9809aa8fd8 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Mon, 6 Jul 2026 22:20:38 +0300 Subject: [PATCH 06/20] refactor: unify ad-hoc similarity search on one scored shape MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Model2VecUtils had three near-identical searches returning three shapes: similaritySearch and similaritySearchWithThreshold handed back bare List indices (forcing callers to re-associate), while similaritySearchWithScores returned (index, score) records. Make similaritySearchWithScores the single scored search — it now takes an optional score threshold alongside topK (idiomatic topK + min-score), so it subsumes both query modes. The two bare-index methods become @Deprecated one-line delegates over it, so existing callers keep working for one release. Removes the duplicated scoring loops and the private _IndexedSimilarity type; migrates the example and README to the scored shape. --- README.md | 13 +++--- example/main.dart | 8 +++- lib/src/utils.dart | 109 +++++++++++++++++++------------------------ test/utils_test.dart | 47 ++++++++++++++++++- 4 files changed, 106 insertions(+), 71 deletions(-) diff --git a/README.md b/README.md index 159e5a6..8b7a651 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ - **Massive Streaming:** Built-in `generateEmbeddingStream` for processing millions of rows without blocking the Event Loop or overflowing RAM. - **Hugging Face Integration:** Automatically downloads and caches models directly from the Hugging Face Hub. - **Zero-Stutter Async:** Transparently runs heavy tokenization and math in background Dart Isolates using `Async` methods. -- **Vector Utilities:** Ships with high-performance mathematical tools (`cosineSimilarity`, `quantizeToInt8`, `similaritySearch`, etc.). +- **Vector Utilities:** Ships with high-performance mathematical tools (`cosineSimilarity`, `quantizeToInt8`, `similaritySearchWithScores`, etc.). ## Recommended Models @@ -162,9 +162,9 @@ final candidates = [ // 1. Semantic Similarity (Cosine) final sim = Model2VecUtils.cosineSimilarity(query, candidates[0]); -// 2. Threshold Searching (Find all matches > 80%) -final matches = Model2VecUtils.similaritySearchWithThreshold( - query, candidates, threshold: 0.8, +// 2. Threshold Searching — (index, score) pairs for all matches > 80% +final matches = Model2VecUtils.similaritySearchWithScores( + query, candidates, threshold: 0.8, topK: candidates.length, ); // 3. Scalar Quantization (Compress Float32 to Int8 to save 4x RAM) @@ -259,8 +259,7 @@ Model2Vec.unloadModel(); // releases the native model | `cosineSimilarity(a, b)` | Calculates cosine similarity (-1.0 to 1.0) between two vectors. | | `cosineDistance(a, b)` | Calculates cosine distance (0.0 to 2.0). | | `euclideanDistance(a, b)` | Calculates Euclidean (L2) distance. | -| `similaritySearch(query, docs)` | Returns the indices of the Top-K most similar vectors in a database. | -| `similaritySearchWithThreshold` | Returns all indices with similarity above a given threshold. | +| `similaritySearchWithScores(query, docs)` | Ranks an in-memory vector list by cosine similarity; returns `(index, score)` pairs, top-K with an optional score `threshold`. | | `quantizeToInt8(vector)` | Compresses a `Float32List` into an `Int8List` (4x memory savings). | | `normalize(vector)` | Applies L2 normalization to a vector. | | `meanPooling(vectors)` | Averages multiple vectors into a single vector. | @@ -282,7 +281,7 @@ Here is a performance benchmark run on a typical machine (AOT compiled): > *Note: Initial load times may vary slightly based on the disk speed. Generating an embedding takes just **a few microseconds** per string.* -- `similaritySearch` over 100,000 vectors takes **<100ms** in pure Dart. +- `similaritySearchWithScores` over 100,000 vectors takes **<100ms** in pure Dart. ## Development & Contributing diff --git a/example/main.dart b/example/main.dart index b6d2d12..d77db94 100644 --- a/example/main.dart +++ b/example/main.dart @@ -78,9 +78,13 @@ Future main() async { ' - Sim(kitten, space): ${(simSpace * 100).toStringAsFixed(1)}%', ); - final topMatch = Model2VecUtils.similaritySearch(query, db, topK: 1); + final topMatch = Model2VecUtils.similaritySearchWithScores( + query, + db, + topK: 1, + ); stdout - ..writeln(' - Best match index: ${topMatch.first}') + ..writeln(' - Best match index: ${topMatch.first.index}') // 8. Streaming API for Huge Datasets ..writeln('\n🌊 Streaming API (1000 items):'); final stream = Stream.fromIterable(List.generate(1000, (i) => 'Item $i')); diff --git a/lib/src/utils.dart b/lib/src/utils.dart index f43a1e5..e0a81a6 100644 --- a/lib/src/utils.dart +++ b/lib/src/utils.dart @@ -2,9 +2,6 @@ import 'dart:convert'; import 'dart:math' as math; import 'dart:typed_data'; -/// -typedef _IndexedSimilarity = ({int i, double s}); - /// Utilities for working with embedding vectors. // ignore: avoid_classes_with_only_static_members final class Model2VecUtils { @@ -72,66 +69,19 @@ final class Model2VecUtils { return math.sqrt(sum); } - /// Finds the indices of the top [topK] most similar vectors in [candidates] - - /// to the [query] vector. + /// Ranks [candidates] by cosine similarity to [query], most similar first, + /// pairing each kept candidate's index with its score. /// - /// Returns a list of indices sorted by similarity (descending). - static List similaritySearch( - Float32List query, - List candidates, { - int topK = 5, - }) { - if (candidates.isEmpty) { - return []; - } - - final similarities = <_IndexedSimilarity>[]; - for (var i = 0; i < candidates.length; i++) { - similarities.add( - (i: i, s: cosineSimilarity(query, candidates[i])), - ); - } - - similarities.sort((a, b) => b.s.compareTo(a.s)); - - return similarities - .take(math.min(topK, similarities.length)) - .map((s) => s.i) - .toList(growable: false); - } - - /// Finds indices of all vectors in [candidates] that have a similarity - /// score greater than or equal to [threshold] with the [query]. - /// - /// Returns a list of indices sorted by similarity (descending). - static List similaritySearchWithThreshold( - Float32List query, - List candidates, { - required double threshold, - }) { - if (candidates.isEmpty) { - return []; - } - - final similarities = <_IndexedSimilarity>[]; - for (var i = 0; i < candidates.length; i++) { - final sim = cosineSimilarity(query, candidates[i]); - if (sim >= threshold) { - similarities.add((i: i, s: sim)); - } - } - - similarities.sort((a, b) => b.s.compareTo(a.s)); - return similarities.map((s) => s.i).toList(growable: false); - } - - /// Like [similaritySearch] but returns each candidate index paired with its - /// cosine similarity to [query], most similar first. + /// Keeps at most [topK] results; pass a [threshold] to also drop any result + /// scoring below it (raise [topK] to keep every match above the floor). This + /// is the ad-hoc counterpart to `EmbeddingIndex.search` for a list of vectors + /// you already hold in memory — it is the single scored search on this class; + /// the older index-only variants are deprecated. static List<({int index, double score})> similaritySearchWithScores( Float32List query, List candidates, { int topK = 5, + double? threshold, }) { if (candidates.isEmpty || topK <= 0) { return []; @@ -139,10 +89,49 @@ final class Model2VecUtils { final scored = <({int index, double score})>[ for (var i = 0; i < candidates.length; i++) (index: i, score: cosineSimilarity(query, candidates[i])), - ]..sort((a, b) => b.score.compareTo(a.score)); - return scored.take(math.min(topK, scored.length)).toList(growable: false); + ]; + final ranked = + (threshold == null + ? scored + : scored.where((r) => r.score >= threshold).toList()) + ..sort((a, b) => b.score.compareTo(a.score)); + return ranked.take(math.min(topK, ranked.length)).toList(growable: false); } + /// Finds the indices of the top [topK] most similar vectors in [candidates], + /// sorted by similarity (descending). + // Kept one release as a delegating shim; scheduled for removal in 3.0.0. + // ignore: remove_deprecations_in_breaking_versions + @Deprecated( + 'Use similaritySearchWithScores and read .index. Will be removed in 3.0.0.', + ) + static List similaritySearch( + Float32List query, + List candidates, { + int topK = 5, + }) => similaritySearchWithScores(query, candidates, topK: topK) + .map((r) => r.index) + .toList(growable: false); + + /// Finds the indices of every vector in [candidates] whose similarity to + /// [query] is `>= threshold`, sorted by similarity (descending). + // Kept one release as a delegating shim; scheduled for removal in 3.0.0. + // ignore: remove_deprecations_in_breaking_versions + @Deprecated( + 'Use similaritySearchWithScores(threshold: ...) and read .index. ' + 'Will be removed in 3.0.0.', + ) + static List similaritySearchWithThreshold( + Float32List query, + List candidates, { + required double threshold, + }) => similaritySearchWithScores( + query, + candidates, + topK: candidates.length, + threshold: threshold, + ).map((r) => r.index).toList(growable: false); + /// Reranks [candidates] against [query] with Maximal Marginal Relevance, /// trading off relevance (similarity to the query) against diversity /// (dissimilarity to already-picked results). diff --git a/test/utils_test.dart b/test/utils_test.dart index a406302..4b7da53 100644 --- a/test/utils_test.dart +++ b/test/utils_test.dart @@ -44,7 +44,7 @@ void main() { expect(() => Model2VecUtils.cosineSimilarity(a, b), throwsArgumentError); }); - test('similaritySearch finds most similar vectors', () { + test('deprecated similaritySearch still delegates to the scored shape', () { final query = Float32List.fromList([1.0, 0.0]); final candidates = [ Float32List.fromList([0.0, 1.0]), // orthogonal @@ -53,6 +53,8 @@ void main() { Float32List.fromList([0.5, 0.5]), // somewhat similar ]; + // Deliberately exercises the deprecated shim to prove it delegates. + // ignore: deprecated_member_use_from_same_package final results = Model2VecUtils.similaritySearch( query, candidates, @@ -64,7 +66,7 @@ void main() { expect(results[1], equals(3)); // [0.5, 0.5] }); - test('similaritySearchWithThreshold works correctly', () { + test('deprecated similaritySearchWithThreshold still delegates', () { final query = Float32List.fromList([1.0, 0.0]); final candidates = [ Float32List.fromList([0.0, 1.0]), // sim: 0.0 @@ -73,6 +75,8 @@ void main() { Float32List.fromList([0.5, 0.5]), // sim: 0.707 ]; + // Deliberately exercises the deprecated shim to prove it delegates. + // ignore: deprecated_member_use_from_same_package final results = Model2VecUtils.similaritySearchWithThreshold( query, candidates, @@ -158,6 +162,45 @@ void main() { ); }); + test('similaritySearchWithScores drops results below the threshold', () { + final query = Float32List.fromList([1.0, 0.0]); + final candidates = [ + Float32List.fromList([0.0, 1.0]), // sim: 0.0 + Float32List.fromList([0.9, 0.1]), // sim: ~0.994 + Float32List.fromList([-1.0, 0.0]), // sim: -1.0 + Float32List.fromList([0.5, 0.5]), // sim: ~0.707 + ]; + + final results = Model2VecUtils.similaritySearchWithScores( + query, + candidates, + topK: 10, + threshold: 0.8, + ); + + expect(results.map((r) => r.index).toList(), [1]); + expect(results.single.score, greaterThan(0.8)); + }); + + test('similaritySearchWithScores caps threshold matches at topK', () { + final query = Float32List.fromList([1.0, 0.0]); + final candidates = [ + Float32List.fromList([1.0, 0.0]), // sim 1.0 + Float32List.fromList([0.99, 0.01]), // sim ~0.9999 + Float32List.fromList([0.98, 0.02]), // sim ~0.9998 + ]; + + final results = Model2VecUtils.similaritySearchWithScores( + query, + candidates, + topK: 2, + threshold: 0.5, // all three clear the floor, but topK caps to 2 + ); + + expect(results, hasLength(2)); + expect(results.map((r) => r.index).toList(), [0, 1]); + }); + test('MMR prefers a diverse result over a near-duplicate', () { final query = Float32List.fromList([1, 0, 0]); final candidates = [ From 8d8e87164ce9f86e440efd2d546ea848b30d21b1 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Mon, 6 Jul 2026 22:26:23 +0300 Subject: [PATCH 07/20] refactor: rename initEmbedder* to loadModel*, pairing load/unload MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The public lifecycle verbs mixed two nouns — initEmbedder* to load vs unloadModel/modelInfo — and 'Embedder' is a term CONTEXT.md explicitly says to avoid; the load/unload inverse was hidden in the names. Rename the five loaders to loadModel, loadModelAdvanced, loadModelFromBytes, loadModelAsync and loadModelAdvancedAsync, so the surface reads loadModel <-> unloadModel with one noun (Model). The old initEmbedder* names stay as @Deprecated delegating aliases for one release (removed in 3.0.0). The native init_embedder entry points stay internal. Adds a Load/Unload term to CONTEXT.md and migrates the example, benchmark, tests and README. --- README.md | 22 +++--- benchmark/embedding_benchmark.dart | 4 +- example/main.dart | 2 +- example/rag_example.dart | 2 +- lib/src/embedding_pool.dart | 2 +- lib/src/exception.dart | 2 +- lib/src/model2vec_base.dart | 104 ++++++++++++++++++++++++----- test/error_path_test.dart | 4 +- test/init_async_test.dart | 10 +-- test/lifecycle_test.dart | 10 ++- test/model2vec_test.dart | 8 +-- test/utils_test.dart | 2 +- 12 files changed, 128 insertions(+), 44 deletions(-) diff --git a/README.md b/README.md index 8b7a651..817dc2f 100644 --- a/README.md +++ b/README.md @@ -74,8 +74,8 @@ void main() { // Model2Vec is a stateless namespace of static methods — there is one // active model per process, loaded automatically via Native Assets. - // Initialize with a model from Hugging Face - Model2Vec.initEmbedder('minishlab/potion-base-2M'); + // Load a model from Hugging Face + Model2Vec.loadModel('minishlab/potion-base-2M'); // Generate an embedding final embedding = Model2Vec.generateEmbedding('Dart FFI is blazingly fast 🚀'); @@ -138,11 +138,11 @@ Future processHugeFile() async { Never block the main thread. If you are building a Flutter app, always use the `Async` variants to perform generation in a background `Isolate`. -Loading is the heaviest step of all — the first time a model is fetched it downloads tens to hundreds of MB. `initEmbedderAsync` runs that load on a background isolate; because the native model is a single process-global, the loaded model is then visible to every isolate. +Loading is the heaviest step of all — the first time a model is fetched it downloads tens to hundreds of MB. `loadModelAsync` runs that load on a background isolate; because the native model is a single process-global, the loaded model is then visible to every isolate. ```dart // Load off the main thread so the first download never freezes the UI. -await Model2Vec.initEmbedderAsync('minishlab/potion-base-2M'); +await Model2Vec.loadModelAsync('minishlab/potion-base-2M'); final embedding = await Model2Vec.generateEmbeddingAsync('A very long text...'); final batch = await Model2Vec.generateBatchEmbeddingsAsync(['A', 'B', 'C']); @@ -235,11 +235,11 @@ Model2Vec.unloadModel(); // releases the native model | Method / Property | Description | | ----------------- | ----------- | -| `initEmbedder(path)` | Initializes the model from a Hugging Face repo ID or local path. | -| `initEmbedderAdvanced(...)` | Advanced initialization with custom `cacheDirectory`, `hfToken`, or `normalize` overrides. | -| `initEmbedderFromBytes(...)` | Initializes the model directly from raw `Uint8List` bytes (`model.safetensors`, `tokenizer.json`, etc). | -| `initEmbedderAsync(path)` | Loads a model on a background isolate; await it so the first download never blocks the UI. | -| `initEmbedderAdvancedAsync(...)` | Async form of `initEmbedderAdvanced` (background isolate). | +| `loadModel(path)` | Loads the model from a Hugging Face repo ID or local path. | +| `loadModelAdvanced(...)` | Advanced load with custom `cacheDirectory`, `hfToken`, or `normalize` overrides. | +| `loadModelFromBytes(...)` | Loads the model directly from raw `Uint8List` bytes (`model.safetensors`, `tokenizer.json`, etc). | +| `loadModelAsync(path)` | Loads a model on a background isolate; await it so the first download never blocks the UI. | +| `loadModelAdvancedAsync(...)` | Async form of `loadModelAdvanced` (background isolate). | | `recommendedModels` | A typed `List` catalog of officially recommended Potion models (offline). | | `tokenize(text)` | Runs the internal BPE tokenizer and returns a `List`. | | `generateEmbedding(text)` | Synchronously generates a `Float32List` embedding vector. | @@ -252,6 +252,10 @@ Model2Vec.unloadModel(); // releases the native model | `embeddingDimension` | Property returning the vector size (e.g., 256, 384, 512). | | `vocabularySize` | Property returning the number of tokens in the model's vocabulary. | +> The `initEmbedder*` methods are deprecated aliases for `loadModel*` — the +> lifecycle pair is now `loadModel` ⇄ `unloadModel`. The old names still work +> and will be removed in 3.0.0. + ### Math Utilities (`Model2VecUtils` class) | Method | Description | diff --git a/benchmark/embedding_benchmark.dart b/benchmark/embedding_benchmark.dart index 5dedfef..c4aa0c9 100644 --- a/benchmark/embedding_benchmark.dart +++ b/benchmark/embedding_benchmark.dart @@ -47,7 +47,7 @@ void main() { for (final model in models) { stdout.write('Warming up $model... '); final watch = Stopwatch()..start(); - Model2Vec.initEmbedder(model); + Model2Vec.loadModel(model); watch.stop(); stdout.writeln('Done in ${watch.elapsedMilliseconds} ms.'); } @@ -62,7 +62,7 @@ void main() { for (final model in models) { // Measure Load Time (Already cached) final watch = Stopwatch()..start(); - Model2Vec.initEmbedder(model); + Model2Vec.loadModel(model); watch.stop(); final loadTime = '${watch.elapsedMilliseconds} ms'; diff --git a/example/main.dart b/example/main.dart index d77db94..758b8f8 100644 --- a/example/main.dart +++ b/example/main.dart @@ -18,7 +18,7 @@ Future main() async { const modelId = 'minishlab/potion-base-2M'; stdout.writeln('\n🚀 Initializing $modelId...'); final sw = Stopwatch()..start(); - Model2Vec.initEmbedder(modelId); + Model2Vec.loadModel(modelId); stdout ..writeln('✨ Initialized in ${sw.elapsedMilliseconds}ms') // 3. Inspect Model Metadata diff --git a/example/rag_example.dart b/example/rag_example.dart index 0b908f7..c547122 100644 --- a/example/rag_example.dart +++ b/example/rag_example.dart @@ -33,7 +33,7 @@ Future main() async { // 1. Load an embedding model. For real retrieval, potion-retrieval-32M is // purpose-built; potion-base-8M is a lighter, still-capable choice. - Model2Vec.initEmbedder('minishlab/potion-base-8M'); + Model2Vec.loadModel('minishlab/potion-base-8M'); // 2. Chunk each article and index the passages. We store the passage text // as the entry's payload, so a hit carries its document directly — no diff --git a/lib/src/embedding_pool.dart b/lib/src/embedding_pool.dart index bf07ea4..d07a339 100644 --- a/lib/src/embedding_pool.dart +++ b/lib/src/embedding_pool.dart @@ -12,7 +12,7 @@ import 'embedding_worker.dart'; /// parallelize embedding across CPU cores. Use this for large workloads where a /// single `Model2Vec.generateEmbeddingStream` worker is the bottleneck. /// -/// A model must be initialized (via `Model2Vec.initEmbedder`) before use, and +/// A model must be loaded (via `Model2Vec.loadModel`) before use, and /// must not be switched while the pool is working (the one-model-per-run /// contract, same as the single worker). class EmbeddingPool { diff --git a/lib/src/exception.dart b/lib/src/exception.dart index a4bacad..9941370 100644 --- a/lib/src/exception.dart +++ b/lib/src/exception.dart @@ -3,7 +3,7 @@ /// Values mirror the stable error codes the Rust layer returns, so callers can /// branch exhaustively with a `switch` instead of matching magic numbers. enum Model2VecErrorKind { - /// No model has been initialized yet (call `initEmbedder` first). + /// No model has been initialized yet (call `loadModel` first). notInitialized, /// A model could not be loaded from the given repo id or path. diff --git a/lib/src/model2vec_base.dart b/lib/src/model2vec_base.dart index 58137c4..82127e2 100644 --- a/lib/src/model2vec_base.dart +++ b/lib/src/model2vec_base.dart @@ -85,22 +85,26 @@ abstract final class Model2Vec { } }); - /// Initializes a model from a Hugging Face repo id or a local directory path. + /// Loads a model from a Hugging Face repo id or a local directory path, + /// replacing any currently active model. /// /// [modelPath] can be a repo id like `minishlab/potion-base-8M` or a path to /// a directory containing `model.safetensors`, `config.json` and /// `tokenizer.json`. - static void initEmbedder(String modelPath) => - initEmbedderAdvanced(modelPath: modelPath); + /// + /// This is synchronous and blocks the calling isolate for the whole load, + /// including the first download. Use [loadModelAsync] to load off-thread. + static void loadModel(String modelPath) => + loadModelAdvanced(modelPath: modelPath); - /// Advanced model initialization with additional options. + /// Advanced model loading with additional options. /// /// - [modelPath]: Repo id or local path. /// - [hfToken]: Optional Hugging Face API token for private repos. /// - [cacheDirectory]: Optional path to store downloaded models. /// - [normalize]: Whether to L2-normalize output embeddings. /// - [subfolder]: Optional subfolder within the repo/path. - static void initEmbedderAdvanced({ + static void loadModelAdvanced({ required String modelPath, String? hfToken, String? cacheDirectory, @@ -133,12 +137,12 @@ abstract final class Model2Vec { ); }); - /// Initializes a model using raw bytes from memory. + /// Loads a model from raw bytes in memory. /// /// - [tokenizerBytes]: Content of `tokenizer.json`. /// - [modelBytes]: Content of `model.safetensors`. /// - [configBytes]: Content of `config.json`. - static void initEmbedderFromBytes({ + static void loadModelFromBytes({ required Uint8List tokenizerBytes, required Uint8List modelBytes, required Uint8List configBytes, @@ -168,28 +172,28 @@ abstract final class Model2Vec { /// Loads a model asynchronously on a background isolate. /// - /// Prefer this over [initEmbedder] when the model may be downloaded from + /// Prefer this over [loadModel] when the model may be downloaded from /// Hugging Face for the first time (tens to hundreds of MB): the synchronous - /// [initEmbedder] blocks the calling isolate for the entire download, which + /// [loadModel] blocks the calling isolate for the entire download, which /// freezes a Flutter UI. The native model is a single process-global, so once /// the background isolate has loaded it the model is visible to every isolate /// — including the one that awaited this call. - static Future initEmbedderAsync(String modelPath) => - Isolate.run(() => initEmbedder(modelPath)); + static Future loadModelAsync(String modelPath) => + Isolate.run(() => loadModel(modelPath)); - /// Advanced counterpart to [initEmbedderAsync]. + /// Advanced counterpart to [loadModelAsync]. /// - /// Takes the same options as [initEmbedderAdvanced] and loads them on a - /// background isolate; see [initEmbedderAsync] for why loading off-thread + /// Takes the same options as [loadModelAdvanced] and loads them on a + /// background isolate; see [loadModelAsync] for why loading off-thread /// matters and why the loaded model is visible on every isolate. - static Future initEmbedderAdvancedAsync({ + static Future loadModelAdvancedAsync({ required String modelPath, String? hfToken, String? cacheDirectory, bool? normalize, String? subfolder, }) => Isolate.run( - () => initEmbedderAdvanced( + () => loadModelAdvanced( modelPath: modelPath, hfToken: hfToken, cacheDirectory: cacheDirectory, @@ -198,6 +202,74 @@ abstract final class Model2Vec { ), ); + // --- Deprecated init* aliases -------------------------------------------- + // "Embedder" is absent from the domain glossary (CONTEXT.md), which pairs + // load/unload over the Model. These shims delegate to the load* methods and + // are scheduled for removal in 3.0.0. + + /// Deprecated alias for [loadModel]. + // Delegating shim; see the section note above. Remove in 3.0.0. + // ignore: remove_deprecations_in_breaking_versions + @Deprecated('Use loadModel. Will be removed in 3.0.0.') + static void initEmbedder(String modelPath) => loadModel(modelPath); + + /// Deprecated alias for [loadModelAdvanced]. + // Delegating shim; see the section note above. Remove in 3.0.0. + // ignore: remove_deprecations_in_breaking_versions + @Deprecated('Use loadModelAdvanced. Will be removed in 3.0.0.') + static void initEmbedderAdvanced({ + required String modelPath, + String? hfToken, + String? cacheDirectory, + bool? normalize, + String? subfolder, + }) => loadModelAdvanced( + modelPath: modelPath, + hfToken: hfToken, + cacheDirectory: cacheDirectory, + normalize: normalize, + subfolder: subfolder, + ); + + /// Deprecated alias for [loadModelFromBytes]. + // Delegating shim; see the section note above. Remove in 3.0.0. + // ignore: remove_deprecations_in_breaking_versions + @Deprecated('Use loadModelFromBytes. Will be removed in 3.0.0.') + static void initEmbedderFromBytes({ + required Uint8List tokenizerBytes, + required Uint8List modelBytes, + required Uint8List configBytes, + }) => loadModelFromBytes( + tokenizerBytes: tokenizerBytes, + modelBytes: modelBytes, + configBytes: configBytes, + ); + + /// Deprecated alias for [loadModelAsync]. + // Delegating shim; see the section note above. Remove in 3.0.0. + // ignore: remove_deprecations_in_breaking_versions + @Deprecated('Use loadModelAsync. Will be removed in 3.0.0.') + static Future initEmbedderAsync(String modelPath) => + loadModelAsync(modelPath); + + /// Deprecated alias for [loadModelAdvancedAsync]. + // Delegating shim; see the section note above. Remove in 3.0.0. + // ignore: remove_deprecations_in_breaking_versions + @Deprecated('Use loadModelAdvancedAsync. Will be removed in 3.0.0.') + static Future initEmbedderAdvancedAsync({ + required String modelPath, + String? hfToken, + String? cacheDirectory, + bool? normalize, + String? subfolder, + }) => loadModelAdvancedAsync( + modelPath: modelPath, + hfToken: hfToken, + cacheDirectory: cacheDirectory, + normalize: normalize, + subfolder: subfolder, + ); + /// Unloads the active model and frees its native memory. /// /// Safe to call when no model is loaded. After this, [isInitialized] is diff --git a/test/error_path_test.dart b/test/error_path_test.dart index f44b782..3967eb1 100644 --- a/test/error_path_test.dart +++ b/test/error_path_test.dart @@ -42,10 +42,10 @@ void main() { }); group('native error path', () { - test('initEmbedderFromBytes with garbage throws a typed exception', () { + test('loadModelFromBytes with garbage throws a typed exception', () { final garbage = Uint8List.fromList([0, 1, 2, 3]); expect( - () => Model2Vec.initEmbedderFromBytes( + () => Model2Vec.loadModelFromBytes( tokenizerBytes: garbage, modelBytes: garbage, configBytes: garbage, diff --git a/test/init_async_test.dart b/test/init_async_test.dart index 28d5650..12376f1 100644 --- a/test/init_async_test.dart +++ b/test/init_async_test.dart @@ -6,10 +6,10 @@ import 'package:test/test.dart'; void main() { group('async model loading', () { test( - 'initEmbedderAsync loads the process-global model, visible on the ' + 'loadModelAsync loads the process-global model, visible on the ' 'calling isolate', () async { - await Model2Vec.initEmbedderAsync('minishlab/potion-base-2M'); + await Model2Vec.loadModelAsync('minishlab/potion-base-2M'); // The load ran on a background isolate; the native model is a // process-global, so a synchronous call on this isolate sees it. @@ -22,10 +22,10 @@ void main() { }, ); - test('initEmbedderAdvancedAsync honors a custom cache directory', () async { + test('loadModelAdvancedAsync honors a custom cache directory', () async { final tempDir = Directory.systemTemp.createTempSync('m2v_async_cache_'); try { - await Model2Vec.initEmbedderAdvancedAsync( + await Model2Vec.loadModelAdvancedAsync( modelPath: 'minishlab/potion-base-2M', cacheDirectory: tempDir.path, ); @@ -38,7 +38,7 @@ void main() { test('a failed async load surfaces a typed Model2VecException', () { expect( - Model2Vec.initEmbedderAsync('definitely/not-a-real-model-xyz'), + Model2Vec.loadModelAsync('definitely/not-a-real-model-xyz'), throwsA(isA()), ); }); diff --git a/test/lifecycle_test.dart b/test/lifecycle_test.dart index fe133fc..efbd2fb 100644 --- a/test/lifecycle_test.dart +++ b/test/lifecycle_test.dart @@ -4,7 +4,7 @@ import 'package:test/test.dart'; void main() { group('model lifecycle', () { test('isInitialized / modelInfo / unloadModel round-trip', () { - Model2Vec.initEmbedder('minishlab/potion-base-2M'); + Model2Vec.loadModel('minishlab/potion-base-2M'); expect(Model2Vec.isInitialized, isTrue); final info = Model2Vec.modelInfo; @@ -30,6 +30,14 @@ void main() { expect(Model2Vec.unloadModel, returnsNormally); // Re-initialize so later suites in this process find a model. + Model2Vec.loadModel('minishlab/potion-base-2M'); + expect(Model2Vec.isInitialized, isTrue); + }); + + test('deprecated initEmbedder alias still loads a model', () { + Model2Vec.unloadModel(); + // Proves the back-compat shim delegates to loadModel. + // ignore: deprecated_member_use_from_same_package Model2Vec.initEmbedder('minishlab/potion-base-2M'); expect(Model2Vec.isInitialized, isTrue); }); diff --git a/test/model2vec_test.dart b/test/model2vec_test.dart index 2b0b1e3..cfe02b3 100644 --- a/test/model2vec_test.dart +++ b/test/model2vec_test.dart @@ -8,7 +8,7 @@ void main() { test('Successful initialization (online)', () { // Must be the first init call so later tests have a loaded model. expect( - () => Model2Vec.initEmbedder('minishlab/potion-base-2M'), + () => Model2Vec.loadModel('minishlab/potion-base-2M'), returnsNormally, ); }); @@ -145,11 +145,11 @@ void main() { group('Model Switching', () { test('can switch between different models successfully', () { // Switch to 8M model (dimension 256) - Model2Vec.initEmbedder('minishlab/potion-base-8M'); + Model2Vec.loadModel('minishlab/potion-base-8M'); expect(Model2Vec.embeddingDimension, equals(256)); // Switch back to 2M model (dimension 64) - Model2Vec.initEmbedder('minishlab/potion-base-2M'); + Model2Vec.loadModel('minishlab/potion-base-2M'); expect(Model2Vec.embeddingDimension, equals(64)); }); }); @@ -180,7 +180,7 @@ void main() { try { // Now that switching is enabled, this should return normally expect( - () => Model2Vec.initEmbedderAdvanced( + () => Model2Vec.loadModelAdvanced( modelPath: 'minishlab/potion-base-2M', cacheDirectory: tempDir.path, ), diff --git a/test/utils_test.dart b/test/utils_test.dart index 4b7da53..0964862 100644 --- a/test/utils_test.dart +++ b/test/utils_test.dart @@ -302,7 +302,7 @@ void main() { group('Model2VecUtils - Real World', () { setUpAll(() { - Model2Vec.initEmbedder('minishlab/potion-base-2M'); + Model2Vec.loadModel('minishlab/potion-base-2M'); }); test('semantic similarity makes sense', () { From ff1f09e2bfeba2ec04e2535a4be65d5338738b08 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Mon, 6 Jul 2026 22:29:12 +0300 Subject: [PATCH 08/20] refactor: drop batchSize from generateBatchEmbeddings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit generateBatchEmbeddings and its async form exposed a batchSize parameter that the README itself described as 'internal chunks sent to the FFI layer' — a native throughput detail the caller has no basis to choose, and one that never affects the result. Remove it from both signatures and pass a fixed internal chunk size (1024, the former default, so results are unchanged) to the native call. The stream's batchSize stays: there it is a meaningful, user-facing chunk size. Internal callers stop threading batch.length through. --- README.md | 5 ++--- lib/src/embedding_worker.dart | 1 - lib/src/model2vec_base.dart | 22 ++++++++-------------- test/model2vec_test.dart | 3 +-- 4 files changed, 11 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index 817dc2f..1daf94a 100644 --- a/README.md +++ b/README.md @@ -95,15 +95,14 @@ void main() { ### 1. Advanced Batch Processing -Process multiple strings at once for maximum hardware utilization. You can control sequence truncation and batch sizes. +Process multiple strings at once for maximum hardware utilization. You can control sequence truncation with `maxLength`. ```dart final texts = ['Dart', 'Rust', 'Flutter']; final embeddings = Model2Vec.generateBatchEmbeddings( texts, - maxLength: 256, // Truncate strings longer than 256 tokens - batchSize: 1024, // Internal chunks sent to the FFI layer + maxLength: 256, // Truncate strings longer than 256 tokens ); ``` diff --git a/lib/src/embedding_worker.dart b/lib/src/embedding_worker.dart index 0bbfc6d..5e54ef9 100644 --- a/lib/src/embedding_worker.dart +++ b/lib/src/embedding_worker.dart @@ -61,7 +61,6 @@ void embeddingWorkerEntryPoint(SendPort mainSendPort) { final results = Model2Vec.generateBatchEmbeddings( request.batch, maxLength: request.maxLength, - batchSize: request.batch.length, ); mainSendPort.send(results); } on Model2VecException catch (e) { diff --git a/lib/src/model2vec_base.dart b/lib/src/model2vec_base.dart index 82127e2..872b26c 100644 --- a/lib/src/model2vec_base.dart +++ b/lib/src/model2vec_base.dart @@ -369,17 +369,20 @@ abstract final class Model2Vec { } }); + /// How many texts the native layer processes per internal chunk. This is a + /// throughput detail with no bearing on the result, so it is not exposed on + /// [generateBatchEmbeddings]. + static const _ffiBatchSize = 1024; + /// Generates embeddings for multiple [texts] in a single batch call. /// /// - [maxLength]: Maximum number of tokens to keep before truncating. - /// - [batchSize]: Size of the internal batches sent to the model. /// /// Returns a list of [Float32List], one per input text. /// Throws a [Model2VecException] if generation fails. static List generateBatchEmbeddings( List texts, { int maxLength = 512, - int batchSize = 1024, }) { if (texts.isEmpty) { return []; @@ -402,7 +405,7 @@ abstract final class Model2Vec { textPointers, count, maxLength, - batchSize, + _ffiBatchSize, outData, outDim, outCount, @@ -442,13 +445,8 @@ abstract final class Model2Vec { static Future> generateBatchEmbeddingsAsync( List texts, { int maxLength = 512, - int batchSize = 1024, }) => Isolate.run( - () => Model2Vec.generateBatchEmbeddings( - texts, - maxLength: maxLength, - batchSize: batchSize, - ), + () => Model2Vec.generateBatchEmbeddings(texts, maxLength: maxLength), ); /// Consumes a stream of [texts] and yields a stream of embeddings. @@ -469,11 +467,7 @@ abstract final class Model2Vec { if (!useIsolate) { await for (final batch in batches) { - final results = generateBatchEmbeddings( - batch, - maxLength: maxLength, - batchSize: batch.length, - ); + final results = generateBatchEmbeddings(batch, maxLength: maxLength); for (final result in results) { yield result; } diff --git a/test/model2vec_test.dart b/test/model2vec_test.dart index cfe02b3..5c93aa3 100644 --- a/test/model2vec_test.dart +++ b/test/model2vec_test.dart @@ -155,7 +155,7 @@ void main() { }); group('Advanced Features', () { - test('supports advanced parameters maxLength and batchSize', () { + test('maxLength truncation applies to single and batch embeddings', () { const text = 'A very long sentence to test truncation with maxLength parameter'; // With very short maxLength, the embedding should be different from @@ -168,7 +168,6 @@ void main() { final batchEmbedding = Model2Vec.generateBatchEmbeddings( [text, text], maxLength: 2, - batchSize: 1, ); expect(batchEmbedding.length, 2); expect(batchEmbedding[0], orderedEquals(shortEmbedding)); From d17dc12cbf0e3a306b3e4641992d36f8aece9da4 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Tue, 7 Jul 2026 11:20:30 +0300 Subject: [PATCH 09/20] cleanup --- CHANGELOG.md | 7 ++++ README.md | 4 --- hook/build.dart | 48 +++++++++++++------------- lib/src/model2vec_base.dart | 68 ------------------------------------- lib/src/utils.dart | 34 ------------------- native/src/lib.rs | 2 +- test/lifecycle_test.dart | 8 ----- test/utils_test.dart | 43 ----------------------- 8 files changed, 32 insertions(+), 182 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 78ec0e7..cc80145 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,11 @@ correctness. **This release is breaking** — see migration below. its constructor is `(kind, message, [code])` and the `fromCode` factory is replaced by `fromNative(code, message)`. Native failures surface the message produced by the Rust layer, each with an exhaustively-switchable `kind`. +- **Lifecycle naming.** The `initEmbedder*` methods are renamed to `loadModel*`, + pairing `loadModel` ⇄ `unloadModel` over the model. `initEmbedder`, + `initEmbedderAdvanced`, `initEmbedderFromBytes` and their async forms are + removed. `Model2VecUtils.similaritySearch` /`similaritySearchWithThreshold` + are removed in favour of `similaritySearchWithScores` (read `.index`). **Improvements:** @@ -57,6 +62,8 @@ correctness. **This release is breaking** — see migration below. | `Model2Vec.instance.generateEmbedding(t)` | `Model2Vec.generateEmbedding(t)` | | `Model2Vec.boot(lib)` / `Model2Vec(lib)` | removed — resolution is automatic | | `Model2Vec.instance.getRecommendedModels()` | `Model2Vec.recommendedModels` (typed) | +| `Model2Vec.instance.initEmbedder(path)` | `Model2Vec.loadModel(path)` | +| `Model2VecUtils.similaritySearch(q, c)` | `similaritySearchWithScores(q, c).map((r) => r.index)` | | `catch (e) { e.code }` | still works; add `e.kind` for exhaustive handling | # 1.2.0 diff --git a/README.md b/README.md index 1daf94a..9af6213 100644 --- a/README.md +++ b/README.md @@ -251,10 +251,6 @@ Model2Vec.unloadModel(); // releases the native model | `embeddingDimension` | Property returning the vector size (e.g., 256, 384, 512). | | `vocabularySize` | Property returning the number of tokens in the model's vocabulary. | -> The `initEmbedder*` methods are deprecated aliases for `loadModel*` — the -> lifecycle pair is now `loadModel` ⇄ `unloadModel`. The old names still work -> and will be removed in 3.0.0. - ### Math Utilities (`Model2VecUtils` class) | Method | Description | diff --git a/hook/build.dart b/hook/build.dart index 279d7d8..b5736ed 100644 --- a/hook/build.dart +++ b/hook/build.dart @@ -17,7 +17,7 @@ void main(List args) async { const buildMode = 'release'; final targetTriple = _mapToTriple(codeConfig); - final linkMode = codeConfig.linkModePreference == LinkModePreference.static + final linkMode = codeConfig.linkModePreference == .static ? StaticLinking() : DynamicLoadingBundled(); @@ -86,7 +86,7 @@ void main(List args) async { package: input.packageName, name: 'model2vec.so', linkMode: linkMode, - file: Uri.file(binaryPath), + file: .file(binaryPath), ), ); @@ -97,10 +97,10 @@ void main(List args) async { output.dependencies.addAll(deps); } else { // Fallback if .d file is missing - output.dependencies.add(Uri.directory(nativeDir)); + output.dependencies.add(.directory(nativeDir)); } // Always track the manifest - output.dependencies.add(Uri.file(manifestPath)); + output.dependencies.add(.file(manifestPath)); }); } @@ -108,15 +108,15 @@ bool _isHost(CodeConfig codeConfig) { final os = codeConfig.targetOS; final arch = codeConfig.targetArchitecture; - if (Platform.isLinux && os == OS.linux && arch == Architecture.x64) { + if (Platform.isLinux && os == .linux && arch == .x64) { return true; } - if (Platform.isMacOS && os == OS.macOS) { - return arch == Architecture.arm64 || arch == Architecture.x64; + if (Platform.isMacOS && os == .macOS) { + return arch == .arm64 || arch == .x64; } - if (Platform.isWindows && os == OS.windows && arch == Architecture.x64) { + if (Platform.isWindows && os == .windows && arch == .x64) { return true; } @@ -128,30 +128,30 @@ String _mapToTriple(CodeConfig codeConfig) { final arch = codeConfig.targetArchitecture; return switch ((os, arch)) { - (OS.linux, Architecture.x64) => 'x86_64-unknown-linux-gnu', - (OS.linux, Architecture.arm64) => 'aarch64-unknown-linux-gnu', - (OS.macOS, Architecture.x64) => 'x86_64-apple-darwin', - (OS.macOS, Architecture.arm64) => 'aarch64-apple-darwin', - (OS.windows, Architecture.x64) => 'x86_64-pc-windows-msvc', - (OS.windows, Architecture.arm64) => 'aarch64-pc-windows-msvc', - (OS.android, Architecture.arm64) => 'aarch64-linux-android', - (OS.android, Architecture.arm) => 'armv7-linux-androideabi', - (OS.android, Architecture.x64) => 'x86_64-linux-android', - (OS.iOS, Architecture.arm64) => - codeConfig.iOS.targetSdk == IOSSdk.iPhoneSimulator + (.linux, .x64) => 'x86_64-unknown-linux-gnu', + (.linux, .arm64) => 'aarch64-unknown-linux-gnu', + (.macOS, .x64) => 'x86_64-apple-darwin', + (.macOS, .arm64) => 'aarch64-apple-darwin', + (.windows, .x64) => 'x86_64-pc-windows-msvc', + (.windows, .arm64) => 'aarch64-pc-windows-msvc', + (.android, .arm64) => 'aarch64-linux-android', + (.android, .arm) => 'armv7-linux-androideabi', + (.android, .x64) => 'x86_64-linux-android', + (.iOS, .arm64) => + codeConfig.iOS.targetSdk == .iPhoneSimulator ? 'aarch64-apple-ios-sim' : 'aarch64-apple-ios', - (OS.iOS, Architecture.x64) => 'x86_64-apple-ios', + (.iOS, .x64) => 'x86_64-apple-ios', _ => throw UnsupportedError('Unsupported target: $os $arch'), }; } String _getLibraryFileName(OS os, String crateName) { - if (os == OS.windows) { + if (os == .windows) { return '$crateName.dll'; } - if (os == OS.macOS || os == OS.iOS) { + if (os == .macOS || os == .iOS) { return 'lib$crateName.dylib'; } @@ -187,7 +187,7 @@ Map _getBuildEnvVars(CodeConfig codeConfig) { final env = {}; final targetOS = codeConfig.targetOS; - if (targetOS == OS.android) { + if (targetOS == .android) { final cCompiler = codeConfig.cCompiler; if (cCompiler != null) { final targetTriple = _mapToTriple(codeConfig); @@ -259,7 +259,7 @@ Map _getBuildEnvVars(CodeConfig codeConfig) { } } - if (targetOS == OS.iOS) { + if (targetOS == .iOS) { try { env['IPHONEOS_DEPLOYMENT_TARGET'] = codeConfig.iOS.targetVersion .toString(); diff --git a/lib/src/model2vec_base.dart b/lib/src/model2vec_base.dart index 872b26c..dc2c5e7 100644 --- a/lib/src/model2vec_base.dart +++ b/lib/src/model2vec_base.dart @@ -202,74 +202,6 @@ abstract final class Model2Vec { ), ); - // --- Deprecated init* aliases -------------------------------------------- - // "Embedder" is absent from the domain glossary (CONTEXT.md), which pairs - // load/unload over the Model. These shims delegate to the load* methods and - // are scheduled for removal in 3.0.0. - - /// Deprecated alias for [loadModel]. - // Delegating shim; see the section note above. Remove in 3.0.0. - // ignore: remove_deprecations_in_breaking_versions - @Deprecated('Use loadModel. Will be removed in 3.0.0.') - static void initEmbedder(String modelPath) => loadModel(modelPath); - - /// Deprecated alias for [loadModelAdvanced]. - // Delegating shim; see the section note above. Remove in 3.0.0. - // ignore: remove_deprecations_in_breaking_versions - @Deprecated('Use loadModelAdvanced. Will be removed in 3.0.0.') - static void initEmbedderAdvanced({ - required String modelPath, - String? hfToken, - String? cacheDirectory, - bool? normalize, - String? subfolder, - }) => loadModelAdvanced( - modelPath: modelPath, - hfToken: hfToken, - cacheDirectory: cacheDirectory, - normalize: normalize, - subfolder: subfolder, - ); - - /// Deprecated alias for [loadModelFromBytes]. - // Delegating shim; see the section note above. Remove in 3.0.0. - // ignore: remove_deprecations_in_breaking_versions - @Deprecated('Use loadModelFromBytes. Will be removed in 3.0.0.') - static void initEmbedderFromBytes({ - required Uint8List tokenizerBytes, - required Uint8List modelBytes, - required Uint8List configBytes, - }) => loadModelFromBytes( - tokenizerBytes: tokenizerBytes, - modelBytes: modelBytes, - configBytes: configBytes, - ); - - /// Deprecated alias for [loadModelAsync]. - // Delegating shim; see the section note above. Remove in 3.0.0. - // ignore: remove_deprecations_in_breaking_versions - @Deprecated('Use loadModelAsync. Will be removed in 3.0.0.') - static Future initEmbedderAsync(String modelPath) => - loadModelAsync(modelPath); - - /// Deprecated alias for [loadModelAdvancedAsync]. - // Delegating shim; see the section note above. Remove in 3.0.0. - // ignore: remove_deprecations_in_breaking_versions - @Deprecated('Use loadModelAdvancedAsync. Will be removed in 3.0.0.') - static Future initEmbedderAdvancedAsync({ - required String modelPath, - String? hfToken, - String? cacheDirectory, - bool? normalize, - String? subfolder, - }) => loadModelAdvancedAsync( - modelPath: modelPath, - hfToken: hfToken, - cacheDirectory: cacheDirectory, - normalize: normalize, - subfolder: subfolder, - ); - /// Unloads the active model and frees its native memory. /// /// Safe to call when no model is loaded. After this, [isInitialized] is diff --git a/lib/src/utils.dart b/lib/src/utils.dart index e0a81a6..11323e2 100644 --- a/lib/src/utils.dart +++ b/lib/src/utils.dart @@ -98,40 +98,6 @@ final class Model2VecUtils { return ranked.take(math.min(topK, ranked.length)).toList(growable: false); } - /// Finds the indices of the top [topK] most similar vectors in [candidates], - /// sorted by similarity (descending). - // Kept one release as a delegating shim; scheduled for removal in 3.0.0. - // ignore: remove_deprecations_in_breaking_versions - @Deprecated( - 'Use similaritySearchWithScores and read .index. Will be removed in 3.0.0.', - ) - static List similaritySearch( - Float32List query, - List candidates, { - int topK = 5, - }) => similaritySearchWithScores(query, candidates, topK: topK) - .map((r) => r.index) - .toList(growable: false); - - /// Finds the indices of every vector in [candidates] whose similarity to - /// [query] is `>= threshold`, sorted by similarity (descending). - // Kept one release as a delegating shim; scheduled for removal in 3.0.0. - // ignore: remove_deprecations_in_breaking_versions - @Deprecated( - 'Use similaritySearchWithScores(threshold: ...) and read .index. ' - 'Will be removed in 3.0.0.', - ) - static List similaritySearchWithThreshold( - Float32List query, - List candidates, { - required double threshold, - }) => similaritySearchWithScores( - query, - candidates, - topK: candidates.length, - threshold: threshold, - ).map((r) => r.index).toList(growable: false); - /// Reranks [candidates] against [query] with Maximal Marginal Relevance, /// trading off relevance (similarity to the query) against diversity /// (dissimilarity to already-picked results). diff --git a/native/src/lib.rs b/native/src/lib.rs index 38fb4df..c47e421 100644 --- a/native/src/lib.rs +++ b/native/src/lib.rs @@ -88,7 +88,7 @@ where .map_err(|_| FfiError::new(CODE_LOCK_POISONED, "model lock poisoned"))?; let model = lock .as_ref() - .ok_or_else(|| FfiError::new(CODE_NOT_INITIALIZED, "no model initialized; call initEmbedder first"))?; + .ok_or_else(|| FfiError::new(CODE_NOT_INITIALIZED, "no model initialized; call loadModel first"))?; f(model) } diff --git a/test/lifecycle_test.dart b/test/lifecycle_test.dart index efbd2fb..0e10e70 100644 --- a/test/lifecycle_test.dart +++ b/test/lifecycle_test.dart @@ -33,13 +33,5 @@ void main() { Model2Vec.loadModel('minishlab/potion-base-2M'); expect(Model2Vec.isInitialized, isTrue); }); - - test('deprecated initEmbedder alias still loads a model', () { - Model2Vec.unloadModel(); - // Proves the back-compat shim delegates to loadModel. - // ignore: deprecated_member_use_from_same_package - Model2Vec.initEmbedder('minishlab/potion-base-2M'); - expect(Model2Vec.isInitialized, isTrue); - }); }); } diff --git a/test/utils_test.dart b/test/utils_test.dart index 0964862..0adc7ff 100644 --- a/test/utils_test.dart +++ b/test/utils_test.dart @@ -44,49 +44,6 @@ void main() { expect(() => Model2VecUtils.cosineSimilarity(a, b), throwsArgumentError); }); - test('deprecated similaritySearch still delegates to the scored shape', () { - final query = Float32List.fromList([1.0, 0.0]); - final candidates = [ - Float32List.fromList([0.0, 1.0]), // orthogonal - Float32List.fromList([0.9, 0.1]), // very similar - Float32List.fromList([-1.0, 0.0]), // opposite - Float32List.fromList([0.5, 0.5]), // somewhat similar - ]; - - // Deliberately exercises the deprecated shim to prove it delegates. - // ignore: deprecated_member_use_from_same_package - final results = Model2VecUtils.similaritySearch( - query, - candidates, - topK: 2, - ); - - expect(results, hasLength(2)); - expect(results[0], equals(1)); // [0.9, 0.1] - expect(results[1], equals(3)); // [0.5, 0.5] - }); - - test('deprecated similaritySearchWithThreshold still delegates', () { - final query = Float32List.fromList([1.0, 0.0]); - final candidates = [ - Float32List.fromList([0.0, 1.0]), // sim: 0.0 - Float32List.fromList([0.9, 0.1]), // sim: ~0.99 - Float32List.fromList([-1.0, 0.0]), // sim: -1.0 - Float32List.fromList([0.5, 0.5]), // sim: 0.707 - ]; - - // Deliberately exercises the deprecated shim to prove it delegates. - // ignore: deprecated_member_use_from_same_package - final results = Model2VecUtils.similaritySearchWithThreshold( - query, - candidates, - threshold: 0.8, - ); - - expect(results, hasLength(1)); - expect(results[0], equals(1)); // only [0.9, 0.1] passes 0.8 - }); - test('cosineDistance calculates correctly', () { final a = Float32List.fromList([1.0, 0.0]); final b = Float32List.fromList([0.0, 1.0]); From 7929f29dea5c9107870ca0973791f90cf260b0d3 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Tue, 7 Jul 2026 12:20:28 +0300 Subject: [PATCH 10/20] feat: stream model load progress MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Loading a model is a single blocking FFI call, and the first fetch pulls tens to hundreds of MB of weights from Hugging Face with no way to observe it — a Flutter app can only show an indeterminate spinner during the heaviest step of all. Add loadModelWithProgress, which loads on a background isolate and returns a Stream reporting the weights download (bytesDownloaded / totalBytes / fraction) and a coarse phase: resolving -> downloading -> parsing -> done. The Rust side routes the weights fetch through hf-hub's download_with_progress into process-global atomics (a cache hit skips it entirely), read over FFI via new get_load_progress / reset_load_progress. A cached model or local path streams straight to done; the stream always ends on done and surfaces a load failure as an error event. Adds an integration test that downloads into a fresh cache, plus README and CHANGELOG entries. Also reorganize example/ so each file teaches one story: main.dart is now a small quickstart (load -> embed -> cosine similarity), scaling_example.dart covers the at-scale patterns (progress-bar load, batch, EmbeddingPool, streaming), and rag_example.dart keeps the retrieval pipeline. Adds example/README.md indexing the three. --- CHANGELOG.md | 5 ++ README.md | 17 ++++ example/README.md | 38 +++++++++ example/main.dart | 129 +++++------------------------- example/scaling_example.dart | 85 ++++++++++++++++++++ lib/model2vec.dart | 1 + lib/src/load_progress.dart | 57 +++++++++++++ lib/src/model2vec_base.dart | 89 +++++++++++++++++++++ lib/src/model2vec_bindings.g.dart | 16 ++++ native/model2vec.h | 15 ++++ native/src/lib.rs | 24 ++++++ native/src/model.rs | 122 +++++++++++++++++++++++++--- test/load_progress_test.dart | 79 ++++++++++++++++++ 13 files changed, 559 insertions(+), 118 deletions(-) create mode 100644 example/README.md create mode 100644 example/scaling_example.dart create mode 100644 lib/src/load_progress.dart create mode 100644 test/load_progress_test.dart diff --git a/CHANGELOG.md b/CHANGELOG.md index cc80145..86ca7f6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -52,6 +52,11 @@ correctness. **This release is breaking** — see migration below. `Model2Vec.unloadModel()` (free the native model), `Model2Vec.modelInfo` (all metadata in one `ModelInfo`), and `Model2VecUtils.dequantizeInt8` (the inverse of `quantizeToInt8`). +- **Load progress.** `Model2Vec.loadModelWithProgress()` loads on a background + isolate and returns a `Stream` reporting the HF weights download + (`bytesDownloaded` / `totalBytes` / `fraction`) plus a coarse `LoadPhase` + (resolving → downloading → parsing → done). A cached model or local path + streams straight to `done`. - **Parallel worker pool.** `EmbeddingPool` fans batches across N worker isolates to embed concurrently across CPU cores. diff --git a/README.md b/README.md index 9af6213..9630cac 100644 --- a/README.md +++ b/README.md @@ -147,6 +147,22 @@ final embedding = await Model2Vec.generateEmbeddingAsync('A very long text...'); final batch = await Model2Vec.generateBatchEmbeddingsAsync(['A', 'B', 'C']); ``` +Want a progress bar for that first download? `loadModelWithProgress` loads on a background isolate and streams `LoadProgress` snapshots. Byte counts appear while the weights download (a cached model or local path jumps straight to `done`); the stream always ends on `LoadPhase.done` and surfaces any load failure as an error event. + +```dart +await for (final p in Model2Vec.loadModelWithProgress('minishlab/potion-base-2M')) { + switch (p.phase) { + case LoadPhase.downloading: + // p.fraction is 0.0..1.0, or null until the total size is known. + print('Downloading ${((p.fraction ?? 0) * 100).toStringAsFixed(0)}%'); + case LoadPhase.resolving || LoadPhase.parsing: + print('Preparing…'); + case LoadPhase.done: + print('Ready'); + } +} +``` + ### 4. Vector Math & Quantization The library ships with `Model2VecUtils` — a powerful suite of math operations tuned for embeddings. @@ -239,6 +255,7 @@ Model2Vec.unloadModel(); // releases the native model | `loadModelFromBytes(...)` | Loads the model directly from raw `Uint8List` bytes (`model.safetensors`, `tokenizer.json`, etc). | | `loadModelAsync(path)` | Loads a model on a background isolate; await it so the first download never blocks the UI. | | `loadModelAdvancedAsync(...)` | Async form of `loadModelAdvanced` (background isolate). | +| `loadModelWithProgress(path)` | Loads on a background isolate and returns a `Stream` reporting download progress; ends on `LoadPhase.done`. | | `recommendedModels` | A typed `List` catalog of officially recommended Potion models (offline). | | `tokenize(text)` | Runs the internal BPE tokenizer and returns a `List`. | | `generateEmbedding(text)` | Synchronously generates a `Float32List` embedding vector. | diff --git a/example/README.md b/example/README.md new file mode 100644 index 0000000..e2beb67 --- /dev/null +++ b/example/README.md @@ -0,0 +1,38 @@ +# Model2Vec examples + +Runnable examples for the `model2vec` package. Each is a standalone script that +imports the package with `package:model2vec/model2vec.dart`, just like your own +code would. + +Each example tells one story. Start with `main.dart`, then reach for the other +two when you need retrieval or scale. + +## Running + +```sh +dart run example/main.dart # quickstart +dart run example/scaling_example.dart # batch, parallel, streaming +dart run example/rag_example.dart # local retrieval (RAG) +``` + +The native Rust library is compiled automatically on first run via the package's +build hook, so there is nothing to set up. The first run also downloads the +embedding model from Hugging Face (a few MB), so it needs network access; later +runs use the local cache. + +## What's here + +- **[`main.dart`](main.dart)** — quickstart. Load a model, embed a few + sentences, and compare them with cosine similarity. The shortest path from + zero to a result. + +- **[`scaling_example.dart`](scaling_example.dart)** — production, at scale: + loading with a live download progress bar (`loadModelWithProgress`), batch + embedding, parallel embedding across CPU cores with an `EmbeddingPool`, and + the streaming API for datasets too large to hold in memory. + +- **[`rag_example.dart`](rag_example.dart)** — a local retrieval (RAG) pipeline: + `chunkText` splits documents, `EmbeddingIndex` stores each passage with its + text as payload, then it answers questions with nearest-neighbour search, + threshold filtering, and MMR reranking for diverse results — and persists the + index to disk with `toBytes` / `fromBytes`. diff --git a/example/main.dart b/example/main.dart index 758b8f8..4431435 100644 --- a/example/main.dart +++ b/example/main.dart @@ -2,112 +2,27 @@ import 'dart:io'; import 'package:model2vec/model2vec.dart'; -/// Example demonstrating production-ready usage of the Model2Vec package. +/// Quickstart: load a model, embed a few sentences, and compare them by cosine +/// similarity. Run with `dart run example/main.dart`. +/// +/// See `scaling_example.dart` for batch/parallel/streaming embedding, and +/// `rag_example.dart` for local retrieval (RAG). Future main() async { - try { - stdout.writeln('=== Model2Vec Example ==='); - - // 1. Explore Recommended Models - const models = Model2Vec.recommendedModels; - stdout.writeln('\n📦 Available models:'); - for (final m in models) { - stdout.writeln(' - ${m.name.padRight(25)} ID: ${m.id}'); - } - - // 2. Initialize a Model - const modelId = 'minishlab/potion-base-2M'; - stdout.writeln('\n🚀 Initializing $modelId...'); - final sw = Stopwatch()..start(); - Model2Vec.loadModel(modelId); - stdout - ..writeln('✨ Initialized in ${sw.elapsedMilliseconds}ms') - // 3. Inspect Model Metadata - ..writeln('\n📊 Model Metadata:') - ..writeln(' - Dimension: ${Model2Vec.embeddingDimension}') - ..writeln(' - Vocabulary Size: ${Model2Vec.vocabularySize}') - ..writeln(' - Is Normalized: ${Model2Vec.isNormalized}') - ..writeln(' - Median Token Length: ${Model2Vec.medianTokenLength}'); - - // 4. Tokenization Demo - const text = 'Model2Vec is incredibly fast!'; - final tokens = Model2Vec.tokenize(text); - stdout - ..writeln('\n🔍 Tokenization: "$text"') - ..writeln(' - Tokens: $tokens') - // 5. Single Embedding - ..writeln('\n🧠 Generating single embedding...'); - final embedding = Model2Vec.generateEmbedding(text); - stdout - ..writeln(' - Vector (first 3): ${embedding.take(3).toList()}') - ..writeln(' - Total Length: ${embedding.length}'); - - // 6. Batch Embedding (Production Optimization) - final texts = [ - 'The first sentence.', - 'A second, slightly longer sentence for the batch.', - 'Third one.', - ]; - stdout.writeln( - '\n⚡ Generating batch embeddings for ${texts.length} sentences...', - ); - final batchStartTime = DateTime.now(); - final batch = Model2Vec.generateBatchEmbeddings(texts); - final batchDuration = DateTime.now().difference(batchStartTime); - - stdout.writeln(' - Processed in ${batchDuration.inMicroseconds}μs'); - for (var i = 0; i < batch.length; i++) { - stdout.writeln(' - Result $i length: ${batch[i].length}'); - } - - // 7. Vector Math & Semantic Search - stdout.writeln('\n🧠 Vector Math & Semantic Search:'); - final query = Model2Vec.generateEmbedding('A cute little kitten'); - final db = [ - Model2Vec.generateEmbedding('A small cat'), - Model2Vec.generateEmbedding('A big dog'), - Model2Vec.generateEmbedding('Space exploration'), - ]; - - final simCat = Model2VecUtils.cosineSimilarity(query, db[0]); - final simSpace = Model2VecUtils.cosineSimilarity(query, db[2]); - stdout - ..writeln( - ' - Sim(kitten, cat): ${(simCat * 100).toStringAsFixed(1)}%', - ) - ..writeln( - ' - Sim(kitten, space): ${(simSpace * 100).toStringAsFixed(1)}%', - ); - - final topMatch = Model2VecUtils.similaritySearchWithScores( - query, - db, - topK: 1, - ); - stdout - ..writeln(' - Best match index: ${topMatch.first.index}') - // 8. Streaming API for Huge Datasets - ..writeln('\n🌊 Streaming API (1000 items):'); - final stream = Stream.fromIterable(List.generate(1000, (i) => 'Item $i')); - final resultStream = Model2Vec.generateEmbeddingStream( - stream, - batchSize: 200, - ); - - var count = 0; - await for (final _ in resultStream) { - count++; - } - stdout - ..writeln( - ' - Successfully streamed and processed $count embeddings.', - ) - ..writeln('\n🎉 All operations completed successfully.'); - } on Model2VecException catch (e) { - stdout.writeln('\n❌ Model2Vec Error: ${e.message}'); - if (e.code != null) { - stdout.writeln(' Error Code: ${e.code}'); - } - } on Object catch (e) { - stdout.writeln('\n💥 Unexpected Error: $e'); - } + // Loads the model on a background isolate, downloading it on first run. + await Model2Vec.loadModelAsync('minishlab/potion-base-2M'); + + final kitten = Model2Vec.generateEmbedding('A cute little kitten'); + final puppy = Model2Vec.generateEmbedding('A small puppy'); + final rocket = Model2Vec.generateEmbedding('Interstellar space travel'); + + // Higher cosine similarity means closer in meaning. + final puppyPct = (Model2VecUtils.cosineSimilarity(kitten, puppy) * 100) + .toStringAsFixed(1); + final rocketPct = (Model2VecUtils.cosineSimilarity(kitten, rocket) * 100) + .toStringAsFixed(1); + + stdout + ..writeln('Embedding dimension: ${kitten.length}') + ..writeln('kitten vs puppy: $puppyPct%') + ..writeln('kitten vs rocket: $rocketPct%'); } diff --git a/example/scaling_example.dart b/example/scaling_example.dart new file mode 100644 index 0000000..8fc788b --- /dev/null +++ b/example/scaling_example.dart @@ -0,0 +1,85 @@ +import 'dart:io'; + +import 'package:model2vec/model2vec.dart'; + +/// Production, at-scale patterns: +/// 1. load with a live download progress bar, +/// 2. batch embedding in a single native call, +/// 3. parallel embedding across CPU cores with an [EmbeddingPool], +/// 4. streaming a huge dataset with bounded memory. +/// +/// Run with `dart run example/scaling_example.dart`. +Future main() async { + // 1. Load with a progress bar. On a cache hit this jumps straight to `Ready`; + // the byte counts only move while the weights actually download. + const modelId = 'minishlab/potion-base-2M'; + stdout.writeln('Loading $modelId…'); + await for (final progress in Model2Vec.loadModelWithProgress(modelId)) { + _renderProgress(progress); + } + stdout.writeln(); + + // 2. Batch embedding — one FFI call, SIMD across the whole batch. + final sentences = [ + 'The quick brown fox.', + 'Jumps over the lazy dog.', + 'Embeddings are just vectors.', + ]; + final batch = Model2Vec.generateBatchEmbeddings(sentences); + stdout.writeln( + 'Batch: ${batch.length} vectors of dim ${batch.first.length}.', + ); + + // 3. Parallel embedding across CPU cores. The pool shares the already-loaded + // process-global model; each batch runs on the least-busy worker. + final pool = await EmbeddingPool.start(size: 4); + try { + final results = await pool.embedBatches([ + ['alpha', 'beta', 'gamma'], + ['delta', 'epsilon'], + ['zeta'], + ]); + final total = results.fold(0, (sum, b) => sum + b.length); + stdout.writeln('Pool: embedded $total texts across ${pool.size} workers.'); + } finally { + await pool.close(); + } + + // 4. Streaming — process a large Stream in batches without holding + // the whole dataset (or all of its vectors) in memory at once. + final huge = Stream.fromIterable(List.generate(1000, (i) => 'Item $i')); + var streamed = 0; + await for (final _ + in Model2Vec.generateEmbeddingStream(huge, batchSize: 200)) { + streamed++; + } + stdout.writeln('Streamed and embedded $streamed items.'); +} + +/// Renders a single load-progress snapshot as an in-place status line. +void _renderProgress(LoadProgress p) { + final line = switch (p.phase) { + LoadPhase.resolving => 'Resolving model files…', + LoadPhase.downloading => _downloadBar(p), + LoadPhase.parsing => 'Parsing & building model…', + LoadPhase.done => 'Ready', + }; + // `\r` rewrites the same line; pad to clear leftovers from a longer one. + stdout.write('\r ${line.padRight(56)}'); +} + +/// A `[███░░░] 42% 3.1/7.5 MB` style bar, or an indeterminate label until the +/// total size is known. +String _downloadBar(LoadProgress p) { + final fraction = p.fraction; + if (fraction == null) { + return 'Downloading…'; + } + const width = 24; + final filled = (fraction * width).round(); + final bar = '${'█' * filled}${'░' * (width - filled)}'; + final mb = (p.bytesDownloaded / (1024 * 1024)).toStringAsFixed(1); + final totalMb = (p.totalBytes / (1024 * 1024)).toStringAsFixed(1); + return '[$bar] ${(fraction * 100).toStringAsFixed(0).padLeft(3)}% ' + '$mb/$totalMb MB'; +} diff --git a/lib/model2vec.dart b/lib/model2vec.dart index fa410c7..a993196 100644 --- a/lib/model2vec.dart +++ b/lib/model2vec.dart @@ -6,6 +6,7 @@ export 'src/chunker.dart' show chunkText; export 'src/embedding_index.dart' show EmbeddingIndex, SearchResult; export 'src/embedding_pool.dart' show EmbeddingPool; export 'src/exception.dart' show Model2VecErrorKind, Model2VecException; +export 'src/load_progress.dart' show LoadPhase, LoadProgress; export 'src/model2vec_base.dart' show Model2Vec; export 'src/model_info.dart' show ModelInfo; export 'src/recommended_model.dart' show RecommendedModel; diff --git a/lib/src/load_progress.dart b/lib/src/load_progress.dart new file mode 100644 index 0000000..fb6c6a9 --- /dev/null +++ b/lib/src/load_progress.dart @@ -0,0 +1,57 @@ +/// The stage a model load is currently in. +/// +/// A load moves forward through these stages: [resolving] → (optionally) +/// [downloading] → [parsing] → [done]. The [downloading] stage is skipped when +/// the model is already cached or loaded from a local path. +enum LoadPhase { + /// Locating the model files — resolving the repo layout and checking the + /// local cache. No bytes have been transferred yet. + resolving, + + /// Downloading the model weights from Hugging Face. This is the only stage + /// with meaningful byte counts. + downloading, + + /// Parsing the files and building the in-memory model. Fast; no byte counts. + parsing, + + /// The model is loaded and ready. Always the final event of the stream. + done, +} + +/// A snapshot of a model load in progress, emitted by +/// `Model2Vec.loadModelWithProgress`. +/// +/// Byte counts are only populated during [LoadPhase.downloading]; they stay `0` +/// when the size is unknown, on a cache hit, or for a local path. Use +/// [fraction] for a progress bar and fall back to an indeterminate spinner when +/// it is `null`. +final class LoadProgress { + /// Creates a load-progress snapshot. + const LoadProgress({ + required this.phase, + this.bytesDownloaded = 0, + this.totalBytes = 0, + }); + + /// The stage the load is currently in. + final LoadPhase phase; + + /// Bytes of the model weights downloaded so far, or `0` when nothing has been + /// downloaded yet (or the model was already cached). + final int bytesDownloaded; + + /// Total size in bytes of the weights being downloaded, or `0` when it is not + /// yet known (before the download starts, on a cache hit, or a local path). + final int totalBytes; + + /// Download completion in `[0.0, 1.0]`, or `null` when the total size is not + /// known — before the download starts, on a cache hit, or for a local path. + double? get fraction => + totalBytes > 0 ? bytesDownloaded / totalBytes : null; + + @override + String toString() => + 'LoadProgress(phase: ${phase.name}, bytesDownloaded: $bytesDownloaded, ' + 'totalBytes: $totalBytes)'; +} diff --git a/lib/src/model2vec_base.dart b/lib/src/model2vec_base.dart index dc2c5e7..4351282 100644 --- a/lib/src/model2vec_base.dart +++ b/lib/src/model2vec_base.dart @@ -8,6 +8,7 @@ import 'package:ffi/ffi.dart'; import 'batcher.dart'; import 'embedding_worker.dart'; import 'exception.dart'; +import 'load_progress.dart'; import 'model2vec_bindings.g.dart' as native; import 'model_info.dart'; import 'recommended_model.dart'; @@ -202,6 +203,94 @@ abstract final class Model2Vec { ), ); + /// Loads a model on a background isolate, reporting progress as a stream of + /// [LoadProgress] snapshots. + /// + /// Like [loadModelAdvancedAsync] the load runs off-thread so the calling + /// isolate stays responsive; this method additionally polls the native load + /// state and yields a snapshot roughly every [pollInterval]. The stream's + /// final event is always [LoadPhase.done]; if the load fails the stream emits + /// that error instead. Options mirror [loadModelAdvanced]. + /// + /// Meaningful byte progress only appears while downloading the weights on a + /// first (uncached) load — an already-cached model or a local path moves + /// straight to [LoadPhase.done] with no byte counts. + /// + /// Only one load may run at a time (the native model is a single process + /// global); do not start another load, or switch/unload the model, while this + /// stream is active. + static Stream loadModelWithProgress( + String modelPath, { + String? hfToken, + String? cacheDirectory, + bool? normalize, + String? subfolder, + Duration pollInterval = const Duration(milliseconds: 100), + }) async* { + // Arm progress on this isolate before the load starts, so a previous load's + // terminal state is never observed as this one's opening event. + native.reset_load_progress(); + + final load = loadModelAdvancedAsync( + modelPath: modelPath, + hfToken: hfToken, + cacheDirectory: cacheDirectory, + normalize: normalize, + subfolder: subfolder, + ); + + // Mirror completion into a flag and a void future the poll loop can race + // against. Errors are swallowed here (re-thrown via `await load` below), so + // this future never completes with an error. + var finished = false; + final done = load.then( + (_) => finished = true, + onError: (_, _) => finished = true, + ); + + while (!finished) { + yield _readLoadProgress(terminal: false); + // Wake on whichever comes first: the next poll tick or load completion. + await Future.any([Future.delayed(pollInterval), done]); + } + + // Surface a load failure to the stream consumer. + await load; + + // Terminal event: the model is loaded and ready. + yield _readLoadProgress(terminal: true); + } + + /// Reads a [LoadProgress] snapshot from native load state. While polling + /// ([terminal] false) a native `done` is clamped to [LoadPhase.parsing] so + /// the only [LoadPhase.done] a consumer observes is the terminal event + /// emitted once the load future has actually completed. + static LoadProgress _readLoadProgress({required bool terminal}) => + using((arena) { + final outPhase = arena(); + final outDownloaded = arena(); + final outTotal = arena(); + native.get_load_progress(outPhase, outDownloaded, outTotal); + + var phase = _loadPhaseFromCode(outPhase.value); + if (!terminal && phase == LoadPhase.done) { + phase = LoadPhase.parsing; + } + return LoadProgress( + phase: phase, + bytesDownloaded: outDownloaded.value, + totalBytes: outTotal.value, + ); + }); + + /// Maps a native load-phase code (see `native/model2vec.h`) to a [LoadPhase]. + static LoadPhase _loadPhaseFromCode(int code) => switch (code) { + 2 => LoadPhase.downloading, + 3 => LoadPhase.parsing, + 4 => LoadPhase.done, + _ => LoadPhase.resolving, // 0 idle, 1 resolving + }; + /// Unloads the active model and frees its native memory. /// /// Safe to call when no model is loaded. After this, [isInitialized] is diff --git a/lib/src/model2vec_bindings.g.dart b/lib/src/model2vec_bindings.g.dart index ad221d3..755b174 100644 --- a/lib/src/model2vec_bindings.g.dart +++ b/lib/src/model2vec_bindings.g.dart @@ -154,3 +154,19 @@ external int is_model_loaded(); external int free_embedder( ffi.Pointer> out_error, ); + +@ffi.Native< + ffi.Void Function( + ffi.Pointer, + ffi.Pointer, + ffi.Pointer, + ) +>() +external void get_load_progress( + ffi.Pointer out_phase, + ffi.Pointer out_downloaded, + ffi.Pointer out_total, +); + +@ffi.Native() +external void reset_load_progress(); diff --git a/native/model2vec.h b/native/model2vec.h index 91b05cb..fa27818 100644 --- a/native/model2vec.h +++ b/native/model2vec.h @@ -78,3 +78,18 @@ void free_floats(float* ptr, size_t len); int is_model_loaded(); int free_embedder(char** out_error); + +/* + * Load progress + * ------------- + * `get_load_progress` reports the current model load as a phase plus download + * byte counts. Phases (out_phase): + * 0 idle 1 resolving 2 downloading 3 parsing 4 done + * `out_downloaded` / `out_total` are 0 when unknown (before a download starts, + * on a cache hit, or for a local path). Any null out-param is skipped. + * `reset_load_progress` arms progress for a new load; call it before starting a + * load so a previous load's terminal state is not observed. Neither fails. + */ +void get_load_progress(int* out_phase, size_t* out_downloaded, size_t* out_total); + +void reset_load_progress(void); diff --git a/native/src/lib.rs b/native/src/lib.rs index c47e421..f7625a1 100644 --- a/native/src/lib.rs +++ b/native/src/lib.rs @@ -358,6 +358,30 @@ pub extern "C" fn free_floats(ptr: *mut f32, len: usize) { } } +/// Writes the current load progress into the out-params. Never fails; any null +/// out-param is skipped. `out_phase` is a load phase code (0 idle, 1 resolving, +/// 2 downloading, 3 parsing, 4 done); the byte counts are 0 when unknown (before +/// a download starts, on a cache hit, or for a local path). +#[no_mangle] +pub extern "C" fn get_load_progress( + out_phase: *mut i32, + out_downloaded: *mut usize, + out_total: *mut usize, +) { + let (phase, downloaded, total) = model::load_progress(); + write_i32(out_phase, phase as i32); + write_usize(out_downloaded, downloaded); + write_usize(out_total, total); +} + +/// Arms progress tracking for a new load. Call this on the polling side before +/// starting the load so a previous load's terminal state isn't observed as the +/// new one's. Never fails. +#[no_mangle] +pub extern "C" fn reset_load_progress() { + model::begin_load(); +} + /// Returns 1 if a model is currently loaded, 0 otherwise. Never fails: a /// poisoned lock is reported as "not loaded". #[no_mangle] diff --git a/native/src/model.rs b/native/src/model.rs index 181705b..7146a55 100644 --- a/native/src/model.rs +++ b/native/src/model.rs @@ -1,16 +1,82 @@ use anyhow::{anyhow, Context, Result}; use half::f16; use hf_hub::api::sync::{ApiBuilder, ApiRepo}; +use hf_hub::api::Progress; +use hf_hub::Cache; use ndarray::{Array2, CowArray, Ix2}; use safetensors::{tensor::Dtype, SafeTensors}; use serde_json::Value; use std::borrow::Cow; +use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering}; use std::{ fs, path::{Path, PathBuf}, }; use tokenizers::Tokenizer; +// --- Load progress ------------------------------------------------------- +// Only one model is ever loading at a time (the model is a process global), so +// progress lives in process-global atomics that any isolate can poll while the +// load itself runs on another. The phase codes below are mirrored on the Dart +// side (see model2vec.h and load_progress.dart). + +/// No load in progress. +pub const PHASE_IDLE: u8 = 0; +/// Locating the model files (resolving repo layout / checking the cache). +pub const PHASE_RESOLVING: u8 = 1; +/// Downloading the model weights from Hugging Face. +pub const PHASE_DOWNLOADING: u8 = 2; +/// Parsing the files and building the in-memory model. +pub const PHASE_PARSING: u8 = 3; +/// The model is loaded and ready. +pub const PHASE_DONE: u8 = 4; + +static LOAD_PHASE: AtomicU8 = AtomicU8::new(PHASE_IDLE); +static LOAD_DOWNLOADED: AtomicUsize = AtomicUsize::new(0); +static LOAD_TOTAL: AtomicUsize = AtomicUsize::new(0); + +/// Arms progress for a fresh load: zeroes the byte counters and enters the +/// resolving phase. Called at the start of a load and by the caller before it, +/// so a previous load's terminal state is never observed as the new one's. +pub fn begin_load() { + LOAD_DOWNLOADED.store(0, Ordering::Relaxed); + LOAD_TOTAL.store(0, Ordering::Relaxed); + LOAD_PHASE.store(PHASE_RESOLVING, Ordering::Relaxed); +} + +fn set_phase(phase: u8) { + LOAD_PHASE.store(phase, Ordering::Relaxed); +} + +/// Snapshot of the current load as `(phase, downloaded_bytes, total_bytes)`. +pub fn load_progress() -> (u8, usize, usize) { + ( + LOAD_PHASE.load(Ordering::Relaxed), + LOAD_DOWNLOADED.load(Ordering::Relaxed), + LOAD_TOTAL.load(Ordering::Relaxed), + ) +} + +/// hf-hub download sink that writes progress into the global load state. +struct AtomicProgress; + +impl Progress for AtomicProgress { + fn init(&mut self, size: usize, _filename: &str) { + LOAD_TOTAL.store(size, Ordering::Relaxed); + LOAD_DOWNLOADED.store(0, Ordering::Relaxed); + LOAD_PHASE.store(PHASE_DOWNLOADING, Ordering::Relaxed); + } + + fn update(&mut self, size: usize) { + LOAD_DOWNLOADED.fetch_add(size, Ordering::Relaxed); + } + + fn finish(&mut self) { + let total = LOAD_TOTAL.load(Ordering::Relaxed); + LOAD_DOWNLOADED.store(total, Ordering::Relaxed); + } +} + /// Static embedding model for Model2Vec #[derive(Debug, Clone)] pub struct StaticModel { @@ -68,6 +134,8 @@ fn is_not_found(e: &hf_hub::api::sync::ApiError) -> bool { fn match_hub_layout( repo: &ApiRepo, + cache: &Cache, + repo_id: &str, config_prefix: &str, model_prefix: &str, config_file: &str, @@ -81,10 +149,31 @@ fn match_hub_layout( }; let Some(config) = fetch(format!("{config_prefix}{config_file}"))? else { return Ok(None); }; let Some(tokenizer) = fetch(format!("{model_prefix}tokenizer.json"))? else { return Ok(None); }; - let Some(model) = fetch(format!("{model_prefix}model.safetensors"))? else { return Ok(None); }; + let model_file = format!("{model_prefix}model.safetensors"); + let Some(model) = fetch_weights_with_progress(repo, cache, repo_id, model_file)? else { return Ok(None); }; Ok(Some(ModelFiles { tokenizer, model, config })) } +/// Fetches the (large) weights file, reporting download progress on a cache +/// miss. A cache hit returns immediately with no download and no progress — +/// mirroring `ApiRepo::get`, which we can't reuse because it downloads without +/// a progress hook. +fn fetch_weights_with_progress( + repo: &ApiRepo, + cache: &Cache, + repo_id: &str, + filename: String, +) -> Result> { + if let Some(path) = cache.model(repo_id.to_owned()).get(&filename) { + return Ok(Some(path)); + } + match repo.download_with_progress(&filename, AtomicProgress) { + Ok(path) => Ok(Some(path)), + Err(e) if is_not_found(&e) => Ok(None), + Err(e) => Err(e.into()), + } +} + fn resolve_local_model_files(folder: &Path) -> Option { match_local_layout(folder, folder, "config.json") .or_else(|| match_local_layout(folder, folder, "config_sentence_transformers.json")) @@ -92,7 +181,7 @@ fn resolve_local_model_files(folder: &Path) -> Option { .or_else(|| folder.parent().and_then(|p| match_local_layout(p, folder, "config_sentence_transformers.json"))) } -fn resolve_hub_model_files(repo: &ApiRepo, prefix: &str) -> Result { +fn resolve_hub_model_files(repo: &ApiRepo, cache: &Cache, repo_id: &str, prefix: &str) -> Result { let sub_prefix = format!("{prefix}0_StaticEmbedding/"); let trimmed = prefix.trim_end_matches('/'); let parent = match Path::new(trimmed).parent() { @@ -100,10 +189,10 @@ fn resolve_hub_model_files(repo: &ApiRepo, prefix: &str) -> Result { _ => String::new(), }; - if let Some(f) = match_hub_layout(repo, prefix, prefix, "config.json")? { return Ok(f); } - if let Some(f) = match_hub_layout(repo, prefix, prefix, "config_sentence_transformers.json")? { return Ok(f); } - if let Some(f) = match_hub_layout(repo, prefix, &sub_prefix, "config_sentence_transformers.json")? { return Ok(f); } - match_hub_layout(repo, &parent, prefix, "config_sentence_transformers.json")? + if let Some(f) = match_hub_layout(repo, cache, repo_id, prefix, prefix, "config.json")? { return Ok(f); } + if let Some(f) = match_hub_layout(repo, cache, repo_id, prefix, prefix, "config_sentence_transformers.json")? { return Ok(f); } + if let Some(f) = match_hub_layout(repo, cache, repo_id, prefix, &sub_prefix, "config_sentence_transformers.json")? { return Ok(f); } + match_hub_layout(repo, cache, repo_id, &parent, prefix, "config_sentence_transformers.json")? .ok_or_else(|| anyhow!("no valid model layout found in '{prefix}'")) } @@ -169,11 +258,15 @@ impl StaticModel { normalize: Option, subfolder: Option<&str>, ) -> Result { + begin_load(); let files = resolve_model_files(repo_or_path, token, cache_dir, subfolder)?; + set_phase(PHASE_PARSING); let tokenizer_bytes = fs::read(&files.tokenizer).context("failed to read tokenizer.json")?; let model_bytes = fs::read(&files.model).context("failed to read model.safetensors")?; let config_bytes = fs::read(&files.config).context("failed to read config.json")?; - Self::from_bytes(tokenizer_bytes, model_bytes, config_bytes, normalize) + let model = Self::from_bytes(tokenizer_bytes, model_bytes, config_bytes, normalize)?; + set_phase(PHASE_DONE); + Ok(model) } pub fn from_owned( @@ -350,15 +443,22 @@ fn resolve_model_files>( } fn download_model_files(repo_id: &str, token: Option<&str>, cache_dir: Option<&Path>, subfolder: Option<&str>) -> Result { - let mut builder = ApiBuilder::new(); + // Own the cache so we can both drive the hf-hub API and check it directly + // for a weights cache hit (matching ApiBuilder: default location, or the + // caller's cache_dir). + let cache = match cache_dir { + Some(path) => Cache::new(path.to_path_buf()), + None => Cache::default(), + }; + let mut builder = ApiBuilder::from_cache(cache.clone()); if let Some(tok) = token { builder = builder.with_token(Some(tok.to_string())); } - if let Some(path) = cache_dir { builder = builder.with_cache_dir(path.to_path_buf()); } - + let result = (|| { let api = builder.build().context("hf-hub API init failed")?; let repo = api.model(repo_id.to_owned()); let prefix = subfolder.map(|s| format!("{s}/")).unwrap_or_default(); - resolve_hub_model_files(&repo, &prefix).with_context(|| format!("could not load '{repo_id}' from HF")) + resolve_hub_model_files(&repo, &cache, repo_id, &prefix) + .with_context(|| format!("could not load '{repo_id}' from HF")) })(); result } diff --git a/test/load_progress_test.dart b/test/load_progress_test.dart new file mode 100644 index 0000000..8eab08a --- /dev/null +++ b/test/load_progress_test.dart @@ -0,0 +1,79 @@ +import 'dart:io'; + +import 'package:model2vec/model2vec.dart'; +import 'package:test/test.dart'; + +void main() { + group('loadModelWithProgress', () { + test('reports download progress and reaches done on a fresh cache', + () async { + // A private temp cache guarantees a cache miss, so the weights actually + // download and the downloading phase is exercised. + final tmp = Directory.systemTemp.createTempSync('m2v_progress_'); + addTearDown(() => tmp.deleteSync(recursive: true)); + + final events = []; + await for (final p in Model2Vec.loadModelWithProgress( + 'minishlab/potion-base-2M', + cacheDirectory: tmp.path, + pollInterval: const Duration(milliseconds: 1), + )) { + events.add(p); + } + + // The load succeeded and the stream ended on `done`. + expect(Model2Vec.isInitialized, isTrue); + expect(events, isNotEmpty); + expect(events.last.phase, LoadPhase.done); + + // `done` is only ever the terminal event. + final nonTerminal = events.sublist(0, events.length - 1); + expect( + nonTerminal, + everyElement( + isNot( + predicate((e) => e.phase == LoadPhase.done), + ), + ), + ); + + // A fresh cache must download the weights: we should have caught at least + // one downloading snapshot with a known total and a real fraction. + final downloading = + events.where((e) => e.phase == LoadPhase.downloading).toList(); + expect( + downloading, + isNotEmpty, + reason: 'a fresh cache should download the weights', + ); + final last = downloading.last; + expect(last.totalBytes, greaterThan(0)); + expect(last.bytesDownloaded, greaterThan(0)); + expect(last.fraction, isNotNull); + expect(last.fraction, inInclusiveRange(0.0, 1.0)); + + // Print the shape of the run for eyeballing. + final byPhase = {}; + for (final e in events) { + byPhase[e.phase] = (byPhase[e.phase] ?? 0) + 1; + } + // Diagnostic output for eyeballing a real run. + // ignore: avoid_print + print('phase counts: $byPhase; last downloading: $last'); + }); + + test('a cached model streams straight to done', () async { + // The prior test cached the model in a temp dir, but the default cache + // may already hold it too; either way this must terminate on `done`. + final events = []; + await for (final p in Model2Vec.loadModelWithProgress( + 'minishlab/potion-base-2M', + )) { + events.add(p); + } + + expect(Model2Vec.isInitialized, isTrue); + expect(events.last.phase, LoadPhase.done); + }); + }); +} From 983e81570a8976d90b049a7e357a81bf4d3478e3 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Tue, 7 Jul 2026 13:10:58 +0300 Subject: [PATCH 11/20] test improvements --- lib/src/channel.dart | 13 +- test/channel_test.dart | 79 +++ test/chunker_test.dart | 9 + test/embedding_index_test.dart | 69 +- test/embedding_pool_test.dart | 188 ++++-- test/embedding_test.dart | 222 +++++++ test/embedding_worker_test.dart | 173 ++--- ...ror_path_test.dart => exception_test.dart} | 24 - test/init_async_test.dart | 46 -- test/lifecycle_test.dart | 37 -- test/load_progress_test.dart | 79 --- test/model2vec_test.dart | 204 ------ test/model_loading_test.dart | 265 ++++++++ test/recommended_models_test.dart | 78 +++ test/support.dart | 20 + test/utils_test.dart | 600 ++++++++++-------- test/worker_protocol_test.dart | 242 ++++--- test/worker_support.dart | 177 ++++++ 18 files changed, 1593 insertions(+), 932 deletions(-) create mode 100644 test/channel_test.dart create mode 100644 test/embedding_test.dart rename test/{error_path_test.dart => exception_test.dart} (64%) delete mode 100644 test/init_async_test.dart delete mode 100644 test/lifecycle_test.dart delete mode 100644 test/load_progress_test.dart delete mode 100644 test/model2vec_test.dart create mode 100644 test/model_loading_test.dart create mode 100644 test/recommended_models_test.dart create mode 100644 test/support.dart create mode 100644 test/worker_support.dart diff --git a/lib/src/channel.dart b/lib/src/channel.dart index b3dac33..c243609 100644 --- a/lib/src/channel.dart +++ b/lib/src/channel.dart @@ -103,7 +103,10 @@ class IsolateChannel implements Channel { receivePort.close(); exitPort.close(); if (!incoming.isClosed) { - await incoming.close(); + // No one has listened to [incoming] yet (start hasn't returned), and a + // single-subscription controller's close() future never completes + // without a listener — so fire-and-forget instead of awaiting a hang. + unawaited(incoming.close()); } isolate.kill(priority: Isolate.immediate); rethrow; @@ -141,7 +144,13 @@ class IsolateChannel implements Channel { _receivePort.close(); _exitPort.close(); if (!_incoming.isClosed) { - await _incoming.close(); + // Awaiting close() only completes once a listener drains the done event; + // an unlistened single-subscription controller would hang here forever. + if (_incoming.hasListener) { + await _incoming.close(); + } else { + unawaited(_incoming.close()); + } } } } diff --git a/test/channel_test.dart b/test/channel_test.dart new file mode 100644 index 0000000..7dc55d0 --- /dev/null +++ b/test/channel_test.dart @@ -0,0 +1,79 @@ +import 'dart:async'; +import 'dart:typed_data'; + +import 'package:model2vec/src/channel.dart'; +import 'package:test/test.dart'; + +import 'worker_support.dart'; + +void main() { + group('IsolateChannel', () { + test('delivers worker replies over incoming', () async { + final channel = await IsolateChannel.start(echoEntryPoint); + addTearDown(channel.close); + + final reply = Completer(); + channel.incoming.listen(reply.complete); + channel.send((batch: ['ab', 'c'], maxLength: 8)); + + final vectors = (await reply.future)! as List; + expect(vectors.map((v) => v.first).toList(), [2.0, 1.0]); + }); + + test('start fails when the worker exits before the handshake', () async { + await expectLater( + IsolateChannel.start(noHandshakeEntryPoint), + throwsA( + isA().having( + (e) => e.message, + 'message', + contains('exited during startup'), + ), + ), + ); + }, timeout: const Timeout(Duration(seconds: 20))); + + test( + 'start rejects a non-SendPort handshake instead of hanging', + () async { + await expectLater( + IsolateChannel.start(nonSendPortHandshakeEntryPoint), + throwsA( + isA().having( + (e) => e.message, + 'message', + contains('did not hand back a SendPort'), + ), + ), + ); + }, + // Regression guard: a bad-handshake message must reject, not hang. + // start()'s cleanup used to await _incoming.close() on a single- + // subscription controller that was never listened (its close future never + // completes) — now fire-and-forget. + timeout: const Timeout(Duration(seconds: 20)), + ); + + test('incoming closes when the worker exits', () async { + final channel = await IsolateChannel.start(dyingEntryPoint); + addTearDown(channel.close); + + final done = Completer(); + channel.incoming.listen((_) {}, onDone: done.complete); + // The dying worker closes its port (exiting) instead of replying; the + // exit must surface as incoming's onDone so a protocol above can't hang. + channel.send((batch: ['a'], maxLength: 8)); + + await done.future; + }, timeout: const Timeout(Duration(seconds: 20))); + + test('close is idempotent', () async { + final channel = await IsolateChannel.start(echoEntryPoint); + // Listen as a real caller (WorkerProtocol) does; close() awaits the + // incoming controller's own close, which only completes once listened. + channel.incoming.listen((_) {}); + await channel.close(); + await channel.close(); + }, timeout: const Timeout(Duration(seconds: 20))); + }); +} diff --git a/test/chunker_test.dart b/test/chunker_test.dart index b79db30..bc8773d 100644 --- a/test/chunker_test.dart +++ b/test/chunker_test.dart @@ -42,6 +42,15 @@ void main() { } }); + test('a word longer than maxChars becomes its own intact chunk', () { + final longWord = 'a' * 50; + final chunks = chunkText('short $longWord', maxChars: 20, overlap: 0); + + expect(chunks, ['short', longWord]); + // The oversized word is kept whole, exceeding maxChars rather than split. + expect(chunks.last.length, 50); + }); + test('validates arguments', () { expect(() => chunkText('a', maxChars: 0), throwsArgumentError); expect( diff --git a/test/embedding_index_test.dart b/test/embedding_index_test.dart index 7f8cc0c..0542821 100644 --- a/test/embedding_index_test.dart +++ b/test/embedding_index_test.dart @@ -57,14 +57,21 @@ void main() { expect(EmbeddingIndex().search(_v([1, 0])), isEmpty); }); - test('addAll and overwrite by id', () { + test('addAll and overwrite by id replaces the stored vector', () { final index = EmbeddingIndex() ..addAll({'a': _v([1, 0]), 'b': _v([0, 1])}); expect(index.length, 2); - index.add('a', _v([0, 1])); // overwrite - expect(index.length, 2); - expect(index.search(_v([0, 1]), topK: 1).first.id, anyOf('a', 'b')); + index.add('a', _v([0, 1])); // overwrite a's vector [1,0] -> [0,1] + expect(index.length, 2); // no new entry, still 2 + + double scoreOf(String id, Float32List query) => + index.search(query, topK: 2).firstWhere((h) => h.id == id).score; + + // The overwrite took effect: 'a' now matches its new vector and no + // longer matches the old one. + expect(scoreOf('a', _v([0, 1])), closeTo(1.0, 1e-6)); + expect(scoreOf('a', _v([1, 0])), closeTo(0.0, 1e-6)); }); test('dimension mismatch throws on add and on query', () { @@ -130,12 +137,14 @@ void main() { expect(restored.contains(longId), isTrue); }); - test('fromBytes rejects a high-byte / truncated blob with ArgumentError', - () { + test('fromBytes rejects a blob with a bad magic byte', () { expect( () => EmbeddingIndex.fromBytes(Uint8List(14)..[0] = 0xFF), throwsArgumentError, ); + }); + + test('fromBytes rejects a truncated blob', () { final full = (EmbeddingIndex()..add('a', _v([1, 2, 3, 4]))).toBytes(); expect( () => EmbeddingIndex.fromBytes(full.sublist(0, full.length - 2)), @@ -143,6 +152,40 @@ void main() { ); }); + test('fromBytes rejects an unsupported version byte', () { + // Magic-valid header, but version 3 is beyond what the decoder reads. + final blob = Uint8List.fromList([ + 0x4D, 0x32, 0x56, 0x49, // magic 'M2VI' + 0x03, // version 3 (unsupported) + 0x00, // flags: not quantized + 0x00, 0x00, 0x00, 0x00, // dim = 0 + 0x00, 0x00, 0x00, 0x00, // count = 0 + ]); + expect( + () => EmbeddingIndex.fromBytes(blob), + throwsA( + isA().having( + (e) => e.message, + 'message', + contains('unsupported EmbeddingIndex version'), + ), + ), + ); + }); + + test('removing the last entry resets the dimension to null', () { + final index = EmbeddingIndex()..add('a', _v([1, 2, 3])); + expect(index.dimension, 3); + + expect(index.remove('a'), isTrue); + expect(index.isEmpty, isTrue); + expect(index.dimension, isNull); + + // A cleared dimension lets a different-dim vector be added afresh. + index.add('b', _v([1, 0])); + expect(index.dimension, 2); + }); + test('negative topK returns empty', () { final index = EmbeddingIndex()..add('a', _v([1, 0])); expect(index.search(_v([1, 0]), topK: -1), isEmpty); @@ -191,6 +234,20 @@ void main() { expect(byId['none'], isNull); }); + test('a multibyte UTF-8 id and payload round-trip exactly', () { + // Guards the byte-vs-char offset math: both id and payload contain + // multibyte code points (CJK + emoji). + const id = '文書-😀'; + const payload = '説明: café ☕'; + final index = EmbeddingIndex()..add(id, _v([1, 0]), payload: payload); + + final restored = EmbeddingIndex.fromBytes(index.toBytes()); + expect(restored.contains(id), isTrue); + final hit = restored.search(_v([1, 0])).single; + expect(hit.id, id); + expect(hit.payload, payload); + }); + test('quantized index carries payloads too', () { final index = EmbeddingIndex(quantized: true) ..add('a', _v([1, 0, 0]), payload: 'alpha'); diff --git a/test/embedding_pool_test.dart b/test/embedding_pool_test.dart index 6f8b326..ed27d84 100644 --- a/test/embedding_pool_test.dart +++ b/test/embedding_pool_test.dart @@ -1,70 +1,140 @@ -import 'dart:isolate'; -import 'dart:typed_data'; - import 'package:model2vec/model2vec.dart'; -import 'package:model2vec/src/worker_protocol.dart'; import 'package:test/test.dart'; -/// Fake worker entry point: one vector per input text, whose only element is -/// the text's length. No model, no native calls. -void echoEntryPoint(SendPort mainSendPort) { - final receivePort = ReceivePort(); - mainSendPort.send(receivePort.sendPort); - receivePort.listen((message) { - if (message == 'close') { - receivePort.close(); - return; - } - final request = message as EmbedRequest; - mainSendPort.send([ - for (final text in request.batch) - Float32List.fromList([text.length.toDouble()]), - ]); - }); -} +import 'worker_support.dart'; void main() { - test('starts with the requested number of workers', () async { - final pool = await EmbeddingPool.start(size: 3, entryPoint: echoEntryPoint); - expect(pool.size, 3); - await pool.close(); - }); + group('EmbeddingPool', () { + test('starts with the requested number of workers', () async { + final pool = await EmbeddingPool.start( + size: 3, + entryPoint: echoEntryPoint, + ); + addTearDown(pool.close); + expect(pool.size, 3); + }); - test('clamps a non-positive size to at least one worker', () async { - final pool = await EmbeddingPool.start(size: 0, entryPoint: echoEntryPoint); - expect(pool.size, 1); - await pool.close(); - }); + test('clamps a non-positive size to at least one worker', () async { + final pool = await EmbeddingPool.start( + size: 0, + entryPoint: echoEntryPoint, + ); + addTearDown(pool.close); + expect(pool.size, 1); + }); - test('embedBatch returns results', () async { - final pool = await EmbeddingPool.start(size: 2, entryPoint: echoEntryPoint); - final result = await pool.embedBatch(['a', 'bb']); - expect(result.map((v) => v.first).toList(), [1.0, 2.0]); - await pool.close(); - }); + test('embedBatch returns results', () async { + final pool = await EmbeddingPool.start( + size: 2, + entryPoint: echoEntryPoint, + ); + addTearDown(pool.close); + final result = await pool.embedBatch(['a', 'bb']); + expect(result.map((v) => v.first).toList(), [1.0, 2.0]); + }); - test('embedBatches preserves input order under parallelism', () async { - final pool = await EmbeddingPool.start(size: 3, entryPoint: echoEntryPoint); - final results = await pool.embedBatches([ - ['a'], - ['bb'], - ['ccc'], - ['dddd'], - ['eeeee'], - ]); - expect( - results.map((r) => r.first.first).toList(), - [1.0, 2.0, 3.0, 4.0, 5.0], - ); - await pool.close(); - }); + test('embedBatches preserves input order under parallelism', () async { + final pool = await EmbeddingPool.start( + size: 3, + entryPoint: echoEntryPoint, + ); + addTearDown(pool.close); + final results = await pool.embedBatches([ + ['a'], + ['bb'], + ['ccc'], + ['dddd'], + ['eeeee'], + ]); + expect( + results.map((r) => r.first.first).toList(), + [1.0, 2.0, 3.0, 4.0, 5.0], + ); + }); + + test('spreads concurrent batches across distinct workers', () async { + final pool = await EmbeddingPool.start( + size: 3, + entryPoint: countingEntryPoint, + ); + addTearDown(pool.close); + + // Fire one batch per worker at once. Least-busy dispatch should hand each + // to a different worker, so every reply is that worker's *first* request. + final results = await Future.wait([ + for (var i = 0; i < pool.size; i++) pool.embedBatch(['x']), + ]); + + final counts = results.map((r) => r.single.first).toList(); + expect(counts, everyElement(1.0)); + expect(counts, hasLength(pool.size)); + }); + + test('close rejects in-flight and queued batches without hanging', + () async { + final pool = await EmbeddingPool.start( + size: 2, + entryPoint: countingEntryPoint, + ); + + // Launch more batches than workers so some are queued behind others, then + // close before any of the (delayed) replies land. + final batches = [for (var i = 0; i < 5; i++) pool.embedBatch(['a'])]; + final expectations = [ + for (final batch in batches) + expectLater(batch, throwsA(isA())), + ]; + + await pool.close(); + await Future.wait(expectations); + }, timeout: const Timeout(Duration(seconds: 20))); + + test('recovers after a failing batch and serves later batches', () async { + // A single worker fails its first request, then echoes. If the pool + // failed to decrement its in-flight counter on the error path, later + // batches would still have to land — assert they do. + final pool = await EmbeddingPool.start( + size: 1, + entryPoint: failThenEchoEntryPoint, + ); + addTearDown(pool.close); + + await expectLater( + pool.embedBatch(['a']), + throwsA(isA()), + ); + + final result = await pool.embedBatch(['abcd']); + expect(result.single.first, 4.0); + }); + + test('start throws and cleans up when a later worker fails to spawn', + () async { + resetPartialStartSlots(); + addTearDown(resetPartialStartSlots); + + // The pool is larger than the number of spawns the entry point lets + // through, so some workers start and at least one fails. start() must + // close the survivors and throw rather than leak orphaned isolates. + await expectLater( + EmbeddingPool.start( + size: partialStartSuccessLimit + 1, + entryPoint: flakyStartEntryPoint, + ), + throwsA(isA()), + ); + }, timeout: const Timeout(Duration(seconds: 20))); - test('close tears down all workers', () async { - final pool = await EmbeddingPool.start(size: 2, entryPoint: echoEntryPoint); - await pool.close(); - await expectLater( - pool.embedBatch(['a']), - throwsA(isA()), - ); + test('close tears down all workers', () async { + final pool = await EmbeddingPool.start( + size: 2, + entryPoint: echoEntryPoint, + ); + await pool.close(); + await expectLater( + pool.embedBatch(['a']), + throwsA(isA()), + ); + }); }); } diff --git a/test/embedding_test.dart b/test/embedding_test.dart new file mode 100644 index 0000000..755a15a --- /dev/null +++ b/test/embedding_test.dart @@ -0,0 +1,222 @@ +import 'package:model2vec/model2vec.dart'; +import 'package:test/test.dart'; + +import 'support.dart'; + +void main() { + setUpAll(() => Model2Vec.loadModel(testModelId)); + + group('tokenize', () { + test('breaks text into lowercased tokens', () { + final tokens = Model2Vec.tokenize('Dart FFI is powerful'); + expect(tokens, isNotEmpty); + expect(tokens, contains('dart')); + }); + + test('returns an empty list for empty input', () { + expect(Model2Vec.tokenize(''), isEmpty); + }); + }); + + group('generateEmbedding', () { + test('returns a model-dimension vector with the expected snapshot', () { + final vector = Model2Vec.generateEmbedding('Hello world'); + expect(vector.length, Model2Vec.embeddingDimension); + expect(vector.length, testModelDim); + + // Exact values for Potion Base 2M on "Hello world"; closeTo absorbs + // minor floating-point differences across platforms. + expect(vector[0], closeTo(-0.06458017975091934, 1e-6)); + expect(vector[1], closeTo(-0.15089768171310425, 1e-6)); + expect(vector[2], closeTo(-0.27086204290390015, 1e-6)); + }); + + test('empty input returns a full-length vector rather than throwing', () { + // Pinned contract: an empty string yields a dimension-length embedding; + // it does not raise emptyResult. + final vector = Model2Vec.generateEmbedding(''); + expect(vector.length, testModelDim); + }); + }); + + group('generateBatchEmbeddings', () { + test('an empty list returns an empty list', () { + expect(Model2Vec.generateBatchEmbeddings([]), isEmpty); + }); + + test('is consistent with a single embedding', () { + const text = 'Consistency check'; + final single = Model2Vec.generateEmbedding(text); + final batch = Model2Vec.generateBatchEmbeddings([text]); + + expect(batch.length, 1); + expect(batch[0], orderedEquals(single)); + }); + + test('handles multiple diverse strings', () { + final texts = [ + 'Short', + 'A much longer sentence to test truncation logic in the underlying ', + '12345 !? @#', + 'Теж має працювати', + ]; + final results = Model2Vec.generateBatchEmbeddings(texts); + expect(results.length, texts.length); + for (final v in results) { + expect(v.length, testModelDim); + } + }); + + test('stitches native chunks for more than 1024 texts', () { + // 2000 exceeds the internal FFI chunk size (1024), so this exercises the + // native multi-chunk stitching. + final texts = List.generate(2000, (i) => 'row number $i'); + final results = Model2Vec.generateBatchEmbeddings(texts); + expect(results.length, 2000); + for (final v in results) { + expect(v.length, testModelDim); + } + }); + }); + + group('async embeddings', () { + test('generateEmbeddingAsync runs in a background isolate', () async { + final vector = await Model2Vec.generateEmbeddingAsync('Async test'); + expect(vector.length, testModelDim); + }); + + test('generateBatchEmbeddingsAsync runs in a background isolate', () async { + final results = await Model2Vec.generateBatchEmbeddingsAsync([ + 'Async 1', + 'Async 2', + ]); + expect(results.length, 2); + expect(results[0].length, testModelDim); + }); + }); + + group('generateEmbeddingStream', () { + test('processes a stream of texts across multiple batches', () async { + const totalItems = 2500; + final stream = Stream.fromIterable( + Iterable.generate(totalItems, (i) => 'test string number $i'), + ); + + final results = await Model2Vec.generateEmbeddingStream( + stream, + batchSize: 1000, + ).toList(); + + expect(results.length, totalItems); + expect(results.first.length, testModelDim); + }); + + test('works with useIsolate: false', () async { + final stream = Stream.fromIterable(['Text 1', 'Text 2', 'Text 3']); + final results = await Model2Vec.generateEmbeddingStream( + stream, + batchSize: 2, + useIsolate: false, + ).toList(); + expect(results.length, 3); + }); + + test('cancels the subscription cleanly when taking a prefix', () async { + final stream = Stream.fromIterable([ + 'Text 1', + 'Text 2', + 'Text 3', + 'Text 4', + ]); + // take(2) cancels the subscription early. + final firstTwo = await Model2Vec.generateEmbeddingStream( + stream, + batchSize: 2, + ).take(2).toList(); + expect(firstTwo.length, 2); + }); + + test('an empty stream yields nothing and completes (isolate)', () async { + final results = await Model2Vec.generateEmbeddingStream( + const Stream.empty(), + ).toList(); + expect(results, isEmpty); + }); + + test('an empty stream yields nothing and completes (no isolate)', () async { + final results = await Model2Vec.generateEmbeddingStream( + const Stream.empty(), + useIsolate: false, + ).toList(); + expect(results, isEmpty); + }); + }); + + group('maxLength', () { + test('truncation applies to single and batch embeddings', () { + const text = + 'A very long sentence to test truncation with maxLength parameter'; + // With a very short maxLength the embedding differs from the default. + final defaultEmbedding = Model2Vec.generateEmbedding(text); + final shortEmbedding = Model2Vec.generateEmbedding(text, maxLength: 2); + expect(defaultEmbedding, isNot(orderedEquals(shortEmbedding))); + + final batch = Model2Vec.generateBatchEmbeddings( + [text, text], + maxLength: 2, + ); + expect(batch.length, 2); + expect(batch[0], orderedEquals(shortEmbedding)); + expect(batch[1], orderedEquals(shortEmbedding)); + }); + }); + + group('semantic sanity', () { + test('related sentences embed closer than unrelated ones', () { + final cat = Model2Vec.generateEmbedding('A small cute cat'); + final kitten = Model2Vec.generateEmbedding('A young little kitten'); + final space = Model2Vec.generateEmbedding( + 'The exploration of outer space and planets', + ); + + final catKitten = Model2VecUtils.cosineSimilarity(cat, kitten); + final catSpace = Model2VecUtils.cosineSimilarity(cat, space); + + // Semantically related things should be more similar. + expect(catKitten, greaterThan(catSpace)); + expect(catKitten, greaterThan(0.5)); + expect(catSpace, lessThan(0.5)); + }); + }); + + group('notInitialized', () { + test('generateEmbedding and tokenize throw after unloadModel', () { + // Restore the model for the rest of the suite, even if an expect fails. + addTearDown(() => Model2Vec.loadModel(testModelId)); + + Model2Vec.unloadModel(); + expect(Model2Vec.isInitialized, isFalse); + + expect( + () => Model2Vec.generateEmbedding('x'), + throwsA( + isA().having( + (e) => e.kind, + 'kind', + Model2VecErrorKind.notInitialized, + ), + ), + ); + expect( + () => Model2Vec.tokenize('x'), + throwsA( + isA().having( + (e) => e.kind, + 'kind', + Model2VecErrorKind.notInitialized, + ), + ), + ); + }); + }); +} diff --git a/test/embedding_worker_test.dart b/test/embedding_worker_test.dart index f051082..9f08c4a 100644 --- a/test/embedding_worker_test.dart +++ b/test/embedding_worker_test.dart @@ -1,123 +1,78 @@ -import 'dart:isolate'; -import 'dart:typed_data'; - import 'package:model2vec/src/embedding_worker.dart'; import 'package:model2vec/src/exception.dart'; -import 'package:model2vec/src/worker_protocol.dart'; import 'package:test/test.dart'; -/// Fake worker entry point: replies with one vector per input text, whose only -/// element is the text's length. No model, no native calls. -void echoEntryPoint(SendPort mainSendPort) { - final receivePort = ReceivePort(); - mainSendPort.send(receivePort.sendPort); - receivePort.listen((message) { - if (message == 'close') { - receivePort.close(); - return; - } - final request = message as EmbedRequest; - final results = [ - for (final text in request.batch) - Float32List.fromList([text.length.toDouble()]), - ]; - mainSendPort.send(results); - }); -} - -/// Fake worker entry point that always replies with a typed error. -void errorEntryPoint(SendPort mainSendPort) { - final receivePort = ReceivePort(); - mainSendPort.send(receivePort.sendPort); - receivePort.listen((message) { - if (message == 'close') { - receivePort.close(); - return; - } - mainSendPort.send( - const Model2VecException( - Model2VecErrorKind.tokenizationFailed, - 'boom', - 6, - ), - ); - }); -} - -/// Fake entry point that dies (closes its port, exiting the isolate) on the -/// first request without ever replying. -void dyingEntryPoint(SendPort mainSendPort) { - final receivePort = ReceivePort(); - mainSendPort.send(receivePort.sendPort); - receivePort.listen((message) { - if (message == 'close') { - receivePort.close(); - return; - } - receivePort.close(); // exit without replying - }); -} - -/// Fake entry point that exits immediately, never completing the handshake. -void noHandshakeEntryPoint(SendPort mainSendPort) { - // Returns at once; the isolate exits before sending its SendPort. -} +import 'worker_support.dart'; void main() { - test('embedBatch returns results from the worker isolate', () async { - final worker = await EmbeddingWorker.start(entryPoint: echoEntryPoint); - final result = await worker.embedBatch(['a', 'bb', 'ccc']); - expect(result.map((v) => v.first).toList(), [1.0, 2.0, 3.0]); - await worker.close(); - }); + group('EmbeddingWorker', () { + test('embedBatch returns results from the worker isolate', () async { + final worker = await EmbeddingWorker.start(entryPoint: echoEntryPoint); + addTearDown(worker.close); + final result = await worker.embedBatch(['a', 'bb', 'ccc']); + expect(result.map((v) => v.first).toList(), [1.0, 2.0, 3.0]); + }); - test('multiple batches preserve order over the real isolate', () async { - final worker = await EmbeddingWorker.start(entryPoint: echoEntryPoint); - final r1 = worker.embedBatch(['x']); - final r2 = worker.embedBatch(['yy']); - expect((await r1).first.first, 1.0); - expect((await r2).first.first, 2.0); - await worker.close(); - }); + test('multiple batches preserve order over the real isolate', () async { + final worker = await EmbeddingWorker.start(entryPoint: echoEntryPoint); + addTearDown(worker.close); + final r1 = worker.embedBatch(['x']); + final r2 = worker.embedBatch(['yy']); + expect((await r1).first.first, 1.0); + expect((await r2).first.first, 2.0); + }); + + test('an empty batch round-trips as an empty list', () async { + final worker = await EmbeddingWorker.start(entryPoint: echoEntryPoint); + addTearDown(worker.close); + expect(await worker.embedBatch([]), isEmpty); + }); - test('a typed error survives the isolate boundary', () async { - final worker = await EmbeddingWorker.start(entryPoint: errorEntryPoint); - await expectLater( - worker.embedBatch(['a']), - throwsA( - isA().having( - (e) => e.kind, - 'kind', - Model2VecErrorKind.tokenizationFailed, + test('a typed error survives the isolate boundary', () async { + final worker = await EmbeddingWorker.start(entryPoint: errorEntryPoint); + addTearDown(worker.close); + await expectLater( + worker.embedBatch(['a']), + throwsA( + isA().having( + (e) => e.kind, + 'kind', + Model2VecErrorKind.tokenizationFailed, + ), ), - ), - ); - await worker.close(); - }); + ); + }); - test('close tears the worker down and later calls fail fast', () async { - final worker = await EmbeddingWorker.start(entryPoint: echoEntryPoint); - await worker.close(); - await expectLater( - worker.embedBatch(['a']), - throwsA(isA()), - ); - }); + test('close tears the worker down and later calls fail fast', () async { + final worker = await EmbeddingWorker.start(entryPoint: echoEntryPoint); + await worker.close(); + await expectLater( + worker.embedBatch(['a']), + throwsA(isA()), + ); + }); - test('a request fails (does not hang) if the worker dies mid-request', - () async { - final worker = await EmbeddingWorker.start(entryPoint: dyingEntryPoint); - await expectLater( - worker.embedBatch(['a']), - throwsA(isA()), - ); - await worker.close(); - }); + test('close is idempotent', () async { + final worker = await EmbeddingWorker.start(entryPoint: echoEntryPoint); + await worker.close(); + await worker.close(); + }); + + test('a request fails (does not hang) if the worker dies mid-request', + () async { + final worker = await EmbeddingWorker.start(entryPoint: dyingEntryPoint); + addTearDown(worker.close); + await expectLater( + worker.embedBatch(['a']), + throwsA(isA()), + ); + }, timeout: const Timeout(Duration(seconds: 20))); - test('start fails if the worker never completes the handshake', () async { - await expectLater( - EmbeddingWorker.start(entryPoint: noHandshakeEntryPoint), - throwsA(isA()), - ); + test('start fails if the worker never completes the handshake', () async { + await expectLater( + EmbeddingWorker.start(entryPoint: noHandshakeEntryPoint), + throwsA(isA()), + ); + }, timeout: const Timeout(Duration(seconds: 20))); }); } diff --git a/test/error_path_test.dart b/test/exception_test.dart similarity index 64% rename from test/error_path_test.dart rename to test/exception_test.dart index 3967eb1..b7f4bb7 100644 --- a/test/error_path_test.dart +++ b/test/exception_test.dart @@ -1,5 +1,3 @@ -import 'dart:typed_data'; - import 'package:model2vec/model2vec.dart'; import 'package:test/test.dart'; @@ -40,26 +38,4 @@ void main() { expect(e.toString(), contains('boom')); }); }); - - group('native error path', () { - test('loadModelFromBytes with garbage throws a typed exception', () { - final garbage = Uint8List.fromList([0, 1, 2, 3]); - expect( - () => Model2Vec.loadModelFromBytes( - tokenizerBytes: garbage, - modelBytes: garbage, - configBytes: garbage, - ), - throwsA( - isA() - .having( - (e) => e.kind, - 'kind', - Model2VecErrorKind.initFromBytesFailed, - ) - .having((e) => e.message, 'message', isNotEmpty), - ), - ); - }); - }); } diff --git a/test/init_async_test.dart b/test/init_async_test.dart deleted file mode 100644 index 12376f1..0000000 --- a/test/init_async_test.dart +++ /dev/null @@ -1,46 +0,0 @@ -import 'dart:io'; - -import 'package:model2vec/model2vec.dart'; -import 'package:test/test.dart'; - -void main() { - group('async model loading', () { - test( - 'loadModelAsync loads the process-global model, visible on the ' - 'calling isolate', - () async { - await Model2Vec.loadModelAsync('minishlab/potion-base-2M'); - - // The load ran on a background isolate; the native model is a - // process-global, so a synchronous call on this isolate sees it. - expect(Model2Vec.isInitialized, isTrue); - expect(Model2Vec.embeddingDimension, equals(64)); - expect( - Model2Vec.generateEmbedding('async load').length, - equals(64), - ); - }, - ); - - test('loadModelAdvancedAsync honors a custom cache directory', () async { - final tempDir = Directory.systemTemp.createTempSync('m2v_async_cache_'); - try { - await Model2Vec.loadModelAdvancedAsync( - modelPath: 'minishlab/potion-base-2M', - cacheDirectory: tempDir.path, - ); - expect(Model2Vec.isInitialized, isTrue); - expect(Model2Vec.embeddingDimension, equals(64)); - } finally { - tempDir.deleteSync(recursive: true); - } - }); - - test('a failed async load surfaces a typed Model2VecException', () { - expect( - Model2Vec.loadModelAsync('definitely/not-a-real-model-xyz'), - throwsA(isA()), - ); - }); - }); -} diff --git a/test/lifecycle_test.dart b/test/lifecycle_test.dart deleted file mode 100644 index 0e10e70..0000000 --- a/test/lifecycle_test.dart +++ /dev/null @@ -1,37 +0,0 @@ -import 'package:model2vec/model2vec.dart'; -import 'package:test/test.dart'; - -void main() { - group('model lifecycle', () { - test('isInitialized / modelInfo / unloadModel round-trip', () { - Model2Vec.loadModel('minishlab/potion-base-2M'); - expect(Model2Vec.isInitialized, isTrue); - - final info = Model2Vec.modelInfo; - expect(info.dimension, equals(64)); - expect(info.vocabularySize, greaterThan(20000)); - expect(info.isNormalized, isTrue); - expect(info.medianTokenLength, isPositive); - - Model2Vec.unloadModel(); - expect(Model2Vec.isInitialized, isFalse); - expect( - () => Model2Vec.embeddingDimension, - throwsA( - isA().having( - (e) => e.kind, - 'kind', - Model2VecErrorKind.notInitialized, - ), - ), - ); - - // Idempotent: unloading again is a no-op. - expect(Model2Vec.unloadModel, returnsNormally); - - // Re-initialize so later suites in this process find a model. - Model2Vec.loadModel('minishlab/potion-base-2M'); - expect(Model2Vec.isInitialized, isTrue); - }); - }); -} diff --git a/test/load_progress_test.dart b/test/load_progress_test.dart deleted file mode 100644 index 8eab08a..0000000 --- a/test/load_progress_test.dart +++ /dev/null @@ -1,79 +0,0 @@ -import 'dart:io'; - -import 'package:model2vec/model2vec.dart'; -import 'package:test/test.dart'; - -void main() { - group('loadModelWithProgress', () { - test('reports download progress and reaches done on a fresh cache', - () async { - // A private temp cache guarantees a cache miss, so the weights actually - // download and the downloading phase is exercised. - final tmp = Directory.systemTemp.createTempSync('m2v_progress_'); - addTearDown(() => tmp.deleteSync(recursive: true)); - - final events = []; - await for (final p in Model2Vec.loadModelWithProgress( - 'minishlab/potion-base-2M', - cacheDirectory: tmp.path, - pollInterval: const Duration(milliseconds: 1), - )) { - events.add(p); - } - - // The load succeeded and the stream ended on `done`. - expect(Model2Vec.isInitialized, isTrue); - expect(events, isNotEmpty); - expect(events.last.phase, LoadPhase.done); - - // `done` is only ever the terminal event. - final nonTerminal = events.sublist(0, events.length - 1); - expect( - nonTerminal, - everyElement( - isNot( - predicate((e) => e.phase == LoadPhase.done), - ), - ), - ); - - // A fresh cache must download the weights: we should have caught at least - // one downloading snapshot with a known total and a real fraction. - final downloading = - events.where((e) => e.phase == LoadPhase.downloading).toList(); - expect( - downloading, - isNotEmpty, - reason: 'a fresh cache should download the weights', - ); - final last = downloading.last; - expect(last.totalBytes, greaterThan(0)); - expect(last.bytesDownloaded, greaterThan(0)); - expect(last.fraction, isNotNull); - expect(last.fraction, inInclusiveRange(0.0, 1.0)); - - // Print the shape of the run for eyeballing. - final byPhase = {}; - for (final e in events) { - byPhase[e.phase] = (byPhase[e.phase] ?? 0) + 1; - } - // Diagnostic output for eyeballing a real run. - // ignore: avoid_print - print('phase counts: $byPhase; last downloading: $last'); - }); - - test('a cached model streams straight to done', () async { - // The prior test cached the model in a temp dir, but the default cache - // may already hold it too; either way this must terminate on `done`. - final events = []; - await for (final p in Model2Vec.loadModelWithProgress( - 'minishlab/potion-base-2M', - )) { - events.add(p); - } - - expect(Model2Vec.isInitialized, isTrue); - expect(events.last.phase, LoadPhase.done); - }); - }); -} diff --git a/test/model2vec_test.dart b/test/model2vec_test.dart deleted file mode 100644 index 5c93aa3..0000000 --- a/test/model2vec_test.dart +++ /dev/null @@ -1,204 +0,0 @@ -import 'dart:io'; - -import 'package:model2vec/model2vec.dart'; -import 'package:test/test.dart'; - -void main() { - group('Model2Vec Production Tests', () { - test('Successful initialization (online)', () { - // Must be the first init call so later tests have a loaded model. - expect( - () => Model2Vec.loadModel('minishlab/potion-base-2M'), - returnsNormally, - ); - }); - - group('Model Metadata', () { - test('returns valid dimension and metadata', () { - expect(Model2Vec.embeddingDimension, equals(64)); - expect(Model2Vec.vocabularySize, greaterThan(20000)); - expect(Model2Vec.isNormalized, isTrue); - expect(Model2Vec.medianTokenLength, isPositive); - }); - }); - - group('Tokenization', () { - test('breaks text into tokens correctly', () { - final tokens = Model2Vec.tokenize('Dart FFI is powerful'); - expect(tokens, isNotEmpty); - expect(tokens, contains('dart')); - }); - - test('handles empty string tokenization', () { - final tokens = Model2Vec.tokenize(''); - expect(tokens, isEmpty); - }); - }); - - group('Embeddings', () { - test( - 'generates vector of correct length and exact values (snapshot)', - () { - final vector = Model2Vec.generateEmbedding('Hello world'); - expect(vector.length, equals(Model2Vec.embeddingDimension)); - - // Exact values for Potion Base 2M on "Hello world" - // Using closeTo to handle minor floating point precision differences - expect(vector[0], closeTo(-0.06458017975091934, 1e-6)); - expect(vector[1], closeTo(-0.15089768171310425, 1e-6)); - expect(vector[2], closeTo(-0.27086204290390015, 1e-6)); - }, - ); - - test('batch embedding is consistent with single embedding', () { - const text = 'Consistency check'; - final single = Model2Vec.generateEmbedding(text); - final batch = Model2Vec.generateBatchEmbeddings([text]); - - expect(batch.length, 1); - expect(batch[0], orderedEquals(single)); - }); - - test('handles batch with multiple diverse strings', () { - final texts = [ - 'Short', - 'A much longer sentence to test truncation logic in the underlying ', - '12345 !? @#', - 'Теж має працювати', - ]; - final results = Model2Vec.generateBatchEmbeddings(texts); - expect(results.length, equals(texts.length)); - for (final v in results) { - expect(v.length, equals(Model2Vec.embeddingDimension)); - } - }); - - test('batch with empty list returns empty list', () { - expect(Model2Vec.generateBatchEmbeddings([]), isEmpty); - }); - }); - - group('Asynchronous API', () { - test('generates embedding asynchronously in isolate', () async { - final vector = await Model2Vec.generateEmbeddingAsync('Async test'); - expect(vector.length, equals(Model2Vec.embeddingDimension)); - }); - - test('generates batch embeddings asynchronously in isolate', () async { - final results = await Model2Vec.generateBatchEmbeddingsAsync([ - 'Async 1', - 'Async 2', - ]); - expect(results.length, 2); - expect(results[0].length, Model2Vec.embeddingDimension); - }); - }); - - group('Streaming API', () { - test('processes a stream of texts and yields embeddings', () async { - // Create a synthetic stream of 2500 strings - const totalItems = 2500; - final stream = Stream.fromIterable( - Iterable.generate(totalItems, (i) => 'This is test string number $i'), - ); - - final resultStream = Model2Vec.generateEmbeddingStream( - stream, - batchSize: 1000, - ); - - final results = await resultStream.toList(); - - expect(results.length, equals(totalItems)); - expect(results.first.length, equals(Model2Vec.embeddingDimension)); - }); - - test('works correctly with useIsolate: false', () async { - final stream = Stream.fromIterable(['Text 1', 'Text 2', 'Text 3']); - final results = await Model2Vec.generateEmbeddingStream( - stream, - batchSize: 2, - useIsolate: false, - ).toList(); - - expect(results.length, equals(3)); - }); - - test('canceling stream early works cleanly', () async { - final stream = Stream.fromIterable([ - 'Text 1', - 'Text 2', - 'Text 3', - 'Text 4', - ]); - final resultStream = Model2Vec.generateEmbeddingStream( - stream, - batchSize: 2, - ); - - // Take only 2 elements, which will cancel the subscription early - final firstTwo = await resultStream.take(2).toList(); - expect(firstTwo.length, 2); - }); - }); - - group('Model Switching', () { - test('can switch between different models successfully', () { - // Switch to 8M model (dimension 256) - Model2Vec.loadModel('minishlab/potion-base-8M'); - expect(Model2Vec.embeddingDimension, equals(256)); - - // Switch back to 2M model (dimension 64) - Model2Vec.loadModel('minishlab/potion-base-2M'); - expect(Model2Vec.embeddingDimension, equals(64)); - }); - }); - - group('Advanced Features', () { - test('maxLength truncation applies to single and batch embeddings', () { - const text = - 'A very long sentence to test truncation with maxLength parameter'; - // With very short maxLength, the embedding should be different from - // default - final defaultEmbedding = Model2Vec.generateEmbedding(text); - final shortEmbedding = Model2Vec.generateEmbedding(text, maxLength: 2); - - expect(defaultEmbedding, isNot(orderedEquals(shortEmbedding))); - - final batchEmbedding = Model2Vec.generateBatchEmbeddings( - [text, text], - maxLength: 2, - ); - expect(batchEmbedding.length, 2); - expect(batchEmbedding[0], orderedEquals(shortEmbedding)); - expect(batchEmbedding[1], orderedEquals(shortEmbedding)); - }); - - test('supports custom cache directory', () { - final tempDir = Directory.systemTemp.createTempSync('m2v_cache_'); - try { - // Now that switching is enabled, this should return normally - expect( - () => Model2Vec.loadModelAdvanced( - modelPath: 'minishlab/potion-base-2M', - cacheDirectory: tempDir.path, - ), - returnsNormally, - ); - } finally { - tempDir.deleteSync(recursive: true); - } - }); - }); - - group('Recommended Models', () { - test('exposes a non-empty typed catalog', () { - expect(Model2Vec.recommendedModels, isNotEmpty); - expect( - Model2Vec.recommendedModels.first.id, - startsWith('minishlab/'), - ); - }); - }); - }); -} diff --git a/test/model_loading_test.dart b/test/model_loading_test.dart new file mode 100644 index 0000000..ffa23c1 --- /dev/null +++ b/test/model_loading_test.dart @@ -0,0 +1,265 @@ +import 'dart:io'; +import 'dart:typed_data'; + +import 'package:model2vec/model2vec.dart'; +import 'package:test/test.dart'; + +import 'support.dart'; + +void main() { + // Load the default fixture once before the suite, and restore it after every + // test. The native model is a single process-global, so tests that unload or + // switch it must not leak that state to their neighbours; reloading the + // (cached) model after each test keeps every test starting from a clean 2M. + setUpAll(() => Model2Vec.loadModel(testModelId)); + tearDown(() => Model2Vec.loadModel(testModelId)); + + group('loadModel', () { + test('loads a model by repo id and reports it initialized', () { + Model2Vec.loadModel(testModelId); + expect(Model2Vec.isInitialized, isTrue); + }); + }); + + group('loadModelAdvanced', () { + test('populates a fresh custom cache directory with the model files', () { + final cacheDir = Directory.systemTemp.createTempSync('m2v_cache_'); + addTearDown(() => cacheDir.deleteSync(recursive: true)); + + Model2Vec.loadModelAdvanced( + modelPath: testModelId, + cacheDirectory: cacheDir.path, + ); + + expect(Model2Vec.isInitialized, isTrue); + // The download actually landed in the directory we handed it. + expect(cacheDir.listSync(), isNotEmpty); + }); + + test('with normalize:false disables output normalization', () { + Model2Vec.loadModelAdvanced(modelPath: testModelId, normalize: false); + expect(Model2Vec.isInitialized, isTrue); + expect(Model2Vec.isNormalized, isFalse); + // tearDown restores the normal (normalized) model for later tests. + }); + }); + + group('loadModelFromBytes', () { + test('loads a model from the cached tokenizer/model/config bytes', () { + final bytes = _cachedModelBytes(testModelId); + if (bytes == null) { + markTestSkipped('cached files for $testModelId not found in HF cache'); + return; + } + + Model2Vec.loadModelFromBytes( + tokenizerBytes: bytes.tokenizer, + modelBytes: bytes.model, + configBytes: bytes.config, + ); + + expect(Model2Vec.isInitialized, isTrue); + expect(Model2Vec.embeddingDimension, testModelDim); + }); + + test('with garbage bytes throws a typed initFromBytesFailed', () { + final garbage = Uint8List.fromList([0, 1, 2, 3]); + expect( + () => Model2Vec.loadModelFromBytes( + tokenizerBytes: garbage, + modelBytes: garbage, + configBytes: garbage, + ), + throwsA( + isA() + .having( + (e) => e.kind, + 'kind', + Model2VecErrorKind.initFromBytesFailed, + ) + .having((e) => e.message, 'message', isNotEmpty), + ), + ); + }); + }); + + group('async loading', () { + test('loadModelAsync loads the process-global model', () async { + await Model2Vec.loadModelAsync(testModelId); + + // The load ran on a background isolate; the native model is a + // process-global, so a synchronous call on this isolate sees it. + expect(Model2Vec.isInitialized, isTrue); + expect(Model2Vec.embeddingDimension, testModelDim); + expect(Model2Vec.generateEmbedding('async load').length, testModelDim); + }); + + test('loadModelAdvancedAsync honors a custom cache directory', () async { + final cacheDir = Directory.systemTemp.createTempSync('m2v_async_cache_'); + addTearDown(() => cacheDir.deleteSync(recursive: true)); + + await Model2Vec.loadModelAdvancedAsync( + modelPath: testModelId, + cacheDirectory: cacheDir.path, + ); + + expect(Model2Vec.isInitialized, isTrue); + expect(Model2Vec.embeddingDimension, testModelDim); + }); + + test('a failed async load surfaces a typed Model2VecException', () async { + // Awaited so the rejected future is asserted here, not left to surface + // later as an unhandled async error. + await expectLater( + Model2Vec.loadModelAsync('bad/nope'), + throwsA(isA()), + ); + }); + }); + + group('loadModelWithProgress', () { + test('downloads and reaches done on a fresh cache', () async { + // A private temp cache guarantees a cache miss, so the weights actually + // download and the downloading phase is exercised. + final tmp = Directory.systemTemp.createTempSync('m2v_progress_'); + addTearDown(() => tmp.deleteSync(recursive: true)); + + final events = []; + await for (final p in Model2Vec.loadModelWithProgress( + testModelId, + cacheDirectory: tmp.path, + pollInterval: const Duration(milliseconds: 1), + )) { + events.add(p); + } + + // The load succeeded and the stream ended on `done`. + expect(Model2Vec.isInitialized, isTrue); + expect(events, isNotEmpty); + expect(events.last.phase, LoadPhase.done); + + // `done` is only ever the terminal event. + final nonTerminal = events.sublist(0, events.length - 1); + expect( + nonTerminal, + everyElement( + isNot(predicate((e) => e.phase == LoadPhase.done)), + ), + ); + + // A fresh cache must download the weights: we should have caught at least + // one downloading snapshot with a known total and a real fraction. + final downloading = events + .where((e) => e.phase == LoadPhase.downloading) + .toList(); + expect( + downloading, + isNotEmpty, + reason: 'a fresh cache should download the weights', + ); + final last = downloading.last; + expect(last.totalBytes, greaterThan(0)); + expect(last.bytesDownloaded, greaterThan(0)); + expect(last.fraction, isNotNull); + expect(last.fraction, inInclusiveRange(0.0, 1.0)); + }); + + test('a cached model streams straight to done', () async { + // The prior test cached the model in a temp dir, but the default cache + // may already hold it too; either way this must terminate on `done`. + final events = []; + await for (final p in Model2Vec.loadModelWithProgress(testModelId)) { + events.add(p); + } + + expect(Model2Vec.isInitialized, isTrue); + expect(events.last.phase, LoadPhase.done); + }); + + test('a failed load surfaces a typed exception on the stream', () async { + await expectLater( + Model2Vec.loadModelWithProgress('definitely/not-a-real-model').toList(), + throwsA(isA()), + ); + }); + }); + + group('isInitialized and unloadModel', () { + test('unloadModel clears the model; metadata getters then throw', () { + expect(Model2Vec.isInitialized, isTrue); + + Model2Vec.unloadModel(); + expect(Model2Vec.isInitialized, isFalse); + expect( + () => Model2Vec.embeddingDimension, + throwsA( + isA().having( + (e) => e.kind, + 'kind', + Model2VecErrorKind.notInitialized, + ), + ), + ); + + // Idempotent: unloading again is a no-op. + expect(Model2Vec.unloadModel, returnsNormally); + // tearDown reloads the model for later tests in this process. + }); + }); + + group('modelInfo', () { + test('reports dimension, vocab, normalization and median token length', () { + final info = Model2Vec.modelInfo; + expect(info.dimension, testModelDim); + expect(info.vocabularySize, greaterThan(testModelVocabMin)); + expect(info.isNormalized, isTrue); + expect(info.medianTokenLength, isPositive); + }); + + test('the four metadata getters agree with modelInfo', () { + expect(Model2Vec.embeddingDimension, testModelDim); + expect(Model2Vec.vocabularySize, greaterThan(testModelVocabMin)); + expect(Model2Vec.isNormalized, isTrue); + expect(Model2Vec.medianTokenLength, isPositive); + }); + }); + + group('model switching', () { + test('switches between the 2M and 8M models', () { + Model2Vec.loadModel(largeModelId); + expect(Model2Vec.embeddingDimension, largeModelDim); + + Model2Vec.loadModel(testModelId); + expect(Model2Vec.embeddingDimension, testModelDim); + }); + }); +} + +/// Reads the cached `tokenizer.json`, `model.safetensors` and `config.json` +/// bytes for [modelId] from the local Hugging Face cache, or `null` when the +/// model is not cached. Used by the `loadModelFromBytes` happy-path test. +({Uint8List tokenizer, Uint8List model, Uint8List config})? _cachedModelBytes( + String modelId, +) { + final env = Platform.environment; + final hfHome = env['HF_HOME'] ?? '${env['HOME']}/.cache/huggingface'; + final repoDir = 'models--${modelId.replaceAll('/', '--')}'; + final snapshots = Directory('$hfHome/hub/$repoDir/snapshots'); + if (!snapshots.existsSync()) { + return null; + } + + for (final entry in snapshots.listSync().whereType()) { + final tokenizer = File('${entry.path}/tokenizer.json'); + final model = File('${entry.path}/model.safetensors'); + final config = File('${entry.path}/config.json'); + if (tokenizer.existsSync() && model.existsSync() && config.existsSync()) { + return ( + tokenizer: tokenizer.readAsBytesSync(), + model: model.readAsBytesSync(), + config: config.readAsBytesSync(), + ); + } + } + return null; +} diff --git a/test/recommended_models_test.dart b/test/recommended_models_test.dart new file mode 100644 index 0000000..96ded91 --- /dev/null +++ b/test/recommended_models_test.dart @@ -0,0 +1,78 @@ +import 'package:model2vec/model2vec.dart'; +import 'package:test/test.dart'; + +/// A well-formed catalog entry to reuse across the value-type tests. +const _entry = RecommendedModel( + id: 'minishlab/potion-base-8M', + name: 'Potion Base 8M', + lang: 'English', + params: '7.5M', + description: 'Balanced English model.', +); + +void main() { + group('RecommendedModel value type', () { + test('two instances with the same fields are equal', () { + const other = RecommendedModel( + id: 'minishlab/potion-base-8M', + name: 'Potion Base 8M', + lang: 'English', + params: '7.5M', + description: 'Balanced English model.', + ); + + expect(_entry, equals(other)); + expect(_entry.hashCode, equals(other.hashCode)); + }); + + test('a differing field breaks equality', () { + const differ = RecommendedModel( + id: 'minishlab/potion-base-8M', + name: 'Potion Base 8M', + lang: 'English', + params: '7.5M', + description: 'A different description.', + ); + + expect(_entry, isNot(equals(differ))); + }); + + test('toString contains the id', () { + expect(_entry.toString(), contains(_entry.id)); + }); + }); + + group('recommendedModels catalog', () { + const catalog = Model2Vec.recommendedModels; + + test('is non-empty', () { + expect(catalog, isNotEmpty); + }); + + test('every id is unique', () { + final ids = catalog.map((m) => m.id).toList(); + expect(ids.toSet(), hasLength(ids.length)); + }); + + test('every field is non-empty', () { + for (final m in catalog) { + expect(m.id, isNotEmpty, reason: 'id of $m'); + expect(m.name, isNotEmpty, reason: 'name of $m'); + expect(m.lang, isNotEmpty, reason: 'lang of $m'); + expect(m.params, isNotEmpty, reason: 'params of $m'); + expect(m.description, isNotEmpty, reason: 'description of $m'); + } + }); + + test('every id is a well-formed "owner/model" repo id', () { + final repoId = RegExp(r'^[\w.-]+/[\w.-]+$'); + for (final m in catalog) { + expect( + m.id, + matches(repoId), + reason: '${m.id} should look like a Hugging Face repo id', + ); + } + }); + }); +} diff --git a/test/support.dart b/test/support.dart new file mode 100644 index 0000000..2a68844 --- /dev/null +++ b/test/support.dart @@ -0,0 +1,20 @@ +// Shared fixtures for the core Model2Vec integration tests. +// +// The native model is a single process-global, so every core test file loads +// one of these known models. Centralizing the ids and their dimensions here +// removes the magic numbers that used to be duplicated across the suite. + +/// Small, fast English model used as the default fixture across the suite. +const testModelId = 'minishlab/potion-base-2M'; + +/// Embedding dimension of [testModelId]. +const testModelDim = 64; + +/// Conservative lower bound on the vocabulary size of [testModelId]. +const testModelVocabMin = 20000; + +/// A second, larger model used to exercise model switching. +const largeModelId = 'minishlab/potion-base-8M'; + +/// Embedding dimension of [largeModelId]. +const largeModelDim = 256; diff --git a/test/utils_test.dart b/test/utils_test.dart index 0adc7ff..d6bf96f 100644 --- a/test/utils_test.dart +++ b/test/utils_test.dart @@ -1,316 +1,384 @@ -import 'dart:io'; import 'dart:typed_data'; import 'package:model2vec/model2vec.dart'; import 'package:test/test.dart'; void main() { - group('Model2VecUtils - Math', () { - test('cosineSimilarity calculates correctly', () { - final a = Float32List.fromList([1.0, 0.0, 0.0]); - final b = Float32List.fromList([0.0, 1.0, 0.0]); - final c = Float32List.fromList([1.0, 1.0, 0.0]); - - // Orthogonal vectors - expect(Model2VecUtils.cosineSimilarity(a, b), closeTo(0.0, 1e-6)); - - // Identical vectors - expect(Model2VecUtils.cosineSimilarity(a, a), closeTo(1.0, 1e-6)); - - // 45 degree angle - expect(Model2VecUtils.cosineSimilarity(a, c), closeTo(0.707106, 1e-4)); + group('Model2VecUtils', () { + group('cosineSimilarity', () { + test('scores orthogonal, identical, and 45-degree vectors', () { + final a = Float32List.fromList([1.0, 0.0, 0.0]); + final b = Float32List.fromList([0.0, 1.0, 0.0]); + final c = Float32List.fromList([1.0, 1.0, 0.0]); + + expect(Model2VecUtils.cosineSimilarity(a, b), closeTo(0.0, 1e-6)); + expect(Model2VecUtils.cosineSimilarity(a, a), closeTo(1.0, 1e-6)); + expect(Model2VecUtils.cosineSimilarity(a, c), closeTo(0.707106, 1e-4)); + }); + + test('returns 0.0 for a zero vector instead of NaN', () { + final zero = Float32List.fromList([0.0, 0.0, 0.0]); + final a = Float32List.fromList([1.0, 0.0, 0.0]); + + final score = Model2VecUtils.cosineSimilarity(zero, a); + expect(score, equals(0.0)); + expect(score.isNaN, isFalse); + }); + + test('throws on mismatched lengths', () { + final a = Float32List.fromList([1.0, 2.0]); + final b = Float32List.fromList([1.0, 2.0, 3.0]); + + expect( + () => Model2VecUtils.cosineSimilarity(a, b), + throwsArgumentError, + ); + }); }); - test('dotProduct calculates correctly', () { - final a = Float32List.fromList([1.0, 2.0, 3.0]); - final b = Float32List.fromList([4.0, 5.0, 6.0]); + group('dotProduct', () { + test('multiplies and sums element-wise', () { + final a = Float32List.fromList([1.0, 2.0, 3.0]); + final b = Float32List.fromList([4.0, 5.0, 6.0]); - // 1*4 + 2*5 + 3*6 = 4 + 10 + 18 = 32 - expect(Model2VecUtils.dotProduct(a, b), equals(32.0)); - }); + // 1*4 + 2*5 + 3*6 = 32 + expect(Model2VecUtils.dotProduct(a, b), equals(32.0)); + }); - test('euclideanDistance calculates correctly', () { - final a = Float32List.fromList([0.0, 0.0]); - final b = Float32List.fromList([3.0, 4.0]); + test('throws on mismatched lengths', () { + final a = Float32List.fromList([1.0, 2.0]); + final b = Float32List.fromList([1.0, 2.0, 3.0]); - // sqrt(3^2 + 4^2) = 5 - expect(Model2VecUtils.euclideanDistance(a, b), equals(5.0)); + expect(() => Model2VecUtils.dotProduct(a, b), throwsArgumentError); + }); }); - test('throws ArgumentError on mismatched lengths', () { - final a = Float32List.fromList([1.0, 2.0]); - final b = Float32List.fromList([1.0, 2.0, 3.0]); + group('euclideanDistance', () { + test('computes the straight-line distance', () { + final a = Float32List.fromList([0.0, 0.0]); + final b = Float32List.fromList([3.0, 4.0]); - expect(() => Model2VecUtils.cosineSimilarity(a, b), throwsArgumentError); - }); + // sqrt(3^2 + 4^2) = 5 + expect(Model2VecUtils.euclideanDistance(a, b), equals(5.0)); + }); - test('cosineDistance calculates correctly', () { - final a = Float32List.fromList([1.0, 0.0]); - final b = Float32List.fromList([0.0, 1.0]); - expect(Model2VecUtils.cosineDistance(a, a), closeTo(0.0, 1e-6)); - expect(Model2VecUtils.cosineDistance(a, b), closeTo(1.0, 1e-6)); - }); + test('throws on mismatched lengths', () { + final a = Float32List.fromList([1.0, 2.0]); + final b = Float32List.fromList([1.0, 2.0, 3.0]); - test('normalize works correctly', () { - final a = Float32List.fromList([3.0, 4.0]); - final normalized = Model2VecUtils.normalize(a); - - expect(normalized[0], closeTo(0.6, 1e-6)); - expect(normalized[1], closeTo(0.8, 1e-6)); - expect( - Model2VecUtils.euclideanDistance( - normalized, - Float32List.fromList([0, 0]), - ), - closeTo(1.0, 1e-6), - ); + expect( + () => Model2VecUtils.euclideanDistance(a, b), + throwsArgumentError, + ); + }); }); - test('meanPooling averages correctly', () { - final a = Float32List.fromList([1.0, 2.0]); - final b = Float32List.fromList([3.0, 4.0]); - final c = Float32List.fromList([5.0, 6.0]); - - final mean = Model2VecUtils.meanPooling([a, b, c]); - expect(mean[0], closeTo(3.0, 1e-6)); - expect(mean[1], closeTo(4.0, 1e-6)); + group('cosineDistance', () { + test('is 1 minus the cosine similarity', () { + final a = Float32List.fromList([1.0, 0.0]); + final b = Float32List.fromList([0.0, 1.0]); + expect(Model2VecUtils.cosineDistance(a, a), closeTo(0.0, 1e-6)); + expect(Model2VecUtils.cosineDistance(a, b), closeTo(1.0, 1e-6)); + }); }); - test('quantizeToInt8 scales correctly', () { - final vector = Float32List.fromList([1.0, 0.5, -0.5, -1.0]); - final quantized = Model2VecUtils.quantizeToInt8(vector); - - expect(quantized[0], equals(127)); - expect(quantized[1], equals(64)); - expect(quantized[2], equals(-64)); // -0.5 * 127 = -63.5 -> -64 - expect(quantized[3], equals(-127)); - }); + group('normalize', () { + test('scales a vector to unit length', () { + final a = Float32List.fromList([3.0, 4.0]); + final normalized = Model2VecUtils.normalize(a); + + expect(normalized[0], closeTo(0.6, 1e-6)); + expect(normalized[1], closeTo(0.8, 1e-6)); + expect( + Model2VecUtils.euclideanDistance( + normalized, + Float32List.fromList([0, 0]), + ), + closeTo(1.0, 1e-6), + ); + }); - test('dequantizeInt8 approximately inverts quantizeToInt8', () { - final vector = Float32List.fromList([0.5, -0.25, 1.0, -1.0, 0.0]); - final restored = Model2VecUtils.dequantizeInt8( - Model2VecUtils.quantizeToInt8(vector), - ); + test('returns all-zeros for a zero vector with no NaN', () { + final zero = Float32List.fromList([0.0, 0.0, 0.0]); + final normalized = Model2VecUtils.normalize(zero); - expect(restored.length, equals(vector.length)); - for (var i = 0; i < vector.length; i++) { - expect(restored[i], closeTo(vector[i], 0.01)); - } + expect(normalized, [0.0, 0.0, 0.0]); + expect(normalized.any((e) => e.isNaN), isFalse); + }); }); - test('similaritySearchWithScores returns sorted index+score pairs', () { - final query = Float32List.fromList([1, 0, 0]); - final candidates = [ - Float32List.fromList([1, 0.2, 0]), - Float32List.fromList([1, 0.25, 0]), - Float32List.fromList([0.6, 0, 0.8]), - ]; - - final results = Model2VecUtils.similaritySearchWithScores( - query, - candidates, - topK: 2, - ); - expect(results.map((r) => r.index).toList(), [0, 1]); - expect(results.first.score, greaterThan(results[1].score)); - expect( - Model2VecUtils.similaritySearchWithScores(query, const []), - isEmpty, - ); + group('meanPooling', () { + test('averages vectors element-wise', () { + final a = Float32List.fromList([1.0, 2.0]); + final b = Float32List.fromList([3.0, 4.0]); + final c = Float32List.fromList([5.0, 6.0]); + + final mean = Model2VecUtils.meanPooling([a, b, c]); + expect(mean[0], closeTo(3.0, 1e-6)); + expect(mean[1], closeTo(4.0, 1e-6)); + }); + + test('throws on an empty list', () { + expect( + () => Model2VecUtils.meanPooling(const []), + throwsArgumentError, + ); + }); + + test('throws on mismatched dimensions', () { + final a = Float32List.fromList([1.0, 2.0]); + final b = Float32List.fromList([3.0, 4.0, 5.0]); + expect( + () => Model2VecUtils.meanPooling([a, b]), + throwsArgumentError, + ); + }); }); - test('similaritySearchWithScores drops results below the threshold', () { - final query = Float32List.fromList([1.0, 0.0]); - final candidates = [ - Float32List.fromList([0.0, 1.0]), // sim: 0.0 - Float32List.fromList([0.9, 0.1]), // sim: ~0.994 - Float32List.fromList([-1.0, 0.0]), // sim: -1.0 - Float32List.fromList([0.5, 0.5]), // sim: ~0.707 - ]; - - final results = Model2VecUtils.similaritySearchWithScores( - query, - candidates, - topK: 10, - threshold: 0.8, - ); - - expect(results.map((r) => r.index).toList(), [1]); - expect(results.single.score, greaterThan(0.8)); + group('quantizeToInt8', () { + test('scales values into the int8 range', () { + final vector = Float32List.fromList([1.0, 0.5, -0.5, -1.0]); + final quantized = Model2VecUtils.quantizeToInt8(vector); + + expect(quantized[0], equals(127)); + expect(quantized[1], equals(64)); + expect(quantized[2], equals(-64)); // -0.5 * 127 = -63.5 -> -64 + expect(quantized[3], equals(-127)); + }); + + test('clamps out-of-range inputs to [-128, 127]', () { + final vector = Float32List.fromList([1.5, -1.5, 2.0]); + final quantized = Model2VecUtils.quantizeToInt8(vector); + + // 1.5*127=190.5 -> 191 -> clamp 127; -1.5*127=-190.5 -> -191 -> + // clamp -128; 2.0*127=254 -> clamp 127. + expect(quantized, [127, -128, 127]); + }); }); - test('similaritySearchWithScores caps threshold matches at topK', () { - final query = Float32List.fromList([1.0, 0.0]); - final candidates = [ - Float32List.fromList([1.0, 0.0]), // sim 1.0 - Float32List.fromList([0.99, 0.01]), // sim ~0.9999 - Float32List.fromList([0.98, 0.02]), // sim ~0.9998 - ]; - - final results = Model2VecUtils.similaritySearchWithScores( - query, - candidates, - topK: 2, - threshold: 0.5, // all three clear the floor, but topK caps to 2 - ); - - expect(results, hasLength(2)); - expect(results.map((r) => r.index).toList(), [0, 1]); - }); + group('dequantizeInt8', () { + test('approximately inverts quantizeToInt8', () { + final vector = Float32List.fromList([0.5, -0.25, 1.0, -1.0, 0.0]); + final restored = Model2VecUtils.dequantizeInt8( + Model2VecUtils.quantizeToInt8(vector), + ); - test('MMR prefers a diverse result over a near-duplicate', () { - final query = Float32List.fromList([1, 0, 0]); - final candidates = [ - Float32List.fromList([1, 0.2, 0]), // 0: most relevant - Float32List.fromList([1, 0.25, 0]), // 1: near-duplicate of 0 - Float32List.fromList([0.6, 0, 0.8]), // 2: relevant-ish but diverse - ]; - - final selected = Model2VecUtils.maximalMarginalRelevance( - query, - candidates, - topK: 2, - ); - expect(selected.first, 0); // most relevant first - expect(selected[1], 2); // diversity beats the near-duplicate + expect(restored.length, equals(vector.length)); + for (var i = 0; i < vector.length; i++) { + expect(restored[i], closeTo(vector[i], 0.01)); + } + }); }); - test('MMR with lambda 1.0 is pure relevance order', () { - final query = Float32List.fromList([1, 0, 0]); - final candidates = [ - Float32List.fromList([0.6, 0, 0.8]), - Float32List.fromList([1, 0.2, 0]), - Float32List.fromList([1, 0.25, 0]), - ]; - - final selected = Model2VecUtils.maximalMarginalRelevance( - query, - candidates, - topK: 3, - lambda: 1, - ); - expect(selected, [1, 2, 0]); - expect( - Model2VecUtils.maximalMarginalRelevance(query, const []), - isEmpty, - ); - }); + group('similaritySearchWithScores', () { + test('returns sorted index+score pairs', () { + final query = Float32List.fromList([1, 0, 0]); + final candidates = [ + Float32List.fromList([1, 0.2, 0]), + Float32List.fromList([1, 0.25, 0]), + Float32List.fromList([0.6, 0, 0.8]), + ]; - test('MMR rejects lambda outside [0, 1]', () { - final query = Float32List.fromList([1, 0]); - final candidates = [ - Float32List.fromList([1, 0]), - ]; - expect( - () => Model2VecUtils.maximalMarginalRelevance( + final results = Model2VecUtils.similaritySearchWithScores( query, candidates, - lambda: 1.5, - ), - throwsArgumentError, - ); - expect( - () => Model2VecUtils.maximalMarginalRelevance( + topK: 2, + ); + expect(results.map((r) => r.index).toList(), [0, 1]); + expect(results.first.score, greaterThan(results[1].score)); + }); + + test('returns empty for no candidates', () { + final query = Float32List.fromList([1, 0, 0]); + expect( + Model2VecUtils.similaritySearchWithScores(query, const []), + isEmpty, + ); + }); + + test('drops results below the threshold', () { + final query = Float32List.fromList([1.0, 0.0]); + final candidates = [ + Float32List.fromList([0.0, 1.0]), // sim: 0.0 + Float32List.fromList([0.9, 0.1]), // sim: ~0.994 + Float32List.fromList([-1.0, 0.0]), // sim: -1.0 + Float32List.fromList([0.5, 0.5]), // sim: ~0.707 + ]; + + final results = Model2VecUtils.similaritySearchWithScores( query, candidates, - lambda: -0.1, - ), - throwsArgumentError, - ); - }); + topK: 10, + threshold: 0.8, + ); - test('Base64 serialization works bidirectionally', () { - final vector = Float32List.fromList([0.1, 0.2, -0.3, 100.5]); - final base64 = Model2VecUtils.toBase64(vector); - final restored = Model2VecUtils.fromBase64(base64); + expect(results.map((r) => r.index).toList(), [1]); + expect(results.single.score, greaterThan(0.8)); + }); - expect(restored.length, equals(vector.length)); - for (var i = 0; i < vector.length; i++) { - expect(restored[i], equals(vector[i])); - } - }); + test('keeps a candidate scoring exactly the threshold', () { + final query = Float32List.fromList([1.0, 0.0]); + final candidates = [ + Float32List.fromList([1.0, 0.0]), // identical -> cosine 1.0 + Float32List.fromList([0.0, 1.0]), // orthogonal -> cosine 0.0 + ]; - test('pairwiseSimilarity calculates matrix correctly', () { - final listA = [ - Float32List.fromList([1.0, 0.0]), - Float32List.fromList([0.0, 1.0]), - ]; - final listB = [ - Float32List.fromList([1.0, 0.0]), - Float32List.fromList([0.0, 1.0]), - Float32List.fromList([0.707, 0.707]), - ]; - - final result = Model2VecUtils.pairwiseSimilarity(listA, listB); - expect(result.length, 2); - expect(result[0].length, 3); - - // listA[0] vs listB - expect(result[0][0], closeTo(1.0, 1e-6)); - expect(result[0][1], closeTo(0.0, 1e-6)); - expect(result[0][2], closeTo(0.707106, 1e-4)); - - // listA[1] vs listB - expect(result[1][0], closeTo(0.0, 1e-6)); - expect(result[1][1], closeTo(1.0, 1e-6)); - expect(result[1][2], closeTo(0.707106, 1e-4)); - }); - }); + final results = Model2VecUtils.similaritySearchWithScores( + query, + candidates, + topK: 10, + threshold: 1, // boundary is inclusive (>=) + ); - group('Model2VecUtils - Real World', () { - setUpAll(() { - Model2Vec.loadModel('minishlab/potion-base-2M'); - }); + expect(results.map((r) => r.index).toList(), [0]); + expect(results.single.score, closeTo(1.0, 1e-6)); + }); + + test('caps threshold matches at topK', () { + final query = Float32List.fromList([1.0, 0.0]); + final candidates = [ + Float32List.fromList([1.0, 0.0]), // sim 1.0 + Float32List.fromList([0.99, 0.01]), // sim ~0.9999 + Float32List.fromList([0.98, 0.02]), // sim ~0.9998 + ]; - test('semantic similarity makes sense', () { - final vCat = Model2Vec.generateEmbedding('A small cute cat'); - final vKitten = Model2Vec.generateEmbedding('A young little kitten'); - final vSpace = Model2Vec.generateEmbedding( - 'The exploration of outer space and planets', - ); - - final simCatKitten = Model2VecUtils.cosineSimilarity(vCat, vKitten); - final simCatSpace = Model2VecUtils.cosineSimilarity(vCat, vSpace); - - stdout - ..writeln( - 'Similarity (Cat vs Kitten): ' - '${(simCatKitten * 100).toStringAsFixed(2)}%', - ) - ..writeln( - 'Similarity (Cat vs Space): ' - '${(simCatSpace * 100).toStringAsFixed(2)}%', + final results = Model2VecUtils.similaritySearchWithScores( + query, + candidates, + topK: 2, + threshold: 0.5, // all three clear the floor, but topK caps to 2 ); - // Semantically related things should be more similar - expect(simCatKitten, greaterThan(simCatSpace)); - expect(simCatKitten, greaterThan(0.5)); // High similarity - expect(simCatSpace, lessThan(0.5)); // Low similarity + expect(results, hasLength(2)); + expect(results.map((r) => r.index).toList(), [0, 1]); + }); }); - test('task similarity example', () { - final vTask1 = Model2Vec.generateEmbedding( - 'Fix the login bug in production', - ); - final vTask2 = Model2Vec.generateEmbedding( - 'Resolve authentication issue on server', - ); - final vTask3 = Model2Vec.generateEmbedding( - 'Order pizza for the team lunch', - ); - - final sim12 = Model2VecUtils.cosineSimilarity(vTask1, vTask2); - final sim13 = Model2VecUtils.cosineSimilarity(vTask1, vTask3); - - stdout - ..writeln( - 'Similarity (Login vs Auth): ${(sim12 * 100).toStringAsFixed(2)}%', - ) - ..writeln( - 'Similarity (Login vs Pizza): ${(sim13 * 100).toStringAsFixed(2)}%', + group('maximalMarginalRelevance', () { + test('prefers a diverse result over a near-duplicate', () { + final query = Float32List.fromList([1, 0, 0]); + final candidates = [ + Float32List.fromList([1, 0.2, 0]), // 0: most relevant + Float32List.fromList([1, 0.25, 0]), // 1: near-duplicate of 0 + Float32List.fromList([0.6, 0, 0.8]), // 2: relevant-ish but diverse + ]; + + final selected = Model2VecUtils.maximalMarginalRelevance( + query, + candidates, + topK: 2, ); + expect(selected.first, 0); // most relevant first + expect(selected[1], 2); // diversity beats the near-duplicate + }); + + test('with lambda 1.0 is pure relevance order', () { + final query = Float32List.fromList([1, 0, 0]); + final candidates = [ + Float32List.fromList([0.6, 0, 0.8]), + Float32List.fromList([1, 0.2, 0]), + Float32List.fromList([1, 0.25, 0]), + ]; + + final selected = Model2VecUtils.maximalMarginalRelevance( + query, + candidates, + topK: 3, + lambda: 1, + ); + expect(selected, [1, 2, 0]); + }); + + test('with lambda 0.0 selects the maximally-dissimilar spread', () { + final query = Float32List.fromList([1, 0, 0]); + final candidates = [ + Float32List.fromList([1, 0, 0]), // 0: identical to the query + Float32List.fromList([0.9, 0.1, 0]), // 1: near-duplicate of 0 + Float32List.fromList([-1, 0, 0]), // 2: opposite of 0 + ]; + + final selected = Model2VecUtils.maximalMarginalRelevance( + query, + candidates, + topK: 2, + lambda: 0, + ); + // Pure diversity: after the first pick, the maximally-dissimilar + // candidate (2) wins over the near-duplicate (1). + expect(selected, [0, 2]); + }); + + test('returns empty for no candidates', () { + final query = Float32List.fromList([1, 0, 0]); + expect( + Model2VecUtils.maximalMarginalRelevance(query, const []), + isEmpty, + ); + }); + + test('rejects lambda outside [0, 1]', () { + final query = Float32List.fromList([1, 0]); + final candidates = [ + Float32List.fromList([1, 0]), + ]; + expect( + () => Model2VecUtils.maximalMarginalRelevance( + query, + candidates, + lambda: 1.5, + ), + throwsArgumentError, + ); + expect( + () => Model2VecUtils.maximalMarginalRelevance( + query, + candidates, + lambda: -0.1, + ), + throwsArgumentError, + ); + }); + }); + + group('base64', () { + test('round-trips a vector bidirectionally', () { + final vector = Float32List.fromList([0.1, 0.2, -0.3, 100.5]); + final base64 = Model2VecUtils.toBase64(vector); + final restored = Model2VecUtils.fromBase64(base64); + + expect(restored.length, equals(vector.length)); + for (var i = 0; i < vector.length; i++) { + expect(restored[i], equals(vector[i])); + } + }); + }); - expect(sim12, greaterThan(sim13)); + group('pairwiseSimilarity', () { + test('computes the full similarity matrix', () { + final listA = [ + Float32List.fromList([1.0, 0.0]), + Float32List.fromList([0.0, 1.0]), + ]; + final listB = [ + Float32List.fromList([1.0, 0.0]), + Float32List.fromList([0.0, 1.0]), + Float32List.fromList([0.707, 0.707]), + ]; + + final result = Model2VecUtils.pairwiseSimilarity(listA, listB); + expect(result.length, 2); + expect(result[0].length, 3); + + // listA[0] vs listB + expect(result[0][0], closeTo(1.0, 1e-6)); + expect(result[0][1], closeTo(0.0, 1e-6)); + expect(result[0][2], closeTo(0.707106, 1e-4)); + + // listA[1] vs listB + expect(result[1][0], closeTo(0.0, 1e-6)); + expect(result[1][1], closeTo(1.0, 1e-6)); + expect(result[1][2], closeTo(0.707106, 1e-4)); + }); }); }); } diff --git a/test/worker_protocol_test.dart b/test/worker_protocol_test.dart index 8f77902..039253f 100644 --- a/test/worker_protocol_test.dart +++ b/test/worker_protocol_test.dart @@ -32,107 +32,149 @@ class FakeChannel implements Channel { Float32List _vec(List xs) => Float32List.fromList(xs); void main() { - late FakeChannel channel; - late WorkerProtocol protocol; - - setUp(() { - channel = FakeChannel(); - protocol = WorkerProtocol(channel); - }); - - test('sends the request and resolves with the delivered result', () async { - final future = protocol.embedBatch(['a'], maxLength: 8); - expect(channel.sent, hasLength(1)); // sent synchronously - - final result = [ - _vec([1, 2, 3]), - ]; - channel.deliver(result); - expect(await future, same(result)); - }); - - test('serializes: the second request waits for the first reply', () async { - final f1 = protocol.embedBatch(['a'], maxLength: 8); - final f2 = protocol.embedBatch(['b'], maxLength: 8); - - // Only the first request is in flight. - expect(channel.sent, hasLength(1)); - - channel.deliver([ - _vec([1]), - ]); - await f1; - - // Now the second is sent. - expect(channel.sent, hasLength(2)); - channel.deliver([ - _vec([2]), - ]); - await f2; - }); - - test('propagates a typed Model2VecException from the worker', () async { - final future = protocol.embedBatch(['a'], maxLength: 8); - channel.deliver( - const Model2VecException( - Model2VecErrorKind.notInitialized, - 'no model', - 1, - ), - ); - - await expectLater( - future, - throwsA( - isA().having( - (e) => e.kind, - 'kind', + group('WorkerProtocol', () { + late FakeChannel channel; + late WorkerProtocol protocol; + + setUp(() { + channel = FakeChannel(); + protocol = WorkerProtocol(channel); + }); + + test('sends the request and resolves with the delivered result', () async { + final future = protocol.embedBatch(['a'], maxLength: 8); + expect(channel.sent, hasLength(1)); // sent synchronously + + final result = [ + _vec([1, 2, 3]), + ]; + channel.deliver(result); + expect(await future, same(result)); + }); + + test('serializes: the second request waits for the first reply', () async { + final f1 = protocol.embedBatch(['a'], maxLength: 8); + final f2 = protocol.embedBatch(['b'], maxLength: 8); + + // Only the first request is in flight. + expect(channel.sent, hasLength(1)); + + channel.deliver([ + _vec([1]), + ]); + await f1; + + // Now the second is sent. + expect(channel.sent, hasLength(2)); + channel.deliver([ + _vec([2]), + ]); + await f2; + }); + + test('an empty batch resolves with an empty list', () async { + final future = protocol.embedBatch([], maxLength: 8); + expect(channel.sent, hasLength(1)); + channel.deliver([]); + expect(await future, isEmpty); + }); + + test('propagates a typed Model2VecException from the worker', () async { + final future = protocol.embedBatch(['a'], maxLength: 8); + channel.deliver( + const Model2VecException( Model2VecErrorKind.notInitialized, + 'no model', + 1, ), - ), - ); - }); - - test('an unexpected reply shape becomes a StateError', () async { - final future = protocol.embedBatch(['a'], maxLength: 8); - channel.deliver('garbage'); - await expectLater(future, throwsStateError); - }); - - test('close fails queued and in-flight requests and closes the channel', - () async { - final f1 = protocol.embedBatch(['a'], maxLength: 8); - final f2 = protocol.embedBatch(['b'], maxLength: 8); - - // Attach listeners before close so the failed futures aren't reported as - // unhandled (real callers await the future they submitted). - final expect1 = expectLater(f1, throwsA(isA())); - final expect2 = expectLater(f2, throwsA(isA())); - - await protocol.close(); - - await expect1; - await expect2; - expect(channel.closed, isTrue); - }); - - test('embedBatch after close fails fast', () async { - await protocol.close(); - await expectLater( - protocol.embedBatch(['a'], maxLength: 8), - throwsA(isA()), - ); - }); - - test('fails a pending request when the channel closes on its own', () async { - final future = protocol.embedBatch(['a'], maxLength: 8); - final expectation = expectLater( - future, - throwsA(isA()), - ); - // Simulates the worker dying: the channel closes while a request is in - // flight, without protocol.close() first, so onDone reaches the protocol. - await channel.close(); - await expectation; + ); + + await expectLater( + future, + throwsA( + isA().having( + (e) => e.kind, + 'kind', + Model2VecErrorKind.notInitialized, + ), + ), + ); + }); + + test('an unexpected reply shape becomes a StateError', () async { + final future = protocol.embedBatch(['a'], maxLength: 8); + channel.deliver('garbage'); + await expectLater(future, throwsStateError); + }); + + test('drops a stray reply when nothing is awaiting', () async { + // A reply arrives with an empty queue — nothing awaits it, so it must be + // dropped silently and leave the protocol able to serve real requests. + channel.deliver([ + _vec([9]), + ]); + // Let the stray be processed (and dropped) before a real request exists; + // otherwise it would land on the request submitted below. + await pumpEventQueue(); + + final future = protocol.embedBatch(['a'], maxLength: 8); + channel.deliver([ + _vec([1]), + ]); + expect((await future).single.first, 1.0); + }); + + test('close fails queued and in-flight requests and closes the channel', + () async { + final f1 = protocol.embedBatch(['a'], maxLength: 8); + final f2 = protocol.embedBatch(['b'], maxLength: 8); + + // Attach listeners before close so the failed futures aren't reported as + // unhandled (real callers await the future they submitted). + final expect1 = expectLater(f1, throwsA(isA())); + final expect2 = expectLater(f2, throwsA(isA())); + + await protocol.close(); + + await expect1; + await expect2; + expect(channel.closed, isTrue); + }); + + test('embedBatch after close fails fast', () async { + await protocol.close(); + await expectLater( + protocol.embedBatch(['a'], maxLength: 8), + throwsA(isA()), + ); + }); + + test('close is idempotent', () async { + await protocol.close(); + await protocol.close(); + }); + + test('fails a pending request when the channel closes on its own', + () async { + final future = protocol.embedBatch(['a'], maxLength: 8); + final expectation = expectLater( + future, + throwsA(isA()), + ); + // Simulates the worker dying: the channel closes while a request is in + // flight, without protocol.close() first, so onDone reaches the protocol. + await channel.close(); + await expectation; + }, timeout: const Timeout(Duration(seconds: 20))); + + test('rejects new requests after the channel closes underneath', () async { + // The channel closes with nothing outstanding; a *new* request must then + // fail fast rather than wait for a reply that will never come. + await channel.close(); + await expectLater( + protocol.embedBatch(['a'], maxLength: 8), + throwsA(isA()), + ); + }, timeout: const Timeout(Duration(seconds: 20))); }); } diff --git a/test/worker_support.dart b/test/worker_support.dart new file mode 100644 index 0000000..00cb48e --- /dev/null +++ b/test/worker_support.dart @@ -0,0 +1,177 @@ +import 'dart:io'; +import 'dart:isolate'; +import 'dart:typed_data'; + +import 'package:model2vec/src/channel.dart'; +import 'package:model2vec/src/exception.dart'; +import 'package:model2vec/src/worker_protocol.dart'; +import 'package:path/path.dart' as p; + +/// Shared fake isolate entry points for the concurrency tests. +/// +/// Isolate entry points must be top-level functions (they can't be closures), +/// so the fakes that several test files need live here instead of being +/// duplicated. None of them touch the native model — they exercise the +/// isolate/protocol wiring only. + +/// Fake worker entry point: replies with one vector per input text, whose only +/// element is the text's length. No model, no native calls. +void echoEntryPoint(SendPort mainSendPort) { + final receivePort = ReceivePort(); + mainSendPort.send(receivePort.sendPort); + receivePort.listen((message) { + if (message == 'close') { + receivePort.close(); + return; + } + final request = message as EmbedRequest; + mainSendPort.send([ + for (final text in request.batch) + Float32List.fromList([text.length.toDouble()]), + ]); + }); +} + +/// Fake worker entry point that always replies with a typed error, so tests can +/// check that a [Model2VecException] survives the isolate boundary. +void errorEntryPoint(SendPort mainSendPort) { + final receivePort = ReceivePort(); + mainSendPort.send(receivePort.sendPort); + receivePort.listen((message) { + if (message == 'close') { + receivePort.close(); + return; + } + mainSendPort.send( + const Model2VecException( + Model2VecErrorKind.tokenizationFailed, + 'boom', + 6, + ), + ); + }); +} + +/// Fake entry point that dies (closes its port, exiting the isolate) on the +/// first request without ever replying. +void dyingEntryPoint(SendPort mainSendPort) { + final receivePort = ReceivePort(); + mainSendPort.send(receivePort.sendPort); + receivePort.listen((message) { + if (message == 'close') { + receivePort.close(); + return; + } + receivePort.close(); // exit without replying + }); +} + +/// Fake entry point that exits immediately, never completing the handshake. +void noHandshakeEntryPoint(SendPort mainSendPort) { + // Returns at once; the isolate exits before sending its SendPort. +} + +/// Fake entry point whose first message is a plain value rather than a +/// [SendPort], to drive the "bad handshake" branch of [IsolateChannel.start]. +void nonSendPortHandshakeEntryPoint(SendPort mainSendPort) { + mainSendPort.send('not a SendPort'); +} + +/// Fake worker that reports, as the single element of each reply vector, how +/// many requests this specific isolate has handled so far, after a short delay +/// so concurrent calls genuinely overlap. +/// +/// Because each spawned isolate has its own counter, a pool that spreads N +/// concurrent batches across N distinct workers yields all-ones (each worker's +/// first request); a pool that funnelled them to one worker would yield +/// 1, 2, 3, … instead. That makes the dispatch spread observable. +void countingEntryPoint(SendPort mainSendPort) { + final receivePort = ReceivePort(); + mainSendPort.send(receivePort.sendPort); + var handled = 0; + receivePort.listen((message) async { + if (message == 'close') { + receivePort.close(); + return; + } + final request = message as EmbedRequest; + handled++; + final count = handled; + await Future.delayed(const Duration(milliseconds: 50)); + mainSendPort.send([ + for (final _ in request.batch) Float32List.fromList([count.toDouble()]), + ]); + }); +} + +/// Fake worker that fails its first request with a typed error and then behaves +/// like [echoEntryPoint]. Used to check the pool/protocol recover after an +/// error rather than wedging. +void failThenEchoEntryPoint(SendPort mainSendPort) { + final receivePort = ReceivePort(); + mainSendPort.send(receivePort.sendPort); + var failed = false; + receivePort.listen((message) { + if (message == 'close') { + receivePort.close(); + return; + } + final request = message as EmbedRequest; + if (!failed) { + failed = true; + mainSendPort.send( + const Model2VecException(Model2VecErrorKind.panic, 'boom'), + ); + return; + } + mainSendPort.send([ + for (final text in request.batch) + Float32List.fromList([text.length.toDouble()]), + ]); + }); +} + +/// How many [flakyStartEntryPoint] spawns complete their handshake before the +/// rest deliberately fail. Start a pool larger than this to force a partial +/// start failure. +const partialStartSuccessLimit = 2; + +final _partialStartSlots = + Directory(p.join(Directory.systemTemp.path, 'm2v_partial_start_slots')); + +/// Fake entry point that succeeds for the first [partialStartSuccessLimit] +/// spawns and then exits before the handshake, so starting a larger pool leaves +/// some workers running and one (or more) failed. +/// +/// The spawns race in parallel, so they claim a sequence number through +/// exclusive file creation — an atomic cross-isolate counter (isolates share no +/// memory). Call [resetPartialStartSlots] before and after use. +void flakyStartEntryPoint(SendPort mainSendPort) { + final sequence = _claimStartSequence(); + if (sequence >= partialStartSuccessLimit) { + return; // exit before the handshake — this spawn fails + } + echoEntryPoint(mainSendPort); +} + +int _claimStartSequence() { + _partialStartSlots.createSync(recursive: true); + var sequence = 0; + while (true) { + try { + File(p.join(_partialStartSlots.path, 'slot_$sequence')) + .createSync(exclusive: true); + return sequence; + } on FileSystemException { + sequence++; + } + } +} + +/// Clears the [flakyStartEntryPoint] sequence counter so each test run starts +/// fresh regardless of leftovers from a previous run. +void resetPartialStartSlots() { + if (_partialStartSlots.existsSync()) { + _partialStartSlots.deleteSync(recursive: true); + } +} From e18d215f20e5093c6a33f3c1c96ce9b89f3ef3d4 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Tue, 7 Jul 2026 13:11:58 +0300 Subject: [PATCH 12/20] dart format . --- example/scaling_example.dart | 6 ++- lib/src/embedding_index.dart | 7 +-- lib/src/embedding_pool.dart | 2 +- lib/src/load_progress.dart | 3 +- lib/src/model2vec_base.dart | 6 +-- lib/src/utils.dart | 4 +- test/batcher_test.dart | 12 +++-- test/channel_test.dart | 50 ++++++++++++--------- test/chunker_test.dart | 18 ++++---- test/embedding_index_test.dart | 5 ++- test/embedding_pool_test.dart | 76 +++++++++++++++++-------------- test/embedding_worker_test.dart | 37 ++++++++------- test/worker_protocol_test.dart | 79 ++++++++++++++++++--------------- test/worker_support.dart | 10 +++-- 14 files changed, 177 insertions(+), 138 deletions(-) diff --git a/example/scaling_example.dart b/example/scaling_example.dart index 8fc788b..a56f949 100644 --- a/example/scaling_example.dart +++ b/example/scaling_example.dart @@ -49,8 +49,10 @@ Future main() async { // the whole dataset (or all of its vectors) in memory at once. final huge = Stream.fromIterable(List.generate(1000, (i) => 'Item $i')); var streamed = 0; - await for (final _ - in Model2Vec.generateEmbeddingStream(huge, batchSize: 200)) { + await for (final _ in Model2Vec.generateEmbeddingStream( + huge, + batchSize: 200, + )) { streamed++; } stdout.writeln('Streamed and embedded $streamed items.'); diff --git a/lib/src/embedding_index.dart b/lib/src/embedding_index.dart index 8a98a58..7fd3658 100644 --- a/lib/src/embedding_index.dart +++ b/lib/src/embedding_index.dart @@ -110,8 +110,7 @@ class EmbeddingIndex { /// Throws [ArgumentError] if [query]'s length differs from the index's /// [dimension]. Returns an empty list for an empty index. List search(Float32List query, {int topK = 5}) { - final scored = _scoreAll(query) - ..sort((a, b) => b.score.compareTo(a.score)); + final scored = _scoreAll(query)..sort((a, b) => b.score.compareTo(a.score)); if (topK <= 0) { return []; } @@ -127,9 +126,7 @@ class EmbeddingIndex { Float32List query, { required double threshold, }) { - final scored = _scoreAll(query) - .where((r) => r.score >= threshold) - .toList() + final scored = _scoreAll(query).where((r) => r.score >= threshold).toList() ..sort((a, b) => b.score.compareTo(a.score)); return scored; } diff --git a/lib/src/embedding_pool.dart b/lib/src/embedding_pool.dart index d07a339..623b097 100644 --- a/lib/src/embedding_pool.dart +++ b/lib/src/embedding_pool.dart @@ -17,7 +17,7 @@ import 'embedding_worker.dart'; /// contract, same as the single worker). class EmbeddingPool { EmbeddingPool._(this._workers) - : _inFlight = List.filled(_workers.length, 0); + : _inFlight = List.filled(_workers.length, 0); final List _workers; final List _inFlight; diff --git a/lib/src/load_progress.dart b/lib/src/load_progress.dart index fb6c6a9..5ed3ba6 100644 --- a/lib/src/load_progress.dart +++ b/lib/src/load_progress.dart @@ -47,8 +47,7 @@ final class LoadProgress { /// Download completion in `[0.0, 1.0]`, or `null` when the total size is not /// known — before the download starts, on a cache hit, or for a local path. - double? get fraction => - totalBytes > 0 ? bytesDownloaded / totalBytes : null; + double? get fraction => totalBytes > 0 ? bytesDownloaded / totalBytes : null; @override String toString() => diff --git a/lib/src/model2vec_base.dart b/lib/src/model2vec_base.dart index 4351282..27e20f6 100644 --- a/lib/src/model2vec_base.dart +++ b/lib/src/model2vec_base.dart @@ -28,8 +28,7 @@ abstract final class Model2Vec { /// Returns the embedding dimension of the currently loaded model. /// /// Throws a [Model2VecException] if no model has been initialized yet. - static int get embeddingDimension => - _readInt(native.get_embedding_dimension); + static int get embeddingDimension => _readInt(native.get_embedding_dimension); /// Returns the total number of unique tokens in the model's vocabulary. /// @@ -44,8 +43,7 @@ abstract final class Model2Vec { /// Returns the median length (in characters) of tokens in the vocabulary. /// /// Throws a [Model2VecException] if no model has been initialized yet. - static int get medianTokenLength => - _readInt(native.get_median_token_length); + static int get medianTokenLength => _readInt(native.get_median_token_length); /// Reads the current model's metadata into a [ModelInfo]. /// diff --git a/lib/src/utils.dart b/lib/src/utils.dart index 11323e2..f76b60f 100644 --- a/lib/src/utils.dart +++ b/lib/src/utils.dart @@ -92,8 +92,8 @@ final class Model2VecUtils { ]; final ranked = (threshold == null - ? scored - : scored.where((r) => r.score >= threshold).toList()) + ? scored + : scored.where((r) => r.score >= threshold).toList()) ..sort((a, b) => b.score.compareTo(a.score)); return ranked.take(math.min(topK, ranked.length)).toList(growable: false); } diff --git a/test/batcher_test.dart b/test/batcher_test.dart index f3aea90..62608b5 100644 --- a/test/batcher_test.dart +++ b/test/batcher_test.dart @@ -4,8 +4,10 @@ import 'package:test/test.dart'; void main() { group('batched', () { test('groups into full batches with a smaller final remainder', () async { - final result = await batched(Stream.fromIterable([1, 2, 3, 4, 5]), 2) - .toList(); + final result = await batched( + Stream.fromIterable([1, 2, 3, 4, 5]), + 2, + ).toList(); expect(result, [ [1, 2], [3, 4], @@ -14,8 +16,10 @@ void main() { }); test('an exact multiple leaves no remainder', () async { - final result = await batched(Stream.fromIterable([1, 2, 3, 4]), 2) - .toList(); + final result = await batched( + Stream.fromIterable([1, 2, 3, 4]), + 2, + ).toList(); expect(result, [ [1, 2], [3, 4], diff --git a/test/channel_test.dart b/test/channel_test.dart index 7dc55d0..1abca16 100644 --- a/test/channel_test.dart +++ b/test/channel_test.dart @@ -20,18 +20,22 @@ void main() { expect(vectors.map((v) => v.first).toList(), [2.0, 1.0]); }); - test('start fails when the worker exits before the handshake', () async { - await expectLater( - IsolateChannel.start(noHandshakeEntryPoint), - throwsA( - isA().having( - (e) => e.message, - 'message', - contains('exited during startup'), + test( + 'start fails when the worker exits before the handshake', + () async { + await expectLater( + IsolateChannel.start(noHandshakeEntryPoint), + throwsA( + isA().having( + (e) => e.message, + 'message', + contains('exited during startup'), + ), ), - ), - ); - }, timeout: const Timeout(Duration(seconds: 20))); + ); + }, + timeout: const Timeout(Duration(seconds: 20)), + ); test( 'start rejects a non-SendPort handshake instead of hanging', @@ -54,18 +58,22 @@ void main() { timeout: const Timeout(Duration(seconds: 20)), ); - test('incoming closes when the worker exits', () async { - final channel = await IsolateChannel.start(dyingEntryPoint); - addTearDown(channel.close); + test( + 'incoming closes when the worker exits', + () async { + final channel = await IsolateChannel.start(dyingEntryPoint); + addTearDown(channel.close); - final done = Completer(); - channel.incoming.listen((_) {}, onDone: done.complete); - // The dying worker closes its port (exiting) instead of replying; the - // exit must surface as incoming's onDone so a protocol above can't hang. - channel.send((batch: ['a'], maxLength: 8)); + final done = Completer(); + channel.incoming.listen((_) {}, onDone: done.complete); + // The dying worker closes its port (exiting) instead of replying; the + // exit must surface as incoming's onDone so a protocol above can't hang. + channel.send((batch: ['a'], maxLength: 8)); - await done.future; - }, timeout: const Timeout(Duration(seconds: 20))); + await done.future; + }, + timeout: const Timeout(Duration(seconds: 20)), + ); test('close is idempotent', () async { final channel = await IsolateChannel.start(echoEntryPoint); diff --git a/test/chunker_test.dart b/test/chunker_test.dart index bc8773d..dce44cb 100644 --- a/test/chunker_test.dart +++ b/test/chunker_test.dart @@ -11,15 +11,17 @@ void main() { expect(chunkText(' \n\t '), isEmpty); }); - test('splits long text and covers every word exactly once (no overlap)', - () { - final words = List.generate(40, (i) => 'token$i'); - final chunks = chunkText(words.join(' '), maxChars: 40, overlap: 0); + test( + 'splits long text and covers every word exactly once (no overlap)', + () { + final words = List.generate(40, (i) => 'token$i'); + final chunks = chunkText(words.join(' '), maxChars: 40, overlap: 0); - expect(chunks.length, greaterThan(1)); - final emitted = chunks.expand((c) => c.split(' ')).toList(); - expect(emitted, equals(words)); // order preserved, no duplication - }); + expect(chunks.length, greaterThan(1)); + final emitted = chunks.expand((c) => c.split(' ')).toList(); + expect(emitted, equals(words)); // order preserved, no duplication + }, + ); test('overlap re-includes boundary words', () { final words = List.generate(40, (i) => 'token$i').join(' '); diff --git a/test/embedding_index_test.dart b/test/embedding_index_test.dart index 0542821..7c936b1 100644 --- a/test/embedding_index_test.dart +++ b/test/embedding_index_test.dart @@ -59,7 +59,10 @@ void main() { test('addAll and overwrite by id replaces the stored vector', () { final index = EmbeddingIndex() - ..addAll({'a': _v([1, 0]), 'b': _v([0, 1])}); + ..addAll({ + 'a': _v([1, 0]), + 'b': _v([0, 1]), + }); expect(index.length, 2); index.add('a', _v([0, 1])); // overwrite a's vector [1,0] -> [0,1] diff --git a/test/embedding_pool_test.dart b/test/embedding_pool_test.dart index ed27d84..330efc8 100644 --- a/test/embedding_pool_test.dart +++ b/test/embedding_pool_test.dart @@ -70,24 +70,29 @@ void main() { expect(counts, hasLength(pool.size)); }); - test('close rejects in-flight and queued batches without hanging', - () async { - final pool = await EmbeddingPool.start( - size: 2, - entryPoint: countingEntryPoint, - ); - - // Launch more batches than workers so some are queued behind others, then - // close before any of the (delayed) replies land. - final batches = [for (var i = 0; i < 5; i++) pool.embedBatch(['a'])]; - final expectations = [ - for (final batch in batches) - expectLater(batch, throwsA(isA())), - ]; - - await pool.close(); - await Future.wait(expectations); - }, timeout: const Timeout(Duration(seconds: 20))); + test( + 'close rejects in-flight and queued batches without hanging', + () async { + final pool = await EmbeddingPool.start( + size: 2, + entryPoint: countingEntryPoint, + ); + + // Launch more batches than workers so some are queued behind others, then + // close before any of the (delayed) replies land. + final batches = [ + for (var i = 0; i < 5; i++) pool.embedBatch(['a']), + ]; + final expectations = [ + for (final batch in batches) + expectLater(batch, throwsA(isA())), + ]; + + await pool.close(); + await Future.wait(expectations); + }, + timeout: const Timeout(Duration(seconds: 20)), + ); test('recovers after a failing batch and serves later batches', () async { // A single worker fails its first request, then echoes. If the pool @@ -108,22 +113,25 @@ void main() { expect(result.single.first, 4.0); }); - test('start throws and cleans up when a later worker fails to spawn', - () async { - resetPartialStartSlots(); - addTearDown(resetPartialStartSlots); - - // The pool is larger than the number of spawns the entry point lets - // through, so some workers start and at least one fails. start() must - // close the survivors and throw rather than leak orphaned isolates. - await expectLater( - EmbeddingPool.start( - size: partialStartSuccessLimit + 1, - entryPoint: flakyStartEntryPoint, - ), - throwsA(isA()), - ); - }, timeout: const Timeout(Duration(seconds: 20))); + test( + 'start throws and cleans up when a later worker fails to spawn', + () async { + resetPartialStartSlots(); + addTearDown(resetPartialStartSlots); + + // The pool is larger than the number of spawns the entry point lets + // through, so some workers start and at least one fails. start() must + // close the survivors and throw rather than leak orphaned isolates. + await expectLater( + EmbeddingPool.start( + size: partialStartSuccessLimit + 1, + entryPoint: flakyStartEntryPoint, + ), + throwsA(isA()), + ); + }, + timeout: const Timeout(Duration(seconds: 20)), + ); test('close tears down all workers', () async { final pool = await EmbeddingPool.start( diff --git a/test/embedding_worker_test.dart b/test/embedding_worker_test.dart index 9f08c4a..473675e 100644 --- a/test/embedding_worker_test.dart +++ b/test/embedding_worker_test.dart @@ -58,21 +58,28 @@ void main() { await worker.close(); }); - test('a request fails (does not hang) if the worker dies mid-request', - () async { - final worker = await EmbeddingWorker.start(entryPoint: dyingEntryPoint); - addTearDown(worker.close); - await expectLater( - worker.embedBatch(['a']), - throwsA(isA()), - ); - }, timeout: const Timeout(Duration(seconds: 20))); + test( + 'a request fails (does not hang) if the worker dies mid-request', + () async { + final worker = await EmbeddingWorker.start(entryPoint: dyingEntryPoint); + addTearDown(worker.close); + await expectLater( + worker.embedBatch(['a']), + throwsA(isA()), + ); + }, + timeout: const Timeout(Duration(seconds: 20)), + ); - test('start fails if the worker never completes the handshake', () async { - await expectLater( - EmbeddingWorker.start(entryPoint: noHandshakeEntryPoint), - throwsA(isA()), - ); - }, timeout: const Timeout(Duration(seconds: 20))); + test( + 'start fails if the worker never completes the handshake', + () async { + await expectLater( + EmbeddingWorker.start(entryPoint: noHandshakeEntryPoint), + throwsA(isA()), + ); + }, + timeout: const Timeout(Duration(seconds: 20)), + ); }); } diff --git a/test/worker_protocol_test.dart b/test/worker_protocol_test.dart index 039253f..c4942f9 100644 --- a/test/worker_protocol_test.dart +++ b/test/worker_protocol_test.dart @@ -124,22 +124,24 @@ void main() { expect((await future).single.first, 1.0); }); - test('close fails queued and in-flight requests and closes the channel', - () async { - final f1 = protocol.embedBatch(['a'], maxLength: 8); - final f2 = protocol.embedBatch(['b'], maxLength: 8); + test( + 'close fails queued and in-flight requests and closes the channel', + () async { + final f1 = protocol.embedBatch(['a'], maxLength: 8); + final f2 = protocol.embedBatch(['b'], maxLength: 8); - // Attach listeners before close so the failed futures aren't reported as - // unhandled (real callers await the future they submitted). - final expect1 = expectLater(f1, throwsA(isA())); - final expect2 = expectLater(f2, throwsA(isA())); + // Attach listeners before close so the failed futures aren't reported as + // unhandled (real callers await the future they submitted). + final expect1 = expectLater(f1, throwsA(isA())); + final expect2 = expectLater(f2, throwsA(isA())); - await protocol.close(); + await protocol.close(); - await expect1; - await expect2; - expect(channel.closed, isTrue); - }); + await expect1; + await expect2; + expect(channel.closed, isTrue); + }, + ); test('embedBatch after close fails fast', () async { await protocol.close(); @@ -154,27 +156,34 @@ void main() { await protocol.close(); }); - test('fails a pending request when the channel closes on its own', - () async { - final future = protocol.embedBatch(['a'], maxLength: 8); - final expectation = expectLater( - future, - throwsA(isA()), - ); - // Simulates the worker dying: the channel closes while a request is in - // flight, without protocol.close() first, so onDone reaches the protocol. - await channel.close(); - await expectation; - }, timeout: const Timeout(Duration(seconds: 20))); - - test('rejects new requests after the channel closes underneath', () async { - // The channel closes with nothing outstanding; a *new* request must then - // fail fast rather than wait for a reply that will never come. - await channel.close(); - await expectLater( - protocol.embedBatch(['a'], maxLength: 8), - throwsA(isA()), - ); - }, timeout: const Timeout(Duration(seconds: 20))); + test( + 'fails a pending request when the channel closes on its own', + () async { + final future = protocol.embedBatch(['a'], maxLength: 8); + final expectation = expectLater( + future, + throwsA(isA()), + ); + // Simulates the worker dying: the channel closes while a request is in + // flight, without protocol.close() first, so onDone reaches the protocol. + await channel.close(); + await expectation; + }, + timeout: const Timeout(Duration(seconds: 20)), + ); + + test( + 'rejects new requests after the channel closes underneath', + () async { + // The channel closes with nothing outstanding; a *new* request must then + // fail fast rather than wait for a reply that will never come. + await channel.close(); + await expectLater( + protocol.embedBatch(['a'], maxLength: 8), + throwsA(isA()), + ); + }, + timeout: const Timeout(Duration(seconds: 20)), + ); }); } diff --git a/test/worker_support.dart b/test/worker_support.dart index 00cb48e..6e6f4dc 100644 --- a/test/worker_support.dart +++ b/test/worker_support.dart @@ -136,8 +136,9 @@ void failThenEchoEntryPoint(SendPort mainSendPort) { /// start failure. const partialStartSuccessLimit = 2; -final _partialStartSlots = - Directory(p.join(Directory.systemTemp.path, 'm2v_partial_start_slots')); +final _partialStartSlots = Directory( + p.join(Directory.systemTemp.path, 'm2v_partial_start_slots'), +); /// Fake entry point that succeeds for the first [partialStartSuccessLimit] /// spawns and then exits before the handshake, so starting a larger pool leaves @@ -159,8 +160,9 @@ int _claimStartSequence() { var sequence = 0; while (true) { try { - File(p.join(_partialStartSlots.path, 'slot_$sequence')) - .createSync(exclusive: true); + File( + p.join(_partialStartSlots.path, 'slot_$sequence'), + ).createSync(exclusive: true); return sequence; } on FileSystemException { sequence++; From e204dc5870262e4aa97fbe6199ef32c9d9aa0f81 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Tue, 7 Jul 2026 13:13:00 +0300 Subject: [PATCH 13/20] fix warnings --- test/channel_test.dart | 3 ++- test/embedding_pool_test.dart | 4 ++-- test/worker_protocol_test.dart | 11 ++++++----- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/test/channel_test.dart b/test/channel_test.dart index 1abca16..d3bda5b 100644 --- a/test/channel_test.dart +++ b/test/channel_test.dart @@ -67,7 +67,8 @@ void main() { final done = Completer(); channel.incoming.listen((_) {}, onDone: done.complete); // The dying worker closes its port (exiting) instead of replying; the - // exit must surface as incoming's onDone so a protocol above can't hang. + // exit must surface as incoming's onDone so a protocol above can't + // hang. channel.send((batch: ['a'], maxLength: 8)); await done.future; diff --git a/test/embedding_pool_test.dart b/test/embedding_pool_test.dart index 330efc8..6bca04e 100644 --- a/test/embedding_pool_test.dart +++ b/test/embedding_pool_test.dart @@ -78,8 +78,8 @@ void main() { entryPoint: countingEntryPoint, ); - // Launch more batches than workers so some are queued behind others, then - // close before any of the (delayed) replies land. + // Launch more batches than workers so some are queued behind others, + // then close before any of the (delayed) replies land. final batches = [ for (var i = 0; i < 5; i++) pool.embedBatch(['a']), ]; diff --git a/test/worker_protocol_test.dart b/test/worker_protocol_test.dart index c4942f9..6932404 100644 --- a/test/worker_protocol_test.dart +++ b/test/worker_protocol_test.dart @@ -130,8 +130,8 @@ void main() { final f1 = protocol.embedBatch(['a'], maxLength: 8); final f2 = protocol.embedBatch(['b'], maxLength: 8); - // Attach listeners before close so the failed futures aren't reported as - // unhandled (real callers await the future they submitted). + // Attach listeners before close so the failed futures aren't reported + // as unhandled (real callers await the future they submitted). final expect1 = expectLater(f1, throwsA(isA())); final expect2 = expectLater(f2, throwsA(isA())); @@ -165,7 +165,8 @@ void main() { throwsA(isA()), ); // Simulates the worker dying: the channel closes while a request is in - // flight, without protocol.close() first, so onDone reaches the protocol. + // flight, without protocol.close() first, so onDone reaches the + // protocol. await channel.close(); await expectation; }, @@ -175,8 +176,8 @@ void main() { test( 'rejects new requests after the channel closes underneath', () async { - // The channel closes with nothing outstanding; a *new* request must then - // fail fast rather than wait for a reply that will never come. + // The channel closes with nothing outstanding; a *new* request must + // then fail fast rather than wait for a reply that will never come. await channel.close(); await expectLater( protocol.embedBatch(['a'], maxLength: 8), From 5898663e880cd92b6352d5f3250c3527736711b6 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Tue, 7 Jul 2026 13:31:26 +0300 Subject: [PATCH 14/20] CI --- .github/workflows/ci.yml | 83 ++++++++++++++++++++++++++++++++++ README.md | 17 ++++++- dart_test.yaml | 8 ++++ test/channel_test.dart | 2 +- test/embedding_pool_test.dart | 2 +- test/embedding_test.dart | 3 ++ test/model_loading_test.dart | 3 ++ test/worker_protocol_test.dart | 6 +-- 8 files changed, 118 insertions(+), 6 deletions(-) create mode 100644 .github/workflows/ci.yml diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..1d57dd1 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,83 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + +# A newer push to the same branch/PR cancels an in-flight run. +concurrency: + group: ci-${{ github.ref }} + cancel-in-progress: true + +jobs: + # Static checks only — no Rust build, no model download. Fast feedback. + analyze: + name: Format & analyze + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dart-lang/setup-dart@v1 + # `stable` resolves to the latest Dart stable; native assets (this + # package's FFI build hook) work flag-free on 3.12+. + with: + sdk: stable + - run: dart pub get + - name: Check formatting + run: dart format --output=none --set-exit-if-changed . + - name: Analyze + run: dart analyze --fatal-infos + + # Everything except the model-downloading tests. Needs the Rust build (the + # native asset), but never hits the network for a model. + test: + name: Unit tests + runs-on: ubuntu-latest + env: + # Override native/rust-toolchain.toml, whose `targets` list is for + # cross-compiled release builds (Android/iOS/Windows/…). CI only needs the + # host target, and this env var makes cargo ignore the toml entirely. + RUSTUP_TOOLCHAIN: "1.95.0" + steps: + - uses: actions/checkout@v4 + - uses: dart-lang/setup-dart@v1 + with: + sdk: stable + - name: Install Rust ${{ env.RUSTUP_TOOLCHAIN }} (host target only) + run: rustup toolchain install "$RUSTUP_TOOLCHAIN" --profile minimal + - name: Cache cargo registry + build + uses: Swatinem/rust-cache@v2 + with: + workspaces: native + - run: dart pub get + - name: Run unit tests (no model download) + run: dart test -x integration + + # The model-dependent tests. The Hugging Face cache is restored between runs, + # so the ~40 MB download happens once and is reused (keyed on the model ids in + # test/support.dart). + integration: + name: Integration tests + runs-on: ubuntu-latest + env: + RUSTUP_TOOLCHAIN: "1.95.0" + steps: + - uses: actions/checkout@v4 + - uses: dart-lang/setup-dart@v1 + with: + sdk: stable + - name: Install Rust ${{ env.RUSTUP_TOOLCHAIN }} (host target only) + run: rustup toolchain install "$RUSTUP_TOOLCHAIN" --profile minimal + - name: Cache cargo registry + build + uses: Swatinem/rust-cache@v2 + with: + workspaces: native + - name: Cache Hugging Face models + uses: actions/cache@v4 + with: + path: ~/.cache/huggingface + key: hf-models-${{ hashFiles('test/support.dart') }} + restore-keys: hf-models- + - run: dart pub get + - name: Run integration tests (model download cached across runs) + run: dart test -t integration diff --git a/README.md b/README.md index 9630cac..2b90805 100644 --- a/README.md +++ b/README.md @@ -309,9 +309,24 @@ To manually re-build bindings if you modify the Rust C-API (`native/src/lib.rs`) dart run ffigen ``` -To run the test suite: +Running the tests: ```bash +dart test # everything, incl. the model-download integration tests +dart test -x integration # fast lane: unit tests only, no model download +dart test -t integration # only the model-dependent tests +``` + +Integration tests (tagged `integration`) download a small model from Hugging +Face on first run and cache it locally. CI runs formatting + analysis + the unit +lane on every push, and the integration lane with the model cache warmed — see +[`.github/workflows/ci.yml`](.github/workflows/ci.yml). + +Before pushing, mirror the CI checks locally: + +```bash +dart format . +dart analyze --fatal-infos dart test ``` diff --git a/dart_test.yaml b/dart_test.yaml index 7bc9d1b..5439e3b 100644 --- a/dart_test.yaml +++ b/dart_test.yaml @@ -2,3 +2,11 @@ # test mutates that global, so suites must run one at a time — otherwise a # parallel suite can read a model of the wrong dimension mid-switch. concurrency: 1 + +tags: + # Tests that download a model from Hugging Face (network + native model + # load). Excluded from the fast lane via `dart test -x integration`; run on + # their own with `dart test -t integration`. Given extra time for cold-cache + # downloads on CI. + integration: + timeout: 2x diff --git a/test/channel_test.dart b/test/channel_test.dart index d3bda5b..d8e8d04 100644 --- a/test/channel_test.dart +++ b/test/channel_test.dart @@ -67,7 +67,7 @@ void main() { final done = Completer(); channel.incoming.listen((_) {}, onDone: done.complete); // The dying worker closes its port (exiting) instead of replying; the - // exit must surface as incoming's onDone so a protocol above can't + // exit must surface as incoming's onDone so a protocol above can't // hang. channel.send((batch: ['a'], maxLength: 8)); diff --git a/test/embedding_pool_test.dart b/test/embedding_pool_test.dart index 6bca04e..97c6fcd 100644 --- a/test/embedding_pool_test.dart +++ b/test/embedding_pool_test.dart @@ -78,7 +78,7 @@ void main() { entryPoint: countingEntryPoint, ); - // Launch more batches than workers so some are queued behind others, + // Launch more batches than workers so some are queued behind others, // then close before any of the (delayed) replies land. final batches = [ for (var i = 0; i < 5; i++) pool.embedBatch(['a']), diff --git a/test/embedding_test.dart b/test/embedding_test.dart index 755a15a..9ed5606 100644 --- a/test/embedding_test.dart +++ b/test/embedding_test.dart @@ -1,3 +1,6 @@ +@Tags(['integration']) +library; + import 'package:model2vec/model2vec.dart'; import 'package:test/test.dart'; diff --git a/test/model_loading_test.dart b/test/model_loading_test.dart index ffa23c1..b588be2 100644 --- a/test/model_loading_test.dart +++ b/test/model_loading_test.dart @@ -1,3 +1,6 @@ +@Tags(['integration']) +library; + import 'dart:io'; import 'dart:typed_data'; diff --git a/test/worker_protocol_test.dart b/test/worker_protocol_test.dart index 6932404..bb35834 100644 --- a/test/worker_protocol_test.dart +++ b/test/worker_protocol_test.dart @@ -130,7 +130,7 @@ void main() { final f1 = protocol.embedBatch(['a'], maxLength: 8); final f2 = protocol.embedBatch(['b'], maxLength: 8); - // Attach listeners before close so the failed futures aren't reported + // Attach listeners before close so the failed futures aren't reported // as unhandled (real callers await the future they submitted). final expect1 = expectLater(f1, throwsA(isA())); final expect2 = expectLater(f2, throwsA(isA())); @@ -165,7 +165,7 @@ void main() { throwsA(isA()), ); // Simulates the worker dying: the channel closes while a request is in - // flight, without protocol.close() first, so onDone reaches the + // flight, without protocol.close() first, so onDone reaches the // protocol. await channel.close(); await expectation; @@ -176,7 +176,7 @@ void main() { test( 'rejects new requests after the channel closes underneath', () async { - // The channel closes with nothing outstanding; a *new* request must + // The channel closes with nothing outstanding; a *new* request must // then fail fast rather than wait for a reply that will never come. await channel.close(); await expectLater( From 3b8e1cdbf1bde78e3430f95082094e50e41323fe Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Tue, 7 Jul 2026 13:50:10 +0300 Subject: [PATCH 15/20] Update readme --- README.md | 366 ++++++++++++++++++++++-------------------------------- 1 file changed, 146 insertions(+), 220 deletions(-) diff --git a/README.md b/README.md index 2b90805..db62d40 100644 --- a/README.md +++ b/README.md @@ -1,160 +1,126 @@ # model2vec [![pub package](https://img.shields.io/pub/v/model2vec.svg)](https://pub.dev/packages/model2vec) +[![CI](https://github.com/pro100andrey/model2vec/actions/workflows/ci.yml/badge.svg)](https://github.com/pro100andrey/model2vec/actions/workflows/ci.yml) [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT) -**High-performance, local text embeddings for Dart and Flutter.** A Dart wrapper around [model2vec-rs](https://github.com/MinishLab/model2vec-rs) using Rust FFI and Native Assets. Model2Vec creates small, fast, and effective text embeddings by distilling knowledge from large language models into a simple vocabulary-based look-up table. - -## Table of Contents - -- [model2vec](#model2vec) - - [Table of Contents](#table-of-contents) - - [Key Features](#key-features) - - [Recommended Models](#recommended-models) - - [Installation](#installation) - - [Quick Start](#quick-start) - - [Recipes \& Patterns](#recipes--patterns) - - [1. Advanced Batch Processing](#1-advanced-batch-processing) - - [2. Massive Data Streaming](#2-massive-data-streaming) - - [3. Asynchronous Isolate Execution](#3-asynchronous-isolate-execution) - - [4. Vector Math \& Quantization](#4-vector-math--quantization) - - [API Reference](#api-reference) - - [Core Methods (`Model2Vec` class)](#core-methods-model2vec-class) - - [Math Utilities (`Model2VecUtils` class)](#math-utilities-model2vecutils-class) - - [Performance](#performance) - - [Development \& Contributing](#development--contributing) - - [License](#license) - -## Key Features - -- **Extreme Performance:** Built on top of a highly optimized Rust engine. Up to **~1.7x faster** than the official Python implementation, generating embeddings in microseconds. -- **Compact & Quantized:** Models are typically 25MB - 100MB. Perfect for edge computing. -- **Massive Streaming:** Built-in `generateEmbeddingStream` for processing millions of rows without blocking the Event Loop or overflowing RAM. -- **Hugging Face Integration:** Automatically downloads and caches models directly from the Hugging Face Hub. -- **Zero-Stutter Async:** Transparently runs heavy tokenization and math in background Dart Isolates using `Async` methods. -- **Vector Utilities:** Ships with high-performance mathematical tools (`cosineSimilarity`, `quantizeToInt8`, `similaritySearchWithScores`, etc.). - -## Recommended Models - -Model2Vec provides a variety of pre-trained models optimized for different use cases. These can be loaded directly via their Hugging Face model ID. - -| Model ID | Language | Distilled From | Params | Dimension | Size | -| -------- | -------- | -------------- | ------ | --------- | ---- | -| [`minishlab/potion-base-32M`](https://huggingface.co/minishlab/potion-base-32M) | English | bge-base-en-v1.5 | 32.3M | 512 | ~150MB | -| [`minishlab/potion-multilingual-128M`](https://huggingface.co/minishlab/potion-multilingual-128M) | Multi | bge-m3 | 128M | 768 | ~500MB | -| [`minishlab/potion-retrieval-32M`](https://huggingface.co/minishlab/potion-retrieval-32M) | English | bge-base-en-v1.5 | 32.3M | 512 | ~150MB | -| [`minishlab/potion-code-16M`](https://huggingface.co/minishlab/potion-code-16M) | Code | CodeRankEmbed | 16M | 384 | ~80MB | -| [`minishlab/potion-base-8M`](https://huggingface.co/minishlab/potion-base-8M) | English | bge-base-en-v1.5 | 7.5M | 256 | ~50MB | -| [`minishlab/potion-base-4M`](https://huggingface.co/minishlab/potion-base-4M) | English | bge-base-en-v1.5 | 3.7M | 128 | ~30MB | -| [`minishlab/potion-base-2M`](https://huggingface.co/minishlab/potion-base-2M) | English | bge-base-en-v1.5 | 1.8M | 64 | ~25MB | +**Fast, on-device text embeddings for Dart & Flutter.** A wrapper around +[model2vec-rs](https://github.com/MinishLab/model2vec-rs) via Rust FFI and Native +Assets. Model2Vec turns text into vectors with a static vocabulary lookup — not a +transformer — so embeddings are generated in **microseconds**, with no server, no +Python, and no network after the model is cached. -## Installation - -Add `model2vec` to your `pubspec.yaml`: +## Features -```yaml -dependencies: - model2vec: any -``` +- **Fast & local** — embeddings in microseconds; fully offline once a model is cached. +- **Hugging Face models** — load any Potion model by id (auto-download + cache), a local directory, or raw bytes. +- **Scales** — batch embedding (Rust SIMD), background isolates, a streaming API for millions of rows, and a worker pool across CPU cores. +- **Built-in retrieval** — an on-device `EmbeddingIndex` (cosine search, int8 quantization, disk persistence) plus vector math (cosine, MMR, pooling, quantization). +- **Native Assets** — the Rust library builds automatically on `pub get`; nothing to link, bundle, or ship. -Or add it using the command line: +## Installation ```bash dart pub add model2vec ``` -*Requires **Dart SDK**: 3.10.0+ and **Rust toolchain**: 1.86.0+ (to build the native library via Native Assets).* +Requirements: -## Quick Start +- **Dart SDK** 3.10+. +- **Rust toolchain** ([rustup](https://rustup.rs)). The native library is compiled + automatically via Native Assets; the exact version is pinned in + [`native/rust-toolchain.toml`](native/rust-toolchain.toml) and installed for you. + +## Quick start ```dart import 'package:model2vec/model2vec.dart'; void main() { - // Model2Vec is a stateless namespace of static methods — there is one - // active model per process, loaded automatically via Native Assets. - - // Load a model from Hugging Face - Model2Vec.loadModel('minishlab/potion-base-2M'); + // One active model per process; the native library loads automatically. + Model2Vec.loadModel('minishlab/potion-base-2M'); // downloads on first run - // Generate an embedding - final embedding = Model2Vec.generateEmbedding('Dart FFI is blazingly fast 🚀'); + final a = Model2Vec.generateEmbedding('I love programming in Dart'); + final b = Model2Vec.generateEmbedding('Dart is a great language'); - print('Vector dimension: ${Model2Vec.embeddingDimension}'); - print('Vocabulary size: ${Model2Vec.vocabularySize}'); + print(Model2VecUtils.cosineSimilarity(a, b)); // ~0.8 — semantically close } ``` -> **Migrating from 1.x?** The instance API is gone. Replace -> `Model2Vec.instance.foo(...)` with `Model2Vec.foo(...)`, drop any -> `Model2Vec.boot(...)` / `Model2Vec(lib)` calls (library resolution is now -> automatic), and read `Model2Vec.recommendedModels` instead of calling -> `getRecommendedModels()`. See the [CHANGELOG](CHANGELOG.md) for the full list. +> In a Flutter app, use `loadModelAsync` instead so the first download never +> blocks the UI (see [Loading off the main thread](#loading-off-the-main-thread)). + +**Migrating from 1.x?** The instance API is gone: replace +`Model2Vec.instance.foo(...)` with `Model2Vec.foo(...)`, drop `Model2Vec.boot(...)` +/ `Model2Vec(lib)` (resolution is automatic), and read +`Model2Vec.recommendedModels` instead of `getRecommendedModels()`. Full list in the +[CHANGELOG](CHANGELOG.md). -## Recipes & Patterns +## Models -### 1. Advanced Batch Processing +Load any of these by id, or browse the typed catalog at `Model2Vec.recommendedModels`. -Process multiple strings at once for maximum hardware utilization. You can control sequence truncation with `maxLength`. +| Model | Params | Dim | Language | Best for | +| ----- | ------ | --- | -------- | -------- | +| [`potion-base-2M`](https://huggingface.co/minishlab/potion-base-2M) | 1.8M | 64 | English | Smallest, very fast | +| [`potion-base-4M`](https://huggingface.co/minishlab/potion-base-4M) | 3.7M | 128 | English | Small and efficient | +| [`potion-base-8M`](https://huggingface.co/minishlab/potion-base-8M) | 7.5M | 256 | English | Balanced default | +| [`potion-base-32M`](https://huggingface.co/minishlab/potion-base-32M) | 32.3M | 512 | English | Large and accurate | +| [`potion-retrieval-32M`](https://huggingface.co/minishlab/potion-retrieval-32M) | 32.3M | 512 | English | RAG / retrieval | +| [`potion-code-16M`](https://huggingface.co/minishlab/potion-code-16M) | 16M | 384 | Code | Code search | +| [`potion-multilingual-128M`](https://huggingface.co/minishlab/potion-multilingual-128M) | 128M | 768 | 101 languages | Multilingual tasks | + +## Guide + +Runnable versions of everything below live in [`example/`](example/). + +### Generating embeddings ```dart -final texts = ['Dart', 'Rust', 'Flutter']; +// Single vector +final v = Model2Vec.generateEmbedding('Hello world'); -final embeddings = Model2Vec.generateBatchEmbeddings( - texts, - maxLength: 256, // Truncate strings longer than 256 tokens +// Batch — one native call, SIMD across the batch. maxLength truncates long input. +final batch = Model2Vec.generateBatchEmbeddings( + ['Dart', 'Rust', 'Flutter'], + maxLength: 256, ); ``` -### 2. Massive Data Streaming - -When reading gigabytes of text from files or databases, loading everything into memory will crash the app. Use the **Streaming API** to handle data in chunks automatically. +For datasets too large to hold in memory, stream them — the input is processed in +batches and the output is a `Stream`: ```dart -import 'dart:convert'; -import 'dart:io'; - -Future processHugeFile() async { - final fileStream = File('massive_dataset.txt') - .openRead() - .transform(utf8.decoder) - .transform(const LineSplitter()); - - // Converts a Stream into a Stream - final embeddingStream = Model2Vec.generateEmbeddingStream( - fileStream, - batchSize: 500, // Process 500 strings at a time - useIsolate: true, // Run math in background threads - ); - - await for (final embedding in embeddingStream) { - saveToDb(embedding); // Memory safe! - } +final vectors = Model2Vec.generateEmbeddingStream( + lines, // a Stream from a file, DB, socket… + batchSize: 500, + useIsolate: true, // run the work off the main isolate +); + +await for (final v in vectors) { + save(v); // bounded memory, whatever the input size } ``` -### 3. Asynchronous Isolate Execution - -Never block the main thread. If you are building a Flutter app, always use the `Async` variants to perform generation in a background `Isolate`. +### Loading off the main thread -Loading is the heaviest step of all — the first time a model is fetched it downloads tens to hundreds of MB. `loadModelAsync` runs that load on a background isolate; because the native model is a single process-global, the loaded model is then visible to every isolate. +The first load of a model downloads tens to hundreds of MB. `loadModelAsync` runs +it on a background isolate; because the native model is a single process-global, it +becomes visible to every isolate once loaded. ```dart -// Load off the main thread so the first download never freezes the UI. await Model2Vec.loadModelAsync('minishlab/potion-base-2M'); - -final embedding = await Model2Vec.generateEmbeddingAsync('A very long text...'); -final batch = await Model2Vec.generateBatchEmbeddingsAsync(['A', 'B', 'C']); ``` -Want a progress bar for that first download? `loadModelWithProgress` loads on a background isolate and streams `LoadProgress` snapshots. Byte counts appear while the weights download (a cached model or local path jumps straight to `done`); the stream always ends on `LoadPhase.done` and surfaces any load failure as an error event. +To show a progress bar for that download, use `loadModelWithProgress` — it streams +`LoadProgress` snapshots and always ends on `LoadPhase.done` (a cached model or +local path jumps straight there; a failed load surfaces as a stream error): ```dart -await for (final p in Model2Vec.loadModelWithProgress('minishlab/potion-base-2M')) { +await for (final p in Model2Vec.loadModelWithProgress('minishlab/potion-base-8M')) { switch (p.phase) { case LoadPhase.downloading: - // p.fraction is 0.0..1.0, or null until the total size is known. - print('Downloading ${((p.fraction ?? 0) * 100).toStringAsFixed(0)}%'); + print('${((p.fraction ?? 0) * 100).round()}%'); // fraction is null until size is known case LoadPhase.resolving || LoadPhase.parsing: print('Preparing…'); case LoadPhase.done: @@ -163,166 +129,126 @@ await for (final p in Model2Vec.loadModelWithProgress('minishlab/potion-base-2M' } ``` -### 4. Vector Math & Quantization +### Similarity & vector math -The library ships with `Model2VecUtils` — a powerful suite of math operations tuned for embeddings. +`Model2VecUtils` is a set of static helpers tuned for embeddings. ```dart final query = Model2Vec.generateEmbedding('cat'); -final candidates = [ - Model2Vec.generateEmbedding('dog'), - Model2Vec.generateEmbedding('space'), +final docs = [ + Model2Vec.generateEmbedding('kitten'), + Model2Vec.generateEmbedding('rocket'), ]; -// 1. Semantic Similarity (Cosine) -final sim = Model2VecUtils.cosineSimilarity(query, candidates[0]); +// Cosine similarity of two vectors (-1.0 … 1.0) +Model2VecUtils.cosineSimilarity(query, docs[0]); -// 2. Threshold Searching — (index, score) pairs for all matches > 80% -final matches = Model2VecUtils.similaritySearchWithScores( - query, candidates, threshold: 0.8, topK: candidates.length, -); +// Rank an in-memory list: (index, score) pairs, top-K with an optional threshold +Model2VecUtils.similaritySearchWithScores(query, docs, topK: 5, threshold: 0.5); -// 3. Scalar Quantization (Compress Float32 to Int8 to save 4x RAM) -final compressed = Model2VecUtils.quantizeToInt8(query); - -// 4. Mean Pooling (Average multiple vectors into one) -final sentenceVector = Model2VecUtils.meanPooling(candidates); - -// 5. DB Serialization -final base64String = Model2VecUtils.toBase64(query); +// Compress Float32 → Int8 (¼ the memory) and serialize for storage +final int8 = Model2VecUtils.quantizeToInt8(query); +final str = Model2VecUtils.toBase64(query); ``` -### 5. Local Retrieval (RAG) with `EmbeddingIndex` +### Local retrieval (RAG) -Build a searchable, persistable index of your documents entirely on-device — chunk, embed, store, and query without a server. +Build a searchable, persistable index entirely on-device — chunk, embed, store, +query. No server. ```dart -// 1. Split long documents into overlapping passages -final passages = chunkText(document, maxChars: 800, overlap: 100); - -// 2. Embed and index them, keeping each passage as the entry's payload so a -// hit carries its text directly (int8 storage cuts memory ~4x) -final index = EmbeddingIndex(quantized: true); -for (var i = 0; i < passages.length; i++) { - index.add( - 'passage-$i', - Model2Vec.generateEmbedding(passages[i]), - payload: passages[i], - ); +// Split documents into overlapping passages, then embed and index them. Storing +// each passage as the entry's payload means a hit carries its text directly. +final index = EmbeddingIndex(quantized: true); // int8 storage, ~¼ the memory +for (final passage in chunkText(document, maxChars: 800, overlap: 100)) { + index.add(passage, Model2Vec.generateEmbedding(passage), payload: passage); } -// 3. Query — returns SearchResult(id, score, payload), most similar first +// Query — SearchResult(id, score, payload), most similar first. final query = Model2Vec.generateEmbedding('How do I reset my password?'); -final hits = index.search(query, topK: 5); -for (final hit in hits) { +for (final hit in index.search(query, topK: 5)) { print('${hit.score.toStringAsFixed(3)} ${hit.payload}'); } -// 4. Persist to disk and reload later -final bytes = index.toBytes(); -final reloaded = EmbeddingIndex.fromBytes(bytes); - -// (Optional) Rerank a candidate vector list for diverse results (MMR): -final diverse = Model2VecUtils.maximalMarginalRelevance( - query, candidateVectors, topK: 5, lambda: 0.5, -); +// Persist and reload — no model needed to load, only to embed new queries. +final reloaded = EmbeddingIndex.fromBytes(index.toBytes()); ``` -### 6. Parallel Embedding & Lifecycle +Use `Model2VecUtils.maximalMarginalRelevance` to rerank for diverse results. + +### Parallel embedding & lifecycle ```dart -// Embed across all CPU cores with a pool of worker isolates -final pool = await EmbeddingPool.start(); // defaults to core count +// Fan batches across CPU cores with a pool of worker isolates. +final pool = await EmbeddingPool.start(); // defaults to the core count final results = await pool.embedBatches(listOfBatches); await pool.close(); -// Non-throwing state check + free native memory when done +// State & teardown. if (Model2Vec.isInitialized) { - final info = Model2Vec.modelInfo; // dim, vocab, normalized, median - print('Loaded model with dimension ${info.dimension}'); + final info = Model2Vec.modelInfo; // dimension, vocabulary, normalized, median + print('dimension ${info.dimension}'); } -Model2Vec.unloadModel(); // releases the native model +Model2Vec.unloadModel(); // frees the native model ``` -## API Reference - -### Core Methods (`Model2Vec` class) - -| Method / Property | Description | -| ----------------- | ----------- | -| `loadModel(path)` | Loads the model from a Hugging Face repo ID or local path. | -| `loadModelAdvanced(...)` | Advanced load with custom `cacheDirectory`, `hfToken`, or `normalize` overrides. | -| `loadModelFromBytes(...)` | Loads the model directly from raw `Uint8List` bytes (`model.safetensors`, `tokenizer.json`, etc). | -| `loadModelAsync(path)` | Loads a model on a background isolate; await it so the first download never blocks the UI. | -| `loadModelAdvancedAsync(...)` | Async form of `loadModelAdvanced` (background isolate). | -| `loadModelWithProgress(path)` | Loads on a background isolate and returns a `Stream` reporting download progress; ends on `LoadPhase.done`. | -| `recommendedModels` | A typed `List` catalog of officially recommended Potion models (offline). | -| `tokenize(text)` | Runs the internal BPE tokenizer and returns a `List`. | -| `generateEmbedding(text)` | Synchronously generates a `Float32List` embedding vector. | -| `generateBatchEmbeddings(texts)` | Synchronously generates embeddings for a `List` using Rust SIMD. | -| `generateEmbeddingAsync(text)` | Asynchronously generates an embedding in a background `Isolate`. | -| `generateEmbeddingStream(stream)` | Processes a huge `Stream` into a `Stream` in batches. | -| `isInitialized` | Non-throwing check for whether a model is currently loaded. | -| `modelInfo` | All model metadata in one `ModelInfo` (dimension, vocabulary, normalized, median). | -| `unloadModel()` | Unloads the active model and frees its native memory. | -| `embeddingDimension` | Property returning the vector size (e.g., 256, 384, 512). | -| `vocabularySize` | Property returning the number of tokens in the model's vocabulary. | - -### Math Utilities (`Model2VecUtils` class) - -| Method | Description | -| ------ | ----------- | -| `cosineSimilarity(a, b)` | Calculates cosine similarity (-1.0 to 1.0) between two vectors. | -| `cosineDistance(a, b)` | Calculates cosine distance (0.0 to 2.0). | -| `euclideanDistance(a, b)` | Calculates Euclidean (L2) distance. | -| `similaritySearchWithScores(query, docs)` | Ranks an in-memory vector list by cosine similarity; returns `(index, score)` pairs, top-K with an optional score `threshold`. | -| `quantizeToInt8(vector)` | Compresses a `Float32List` into an `Int8List` (4x memory savings). | -| `normalize(vector)` | Applies L2 normalization to a vector. | -| `meanPooling(vectors)` | Averages multiple vectors into a single vector. | -| `toBase64` / `fromBase64` | Serializes/Deserializes a vector to/from a Base64 string for DB storage. | +## API reference -## Performance +Full docs are generated on [pub.dev](https://pub.dev/documentation/model2vec/latest/). +The essentials: -`model2vec` uses highly optimized FFI bindings. For mathematical operations on embeddings, Dart handles single-vector math natively with zero-overhead, while batch generation leverages Rust's SIMD (auto-vectorization) capabilities. +**`Model2Vec`** — `loadModel` · `loadModelAdvanced` · `loadModelFromBytes` · +`loadModelAsync` · `loadModelAdvancedAsync` · `loadModelWithProgress` · +`generateEmbedding` · `generateBatchEmbeddings` · `generateEmbeddingAsync` · +`generateBatchEmbeddingsAsync` · `generateEmbeddingStream` · `tokenize` · +`isInitialized` · `modelInfo` · `unloadModel` · `recommendedModels` · +`embeddingDimension` · `vocabularySize` · `isNormalized` · `medianTokenLength`. -Here is a performance benchmark run on a typical machine (AOT compiled): +**`Model2VecUtils`** — `cosineSimilarity` · `cosineDistance` · `euclideanDistance` · +`dotProduct` · `similaritySearchWithScores` · `maximalMarginalRelevance` · +`normalize` · `meanPooling` · `quantizeToInt8` · `dequantizeInt8` · +`pairwiseSimilarity` · `toBase64` · `fromBase64`. -| Model | Load Time (Cache) | Single Embedding | Batch (32) | -| --- | --- | --- | --- | -| `minishlab/potion-base-2M` | ~40 ms | 372.9 μs | 3.85 ms | -| `minishlab/potion-base-4M` | ~40 ms | 363.7 μs | 4.19 ms | -| `minishlab/potion-base-8M` | ~40 ms | 382.1 μs | 5.60 ms | -| `minishlab/potion-base-32M` | ~120 ms | 452.6 μs | 6.79 ms | -| `minishlab/potion-multilingual-128M` | ~1050 ms | 416.1 μs | 5.38 ms | +**Types** — `EmbeddingIndex` (on-device vector store), `EmbeddingPool` (parallel +embedding), `chunkText` (overlapping chunker), `ModelInfo`, `LoadProgress` / +`LoadPhase`, `RecommendedModel`, `Model2VecException` / `Model2VecErrorKind`. -> *Note: Initial load times may vary slightly based on the disk speed. Generating an embedding takes just **a few microseconds** per string.* +## Performance + +Single-vector math runs natively in Dart with zero FFI overhead; batch generation +uses the Rust engine's SIMD. Benchmarked on a typical machine (AOT compiled): -- `similaritySearchWithScores` over 100,000 vectors takes **<100ms** in pure Dart. +| Model | Load (cached) | Single embedding | Batch of 32 | +| ----- | ------------- | ---------------- | ----------- | +| `potion-base-2M` | ~40 ms | 373 μs | 3.85 ms | +| `potion-base-4M` | ~40 ms | 364 μs | 4.19 ms | +| `potion-base-8M` | ~40 ms | 382 μs | 5.60 ms | +| `potion-base-32M` | ~120 ms | 453 μs | 6.79 ms | +| `potion-multilingual-128M` | ~1050 ms | 416 μs | 5.38 ms | -## Development & Contributing +`similaritySearchWithScores` over 100,000 vectors takes **< 100 ms** in pure Dart. -The library uses Dart Native Assets, meaning `cargo build` is invoked automatically when running Dart code. +## Development -To manually re-build bindings if you modify the Rust C-API (`native/src/lib.rs`): +The Rust library builds automatically through Native Assets (`cargo build` runs when +you run Dart code). Regenerate the FFI bindings after changing the C API in +[`native/src/lib.rs`](native/src/lib.rs): ```bash dart run ffigen ``` -Running the tests: +Testing — model-dependent tests are tagged `integration` (they download a small +model on first run and cache it): ```bash -dart test # everything, incl. the model-download integration tests +dart test # everything dart test -x integration # fast lane: unit tests only, no model download dart test -t integration # only the model-dependent tests ``` -Integration tests (tagged `integration`) download a small model from Hugging -Face on first run and cache it locally. CI runs formatting + analysis + the unit -lane on every push, and the integration lane with the model cache warmed — see -[`.github/workflows/ci.yml`](.github/workflows/ci.yml). - -Before pushing, mirror the CI checks locally: +Before pushing, mirror the CI checks +([`.github/workflows/ci.yml`](.github/workflows/ci.yml)): ```bash dart format . @@ -332,4 +258,4 @@ dart test ## License -This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details. +MIT — see [LICENSE](LICENSE). From 18466c83c4a56f1cccc9c85eca802a822e0b8ca9 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Tue, 7 Jul 2026 14:39:02 +0300 Subject: [PATCH 16/20] perf(native): flat-buffer pooling, LTO, drop ndarray/libc/serde; enable clippy --- .github/workflows/ci.yml | 2 + native/Cargo.lock | 76 ------------------------------- native/Cargo.toml | 10 ++-- native/src/lib.rs | 41 ++++++++--------- native/src/model.rs | 98 +++++++++++++++++++++------------------- 5 files changed, 79 insertions(+), 148 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1d57dd1..9b4b0b1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -49,6 +49,8 @@ jobs: uses: Swatinem/rust-cache@v2 with: workspaces: native + - name: Clippy (native) + run: cargo clippy --manifest-path native/Cargo.toml --all-targets -- -D warnings - run: dart pub get - name: Run unit tests (no model download) run: dart test -x integration diff --git a/native/Cargo.lock b/native/Cargo.lock index 2a4342d..b5069cc 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -43,12 +43,6 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" -[[package]] -name = "autocfg" -version = "1.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" - [[package]] name = "base64" version = "0.13.1" @@ -709,10 +703,7 @@ dependencies = [ "anyhow", "half", "hf-hub", - "libc", - "ndarray", "safetensors", - "serde", "serde_json", "tokenizers", "ureq", @@ -734,16 +725,6 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" -[[package]] -name = "matrixmultiply" -version = "0.3.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" -dependencies = [ - "autocfg", - "rawpointer", -] - [[package]] name = "memchr" version = "2.8.0" @@ -788,21 +769,6 @@ dependencies = [ "syn", ] -[[package]] -name = "ndarray" -version = "0.17.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" -dependencies = [ - "matrixmultiply", - "num-complex", - "num-integer", - "num-traits", - "portable-atomic", - "portable-atomic-util", - "rawpointer", -] - [[package]] name = "nom" version = "7.1.3" @@ -813,39 +779,12 @@ dependencies = [ "minimal-lexical", ] -[[package]] -name = "num-complex" -version = "0.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" -dependencies = [ - "num-traits", -] - [[package]] name = "num-conv" version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" -[[package]] -name = "num-integer" -version = "0.1.46" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" -dependencies = [ - "num-traits", -] - -[[package]] -name = "num-traits" -version = "0.2.19" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" -dependencies = [ - "autocfg", -] - [[package]] name = "once_cell" version = "1.21.4" @@ -910,15 +849,6 @@ version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" -[[package]] -name = "portable-atomic-util" -version = "0.2.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" -dependencies = [ - "portable-atomic", -] - [[package]] name = "potential_utf" version = "0.1.5" @@ -996,12 +926,6 @@ dependencies = [ "getrandom 0.3.4", ] -[[package]] -name = "rawpointer" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" - [[package]] name = "rayon" version = "1.12.0" diff --git a/native/Cargo.toml b/native/Cargo.toml index e40c966..0c4f45d 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -7,13 +7,17 @@ edition = "2021" crate-type = ["staticlib", "cdylib"] [dependencies] -libc = "0.2.186" anyhow = "1.0.102" serde_json = "1.0.150" -serde = { version = "1.0.228", features = ["derive"] } tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } safetensors = "0.7.0" -ndarray = "0.17.2" hf-hub = { version = "0.5.0", default-features = false, features = ["ureq"] } ureq = "3.3.0" half = "2.7.1" + +# Release is the profile the build hook uses. `panic = "unwind"` is left at the +# default on purpose — the FFI boundary relies on catch_unwind to turn panics +# into typed errors instead of aborting the host process. +[profile.release] +lto = true +codegen-units = 1 diff --git a/native/src/lib.rs b/native/src/lib.rs index f7625a1..881e98e 100644 --- a/native/src/lib.rs +++ b/native/src/lib.rs @@ -1,3 +1,9 @@ +// FFI entry points take raw pointers from the Dart side and dereference them +// after null-checking — that is the C-ABI contract, not something clippy can +// see the safety of. Marking every export `unsafe` would only move the same +// obligation to the (non-Rust) call sites, so silence the lint crate-wide. +#![allow(clippy::not_unsafe_ptr_arg_deref)] + use anyhow::Result; use std::any::Any; use std::ffi::{CStr, CString}; @@ -47,13 +53,13 @@ where e.code } Err(panic) => { - write_error(out_error, Some(&panic_message(&panic))); + write_error(out_error, Some(&panic_message(&*panic))); CODE_PANIC } } } -fn panic_message(panic: &Box) -> String { +fn panic_message(panic: &(dyn Any + Send)) -> String { if let Some(s) = panic.downcast_ref::<&str>() { format!("panic in native code: {s}") } else if let Some(s) = panic.downcast_ref::() { @@ -272,13 +278,10 @@ pub extern "C" fn generate_embedding( } let text_str = unsafe { CStr::from_ptr(text) }.to_string_lossy().into_owned(); with_model(|model| { - let results = model - .encode_with_args(&[text_str], Some(max_length), 1) + // One sentence in, exactly `dim` floats out. + let embedding = model + .encode_flat(&[text_str], Some(max_length), 1) .map_err(|e| FfiError::new(CODE_TOKENIZATION_FAILED, format!("{e}")))?; - let embedding = results - .into_iter() - .next() - .ok_or_else(|| FfiError::new(CODE_EMPTY_RESULT, "embedding result was empty"))?; write_usize(out_dim, embedding.len()); alloc_floats(out_data, embedding); Ok(()) @@ -311,25 +314,19 @@ pub extern "C" fn generate_batch_embeddings_advanced( } with_model(|model| { - let results = model - .encode_with_args(&texts, Some(max_length), batch_size) + // Pooled straight into one flat buffer. dim + data are produced + // under this single read lock, so the dimension the caller frees + // against always matches the data. + let dim = model.dim(); + let flat = model + .encode_flat(&texts, Some(max_length), batch_size) .map_err(|e| FfiError::new(CODE_TOKENIZATION_FAILED, format!("{e}")))?; - if results.len() != count { + if flat.len() != count * dim { return Err(FfiError::new( CODE_EMPTY_RESULT, - format!("expected {count} embeddings, got {}", results.len()), + format!("expected {} floats, got {}", count * dim, flat.len()), )); } - // dim + data are produced under this single read lock, so the - // dimension the caller frees against always matches the data. - let dim = model.dim(); - let mut flat = Vec::with_capacity(count * dim); - for embedding in &results { - if embedding.len() != dim { - return Err(FfiError::new(CODE_EMPTY_RESULT, "embedding dimension mismatch")); - } - flat.extend_from_slice(embedding); - } write_usize(out_dim, dim); write_usize(out_count, count); alloc_floats(out_data, flat); diff --git a/native/src/model.rs b/native/src/model.rs index 7146a55..c1021df 100644 --- a/native/src/model.rs +++ b/native/src/model.rs @@ -3,7 +3,6 @@ use half::f16; use hf_hub::api::sync::{ApiBuilder, ApiRepo}; use hf_hub::api::Progress; use hf_hub::Cache; -use ndarray::{Array2, CowArray, Ix2}; use safetensors::{tensor::Dtype, SafeTensors}; use serde_json::Value; use std::borrow::Cow; @@ -77,11 +76,18 @@ impl Progress for AtomicProgress { } } -/// Static embedding model for Model2Vec +/// Static embedding model for Model2Vec. +/// +/// `embeddings` is a flat, row-major `rows * cols` buffer: row `r` (a token's +/// vector) is `embeddings[r * cols .. (r + 1) * cols]`. Keeping it a plain slice +/// (rather than an ndarray) lets the pooling loop iterate contiguous `&[f32]` +/// rows, which the compiler auto-vectorizes. #[derive(Debug, Clone)] pub struct StaticModel { pub tokenizer: Tokenizer, - embeddings: CowArray<'static, f32, Ix2>, + embeddings: Cow<'static, [f32]>, + rows: usize, + cols: usize, weights: Option>, token_mapping: Option>, normalize: bool, @@ -282,10 +288,11 @@ impl StaticModel { return Err(anyhow!("embeddings length {} != rows {} * cols {}", embeddings.len(), rows, cols)); } let (median_token_length, unk_token_id) = Self::compute_metadata(&tokenizer)?; - let embeddings = Array2::from_shape_vec((rows, cols), embeddings).context("failed to build embeddings array")?; Ok(Self { tokenizer, - embeddings: CowArray::from(embeddings), + embeddings: Cow::Owned(embeddings), + rows, + cols, weights: weights.map(Cow::Owned), token_mapping: token_mapping.map(Cow::Owned), normalize, @@ -313,13 +320,26 @@ impl StaticModel { s.char_indices().nth(max_tokens.saturating_mul(median_len)).map_or(s, |(byte_idx, _)| &s[..byte_idx]) } - pub fn encode_with_args( + /// The `dim`-length embedding row for token `r`, or an error if `r` is + /// outside the table (a malformed model), so a bad id becomes a typed error + /// rather than an out-of-bounds panic. + fn row(&self, r: usize) -> Result<&[f32]> { + let dim = self.cols; + self.embeddings + .get(r * dim..r * dim + dim) + .ok_or_else(|| anyhow!("token row {r} out of range (vocab size {})", self.rows)) + } + + /// Encodes `sentences` into one flat, row-major `sentences.len() * cols` + /// buffer, pooling each sentence directly into the output — a single + /// allocation for the whole batch, with no per-sentence `Vec`. + pub fn encode_flat( &self, sentences: &[String], max_length: Option, batch_size: usize, - ) -> Result>> { - let mut embeddings = Vec::with_capacity(sentences.len()); + ) -> Result> { + let mut out = Vec::with_capacity(sentences.len() * self.cols); for batch in sentences.chunks(batch_size) { let truncated: Vec<&str> = batch.iter().map(|text| { max_length.map(|max_tok| Self::truncate_str(text, max_tok, self.median_token_length)).unwrap_or(text.as_str()) @@ -335,70 +355,55 @@ impl StaticModel { if let Some(max_tok) = max_length { token_ids.truncate(max_tok); } - embeddings.push(self.pool_ids(token_ids)); + self.pool_into(&token_ids, &mut out)?; } } - Ok(embeddings) - } - - #[allow(dead_code)] - pub fn encode(&self, sentences: &[String]) -> Result>> { - self.encode_with_args(sentences, Some(512), 1024) + Ok(out) } - #[allow(dead_code)] - pub fn encode_single(&self, sentence: &str) -> Vec { - self.encode(&[sentence.to_string()]) - .ok() - .and_then(|v| v.into_iter().next()) - .unwrap_or_default() - } - - fn pool_ids(&self, ids: Vec) -> Vec { - let dim = self.embeddings.ncols(); + /// Mean-pools `ids` and appends the resulting `cols` floats to `out`. Empty + /// input contributes a zero vector. + fn pool_into(&self, ids: &[u32], out: &mut Vec) -> Result<()> { + let dim = self.cols; + let start = out.len(); + out.resize(start + dim, 0.0); if ids.is_empty() { - return vec![0.0_f32; dim]; + return Ok(()); } + let sum = &mut out[start..start + dim]; - let mut sum = vec![0.0_f32; dim]; - match (&self.weights, &self.token_mapping) { (Some(weights), Some(mapping)) => { - for &id in &ids { + for &id in ids { let tok = id as usize; let row_idx = mapping.get(tok).copied().unwrap_or(tok); let scale = weights.get(tok).copied().unwrap_or(1.0); - let row = self.embeddings.row(row_idx); - for (s, &v) in sum.iter_mut().zip(row.iter()) { + for (s, &v) in sum.iter_mut().zip(self.row(row_idx)?) { *s += v * scale; } } } (Some(weights), None) => { - for &id in &ids { + for &id in ids { let tok = id as usize; let scale = weights.get(tok).copied().unwrap_or(1.0); - let row = self.embeddings.row(tok); - for (s, &v) in sum.iter_mut().zip(row.iter()) { + for (s, &v) in sum.iter_mut().zip(self.row(tok)?) { *s += v * scale; } } } (None, Some(mapping)) => { - for &id in &ids { + for &id in ids { let tok = id as usize; let row_idx = mapping.get(tok).copied().unwrap_or(tok); - let row = self.embeddings.row(row_idx); - for (s, &v) in sum.iter_mut().zip(row.iter()) { + for (s, &v) in sum.iter_mut().zip(self.row(row_idx)?) { *s += v; } } } (None, None) => { - for &id in &ids { - let tok = id as usize; - let row = self.embeddings.row(tok); - for (s, &v) in sum.iter_mut().zip(row.iter()) { + for &id in ids { + for (s, &v) in sum.iter_mut().zip(self.row(id as usize)?) { *s += v; } } @@ -406,24 +411,23 @@ impl StaticModel { } let denom = ids.len() as f32; - for x in &mut sum { + for x in sum.iter_mut() { *x /= denom; } - if self.normalize { let norm_sq: f32 = sum.iter().map(|&v| v * v).sum(); if norm_sq > 1e-12 { let norm = norm_sq.sqrt(); - for x in &mut sum { + for x in sum.iter_mut() { *x /= norm; } } } - sum + Ok(()) } - pub fn dim(&self) -> usize { self.embeddings.ncols() } - pub fn vocabulary_size(&self) -> usize { self.embeddings.nrows() } + pub fn dim(&self) -> usize { self.cols } + pub fn vocabulary_size(&self) -> usize { self.rows } pub fn is_normalized(&self) -> bool { self.normalize } pub fn median_token_length(&self) -> usize { self.median_token_length } } From 747d215bf1dde21bf857c8b1d7f479d5155b419d Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Tue, 7 Jul 2026 15:16:38 +0300 Subject: [PATCH 17/20] update readme --- README.md | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index db62d40..a6eebf3 100644 --- a/README.md +++ b/README.md @@ -216,17 +216,28 @@ embedding), `chunkText` (overlapping chunker), `ModelInfo`, `LoadProgress` / ## Performance Single-vector math runs natively in Dart with zero FFI overhead; batch generation -uses the Rust engine's SIMD. Benchmarked on a typical machine (AOT compiled): +uses the Rust engine's SIMD. Measured on an Apple Silicon laptop from an AOT +`dart build` bundle, best of 3 runs (absolute numbers vary by machine): | Model | Load (cached) | Single embedding | Batch of 32 | | ----- | ------------- | ---------------- | ----------- | -| `potion-base-2M` | ~40 ms | 373 μs | 3.85 ms | -| `potion-base-4M` | ~40 ms | 364 μs | 4.19 ms | -| `potion-base-8M` | ~40 ms | 382 μs | 5.60 ms | -| `potion-base-32M` | ~120 ms | 453 μs | 6.79 ms | -| `potion-multilingual-128M` | ~1050 ms | 416 μs | 5.38 ms | +| `potion-base-2M` | 21 ms | 240 μs | 3.8 ms | +| `potion-base-4M` | 22 ms | 248 μs | 3.8 ms | +| `potion-base-8M` | 24 ms | 251 μs | 4.0 ms | +| `potion-base-32M` | 66 ms | 254 μs | 4.3 ms | +| `potion-multilingual-128M` | 841 ms | 312 μs | 4.1 ms | -`similaritySearchWithScores` over 100,000 vectors takes **< 100 ms** in pure Dart. +A single embedding is a few hundred microseconds; `similaritySearchWithScores` +over 100,000 vectors takes **< 100 ms** in pure Dart. + +## Deployment + +The native library ships with your app automatically — nothing to link by hand: + +- **Flutter** (`flutter build …`) and **`dart run`** bundle and resolve it for you. +- **Standalone CLI**: build with `dart build cli` (not `dart compile exe`, which + does not bundle native assets yet). The library is copied into `bundle/lib/` + next to the executable, so the bundle is self-contained. ## Development From a59a21f7403ccdf35f29adbd609f704759def875 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Tue, 7 Jul 2026 15:55:00 +0300 Subject: [PATCH 18/20] cleanup before release --- CHANGELOG.md | 4 ++++ LICENSE | 2 +- README.md | 16 +++++++++++----- lib/src/embedding_index.dart | 14 ++++++++++++++ lib/src/model2vec_base.dart | 6 ++++++ lib/src/utils.dart | 18 ++++++++++++------ pubspec.yaml | 10 +++++----- test/embedding_index_test.dart | 13 +++++++++++++ test/utils_test.dart | 13 +++++++++++++ 9 files changed, 79 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 86ca7f6..d38d29e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,10 @@ correctness. **This release is breaking** — see migration below. `initEmbedderAdvanced`, `initEmbedderFromBytes` and their async forms are removed. `Model2VecUtils.similaritySearch` /`similaritySearchWithThreshold` are removed in favour of `similaritySearchWithScores` (read `.index`). +- **Batch signature.** `generateBatchEmbeddings` no longer takes `batchSize` + (its signature is now `(List texts, {int maxLength})`). The native + layer batches internally; `batchSize` remains only on + `generateEmbeddingStream`, which still controls its per-batch size. **Improvements:** diff --git a/LICENSE b/LICENSE index d88ce2b..b26b439 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2026 Andrii +Copyright (c) 2026 pro100Andrey Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/README.md b/README.md index a6eebf3..2dbd50b 100644 --- a/README.md +++ b/README.md @@ -4,11 +4,11 @@ [![CI](https://github.com/pro100andrey/model2vec/actions/workflows/ci.yml/badge.svg)](https://github.com/pro100andrey/model2vec/actions/workflows/ci.yml) [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT) -**Fast, on-device text embeddings for Dart & Flutter.** A wrapper around -[model2vec-rs](https://github.com/MinishLab/model2vec-rs) via Rust FFI and Native -Assets. Model2Vec turns text into vectors with a static vocabulary lookup — not a -transformer — so embeddings are generated in **microseconds**, with no server, no -Python, and no network after the model is cached. +**Fast, on-device text embeddings for Dart & Flutter.** A Dart implementation of +[Model2Vec](https://github.com/MinishLab/model2vec) with a self-contained Rust +core (FFI + Native Assets). It turns text into vectors with a static vocabulary +lookup — not a transformer — so embeddings are generated in **microseconds**, +with no server, no Python, and no network after the model is cached. ## Features @@ -267,6 +267,12 @@ dart analyze --fatal-infos dart test ``` +## Credits + +Model2Vec and the Potion models are by [MinishLab](https://github.com/MinishLab). +This package's Rust core draws on their +[model2vec-rs](https://github.com/MinishLab/model2vec-rs). + ## License MIT — see [LICENSE](LICENSE). diff --git a/lib/src/embedding_index.dart b/lib/src/embedding_index.dart index 7fd3658..207e0fe 100644 --- a/lib/src/embedding_index.dart +++ b/lib/src/embedding_index.dart @@ -290,6 +290,20 @@ class EmbeddingIndex { final count = data.getUint32(o, Endian.little); o += 4; + // Reject a corrupt header that declares sizes far larger than the blob + // before we allocate a giant list (which would OOM past the RangeError + // guard below). Each entry carries at least a 4-byte id length plus `dim` + // vector elements, so these bounds hold for any valid blob. + final elemSize = quantized ? 1 : 4; + final remaining = bytes.length - o; + if (dim * elemSize > remaining || + (count > 0 && count > remaining ~/ (4 + dim * elemSize))) { + throw ArgumentError( + 'corrupt EmbeddingIndex blob: dim=$dim count=$count exceed ' + '$remaining remaining bytes', + ); + } + final index = EmbeddingIndex(quantized: quantized); for (var e = 0; e < count; e++) { final idLen = data.getUint32(o, Endian.little); diff --git a/lib/src/model2vec_base.dart b/lib/src/model2vec_base.dart index 27e20f6..e0957be 100644 --- a/lib/src/model2vec_base.dart +++ b/lib/src/model2vec_base.dart @@ -217,6 +217,12 @@ abstract final class Model2Vec { /// Only one load may run at a time (the native model is a single process /// global); do not start another load, or switch/unload the model, while this /// stream is active. + /// + /// Cancelling the stream (e.g. `break`ing out of `await for`, or + /// `subscription.cancel()`) stops progress events but does **not** cancel the + /// load: the background isolate keeps downloading and will still swap in the + /// new process-global model when it finishes. There is no way to abort a load + /// in flight — only [unloadModel] afterwards. static Stream loadModelWithProgress( String modelPath, { String? hfToken, diff --git a/lib/src/utils.dart b/lib/src/utils.dart index f76b60f..bd53a4a 100644 --- a/lib/src/utils.dart +++ b/lib/src/utils.dart @@ -72,11 +72,11 @@ final class Model2VecUtils { /// Ranks [candidates] by cosine similarity to [query], most similar first, /// pairing each kept candidate's index with its score. /// - /// Keeps at most [topK] results; pass a [threshold] to also drop any result - /// scoring below it (raise [topK] to keep every match above the floor). This - /// is the ad-hoc counterpart to `EmbeddingIndex.search` for a list of vectors - /// you already hold in memory — it is the single scored search on this class; - /// the older index-only variants are deprecated. + /// Keeps at most [topK] results **even when a [threshold] is given** — the + /// threshold only drops matches scoring below it, it does not lift the [topK] + /// cap (which defaults to 5). To get *every* match above the floor, pass + /// `topK: candidates.length`. This is the ad-hoc counterpart to + /// `EmbeddingIndex.search` for a list of vectors you already hold in memory. static List<({int index, double score})> similaritySearchWithScores( Float32List query, List candidates, { @@ -205,11 +205,17 @@ final class Model2VecUtils { /// Quantizes a normalized [Float32List] to an [Int8List]. /// This saves 4x memory with minimal loss of accuracy for search. + /// + /// Non-finite components are handled defensively rather than throwing: + /// `NaN` maps to `0`, `+Infinity` to `127`, `-Infinity` to `-128`. static Int8List quantizeToInt8(Float32List vector) { final result = Int8List(vector.length); for (var i = 0; i < vector.length; i++) { - result[i] = (vector[i] * 127.0).round().clamp(-128, 127); + final v = vector[i]; + // Map NaN to 0 (neutral) up front — round() throws on NaN. clamp() then + // bounds ±Infinity to the int8 limits (+Inf -> 127, -Inf -> -128). + result[i] = v.isNaN ? 0 : (v * 127.0).clamp(-128, 127).round(); } return result; diff --git a/pubspec.yaml b/pubspec.yaml index b078b7a..83ae19c 100644 --- a/pubspec.yaml +++ b/pubspec.yaml @@ -1,5 +1,5 @@ name: model2vec -description: A high-performance Dart wrapper for model2vec-rs using Rust FFI. Generate fast, local, and static text embeddings with minimal memory footprint using Native Assets. +description: On-device Model2Vec text embeddings for Dart & Flutter — a self-contained Rust core via FFI and Native Assets. Fast, local, static, minimal memory. version: 2.0.0 repository: https://github.com/pro100andrey/model2vec homepage: https://github.com/pro100andrey/model2vec @@ -16,16 +16,16 @@ environment: sdk: "^3.10.0" dependencies: - code_assets: ^1.1.0 + code_assets: ^1.2.1 ffi: ^2.2.0 - hooks: ^2.0.0 + hooks: ^2.0.2 path: ^1.9.1 dev_dependencies: benchmark_harness: ^2.4.0 ffigen: ^20.1.1 - pro_lints: ^6.0.0 - test: ^1.31.1 + pro_lints: ^6.1.0 + test: ^1.31.2 ffigen: name: Model2VecBindings diff --git a/test/embedding_index_test.dart b/test/embedding_index_test.dart index 7c936b1..e4973fc 100644 --- a/test/embedding_index_test.dart +++ b/test/embedding_index_test.dart @@ -176,6 +176,19 @@ void main() { ); }); + test('fromBytes rejects a header dim larger than the blob (no OOM)', () { + // Valid magic + version, but a hostile dim that would allocate ~17 GB + // before any read fails — must be rejected up front as ArgumentError. + final blob = Uint8List.fromList([ + 0x4D, 0x32, 0x56, 0x49, // magic 'M2VI' + 0x02, // version 2 + 0x00, // flags: not quantized + 0xFF, 0xFF, 0xFF, 0xFF, // dim = 0xFFFFFFFF + 0x01, 0x00, 0x00, 0x00, // count = 1 + ]); + expect(() => EmbeddingIndex.fromBytes(blob), throwsArgumentError); + }); + test('removing the last entry resets the dimension to null', () { final index = EmbeddingIndex()..add('a', _v([1, 2, 3])); expect(index.dimension, 3); diff --git a/test/utils_test.dart b/test/utils_test.dart index d6bf96f..958565d 100644 --- a/test/utils_test.dart +++ b/test/utils_test.dart @@ -154,6 +154,19 @@ void main() { // clamp -128; 2.0*127=254 -> clamp 127. expect(quantized, [127, -128, 127]); }); + + test('maps non-finite input to bounds instead of throwing', () { + final q = Model2VecUtils.quantizeToInt8( + Float32List.fromList([ + double.nan, + double.infinity, + double.negativeInfinity, + 0.5, + ]), + ); + // NaN -> 0, +Inf -> 127, -Inf -> -128, 0.5*127=63.5 -> 64. + expect(q, [0, 127, -128, 64]); + }); }); group('dequantizeInt8', () { From 6e81af1b6de747c6de0d7cf12582defd9aa4b178 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Tue, 7 Jul 2026 15:59:03 +0300 Subject: [PATCH 19/20] fix ci.yaml --- .github/workflows/ci.yml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9b4b0b1..fc1d408 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -43,8 +43,9 @@ jobs: - uses: dart-lang/setup-dart@v1 with: sdk: stable - - name: Install Rust ${{ env.RUSTUP_TOOLCHAIN }} (host target only) - run: rustup toolchain install "$RUSTUP_TOOLCHAIN" --profile minimal + - name: Install Rust ${{ env.RUSTUP_TOOLCHAIN }} (host target + clippy) + # --profile minimal omits clippy, so request the component explicitly. + run: rustup toolchain install "$RUSTUP_TOOLCHAIN" --profile minimal --component clippy - name: Cache cargo registry + build uses: Swatinem/rust-cache@v2 with: From 80b4e757c81ad7b6026374d90e4e35a4a9c4d232 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Tue, 7 Jul 2026 16:10:34 +0300 Subject: [PATCH 20/20] code style --- lib/src/channel.dart | 1 + lib/src/chunker.dart | 3 +++ lib/src/embedding_index.dart | 35 +++++++++++++++++++++++----------- lib/src/embedding_pool.dart | 14 +++++++------- lib/src/embedding_worker.dart | 3 ++- lib/src/exception.dart | 18 ++++++++--------- lib/src/utils.dart | 19 ++++++++++++++++-- lib/src/worker_protocol.dart | 16 ++++++++++------ test/exception_test.dart | 18 ++++++++--------- test/worker_protocol_test.dart | 8 ++------ test/worker_support.dart | 4 ++-- 11 files changed, 86 insertions(+), 53 deletions(-) diff --git a/lib/src/channel.dart b/lib/src/channel.dart index c243609..c392a14 100644 --- a/lib/src/channel.dart +++ b/lib/src/channel.dart @@ -45,6 +45,7 @@ class IsolateChannel implements Channel { final StreamSubscription _subscription; final StreamSubscription _exitSubscription; final StreamController _incoming; + var _closed = false; /// Spawns a worker running [entryPoint] and completes the handshake. diff --git a/lib/src/chunker.dart b/lib/src/chunker.dart index 5368680..22086f5 100644 --- a/lib/src/chunker.dart +++ b/lib/src/chunker.dart @@ -63,11 +63,14 @@ String _overlapTail(String chunk, int overlap) { if (overlap <= 0 || chunk.length <= overlap) { return ''; } + var start = chunk.length - overlap; final space = chunk.indexOf(' ', start); if (space == -1) { return ''; // the tail would be a single partial word; skip it } + start = space + 1; + return chunk.substring(start); } diff --git a/lib/src/embedding_index.dart b/lib/src/embedding_index.dart index 207e0fe..d1c91a2 100644 --- a/lib/src/embedding_index.dart +++ b/lib/src/embedding_index.dart @@ -96,6 +96,7 @@ class EmbeddingIndex { if (_store.isEmpty) { _dimension = null; } + return removed; } @@ -114,9 +115,11 @@ class EmbeddingIndex { if (topK <= 0) { return []; } + if (topK >= scored.length) { return scored; } + return scored.sublist(0, topK); } @@ -135,11 +138,13 @@ class EmbeddingIndex { if (_store.isEmpty) { return []; } + if (query.length != _dimension) { throw ArgumentError( 'query length ${query.length} != index dimension $_dimension', ); } + final results = []; for (final entry in _store.entries) { final vector = _asFloat32(entry.value.vector); @@ -151,6 +156,7 @@ class EmbeddingIndex { ), ); } + return results; } @@ -204,15 +210,15 @@ class EmbeddingIndex { o += 1; data.setUint8(o, _quantized ? 1 : 0); o += 1; - data.setUint32(o, dim, Endian.little); + data.setUint32(o, dim, .little); o += 4; - data.setUint32(o, _store.length, Endian.little); + data.setUint32(o, _store.length, .little); o += 4; var e = 0; for (final stored in _store.values) { final idBytes = idByteList[e]; - data.setUint32(o, idBytes.length, Endian.little); + data.setUint32(o, idBytes.length, .little); o += 4; bytes.setRange(o, o + idBytes.length, idBytes); o += idBytes.length; @@ -225,16 +231,16 @@ class EmbeddingIndex { } else { final f = stored.vector as Float32List; for (var i = 0; i < dim; i++) { - data.setFloat32(o, f[i], Endian.little); + data.setFloat32(o, f[i], .little); o += 4; } } final payloadBytes = payloadByteList[e++]; if (payloadBytes == null) { - data.setUint32(o, _noPayload, Endian.little); + data.setUint32(o, _noPayload, .little); o += 4; } else { - data.setUint32(o, payloadBytes.length, Endian.little); + data.setUint32(o, payloadBytes.length, .little); o += 4; bytes.setRange(o, o + payloadBytes.length, payloadBytes); o += payloadBytes.length; @@ -250,6 +256,7 @@ class EmbeddingIndex { if (bytes.length < 14 || !_hasMagic(bytes)) { throw ArgumentError('not a valid EmbeddingIndex blob'); } + try { return _decode(bytes); } on FormatException catch (e) { @@ -269,6 +276,7 @@ class EmbeddingIndex { return false; } } + return true; } @@ -282,12 +290,13 @@ class EmbeddingIndex { if (version != 1 && version != 2) { throw ArgumentError('unsupported EmbeddingIndex version $version'); } + final hasPayloads = version >= 2; final quantized = data.getUint8(o) == 1; o += 1; - final dim = data.getUint32(o, Endian.little); + final dim = data.getUint32(o, .little); o += 4; - final count = data.getUint32(o, Endian.little); + final count = data.getUint32(o, .little); o += 4; // Reject a corrupt header that declares sizes far larger than the blob @@ -306,7 +315,7 @@ class EmbeddingIndex { final index = EmbeddingIndex(quantized: quantized); for (var e = 0; e < count; e++) { - final idLen = data.getUint32(o, Endian.little); + final idLen = data.getUint32(o, .little); o += 4; final id = utf8.decode(bytes.sublist(o, o + idLen)); o += idLen; @@ -321,23 +330,27 @@ class EmbeddingIndex { } else { final f = Float32List(dim); for (var i = 0; i < dim; i++) { - f[i] = data.getFloat32(o, Endian.little); + f[i] = data.getFloat32(o, .little); o += 4; } vector = f; } + String? payload; if (hasPayloads) { - final payloadLen = data.getUint32(o, Endian.little); + final payloadLen = data.getUint32(o, .little); o += 4; if (payloadLen != _noPayload) { payload = utf8.decode(bytes.sublist(o, o + payloadLen)); o += payloadLen; } } + index._store[id] = (vector: vector, payload: payload); } + index._dimension = count > 0 ? dim : null; + return index; } } diff --git a/lib/src/embedding_pool.dart b/lib/src/embedding_pool.dart index 623b097..8ccd4c4 100644 --- a/lib/src/embedding_pool.dart +++ b/lib/src/embedding_pool.dart @@ -16,8 +16,7 @@ import 'embedding_worker.dart'; /// must not be switched while the pool is working (the one-model-per-run /// contract, same as the single worker). class EmbeddingPool { - EmbeddingPool._(this._workers) - : _inFlight = List.filled(_workers.length, 0); + EmbeddingPool._(this._workers) : _inFlight = .filled(_workers.length, 0); final List _workers; final List _inFlight; @@ -37,7 +36,7 @@ class EmbeddingPool { // Start all in parallel; if any worker fails to start, close the ones that // did so we never leak orphaned isolates. final outcomes = await Future.wait( - List.generate(count, (_) async { + .generate(count, (_) async { try { return await EmbeddingWorker.start(entryPoint: entryPoint); } on Object { @@ -50,6 +49,7 @@ class EmbeddingPool { await Future.wait(workers.map((worker) => worker.close())); throw StateError('failed to start $count embedding workers'); } + return EmbeddingPool._(workers); } @@ -61,6 +61,7 @@ class EmbeddingPool { }) { final worker = _leastBusyIndex(); _inFlight[worker]++; + return _workers[worker] .embedBatch(batch, maxLength: maxLength) .whenComplete(() => _inFlight[worker]--); @@ -74,14 +75,12 @@ class EmbeddingPool { Future>> embedBatches( List> batches, { int maxLength = 512, - }) => Future.wait( + }) => .wait( batches.map((batch) => embedBatch(batch, maxLength: maxLength)), ); /// Tears down every worker in the pool. - Future close() async { - await Future.wait(_workers.map((worker) => worker.close())); - } + Future close() => .wait(_workers.map((worker) => worker.close())); int _leastBusyIndex() { var best = 0; @@ -90,6 +89,7 @@ class EmbeddingPool { best = i; } } + return best; } } diff --git a/lib/src/embedding_worker.dart b/lib/src/embedding_worker.dart index 5e54ef9..aead422 100644 --- a/lib/src/embedding_worker.dart +++ b/lib/src/embedding_worker.dart @@ -28,6 +28,7 @@ class EmbeddingWorker { void Function(SendPort) entryPoint = embeddingWorkerEntryPoint, }) async { final channel = await IsolateChannel.start(entryPoint); + return EmbeddingWorker._(WorkerProtocol(channel)); } @@ -66,7 +67,7 @@ void embeddingWorkerEntryPoint(SendPort mainSendPort) { } on Model2VecException catch (e) { mainSendPort.send(e); } on Object catch (e) { - mainSendPort.send(Model2VecException(Model2VecErrorKind.unknown, '$e')); + mainSendPort.send(Model2VecException(.unknown, '$e')); } }); } diff --git a/lib/src/exception.dart b/lib/src/exception.dart index 9941370..b7d4538 100644 --- a/lib/src/exception.dart +++ b/lib/src/exception.dart @@ -60,13 +60,13 @@ class Model2VecException implements Exception { /// /// Keep in sync with the `CODE_*` constants in `native/src/lib.rs`. Model2VecErrorKind _kindFromCode(int code) => switch (code) { - 1 => Model2VecErrorKind.notInitialized, - 2 => Model2VecErrorKind.modelLoadFailed, - 3 => Model2VecErrorKind.initFromBytesFailed, - 4 => Model2VecErrorKind.lockPoisoned, - 5 => Model2VecErrorKind.nullArgument, - 6 => Model2VecErrorKind.tokenizationFailed, - 7 => Model2VecErrorKind.emptyResult, - 8 => Model2VecErrorKind.panic, - _ => Model2VecErrorKind.unknown, + 1 => .notInitialized, + 2 => .modelLoadFailed, + 3 => .initFromBytesFailed, + 4 => .lockPoisoned, + 5 => .nullArgument, + 6 => .tokenizationFailed, + 7 => .emptyResult, + 8 => .panic, + _ => .unknown, }; diff --git a/lib/src/utils.dart b/lib/src/utils.dart index bd53a4a..7efffbb 100644 --- a/lib/src/utils.dart +++ b/lib/src/utils.dart @@ -52,6 +52,7 @@ final class Model2VecUtils { for (var i = 0; i < a.length; i++) { result += a[i] * b[i]; } + return result; } @@ -66,6 +67,7 @@ final class Model2VecUtils { final diff = a[i] - b[i]; sum += diff * diff; } + return math.sqrt(sum); } @@ -86,6 +88,7 @@ final class Model2VecUtils { if (candidates.isEmpty || topK <= 0) { return []; } + final scored = <({int index, double score})>[ for (var i = 0; i < candidates.length; i++) (index: i, score: cosineSimilarity(query, candidates[i])), @@ -95,6 +98,7 @@ final class Model2VecUtils { ? scored : scored.where((r) => r.score >= threshold).toList()) ..sort((a, b) => b.score.compareTo(a.score)); + return ranked.take(math.min(topK, ranked.length)).toList(growable: false); } @@ -113,9 +117,11 @@ final class Model2VecUtils { if (lambda < 0.0 || lambda > 1.0) { throw ArgumentError.value(lambda, 'lambda', 'must be in [0.0, 1.0]'); } + if (candidates.isEmpty) { return []; } + final k = math.min(topK, candidates.length); final relevance = [for (final c in candidates) cosineSimilarity(query, c)]; // Running max similarity of each candidate to any already-selected item, @@ -138,8 +144,10 @@ final class Model2VecUtils { bestIdx = i; } } + selected.add(bestIdx); remaining.remove(bestIdx); + for (final i in remaining) { final sim = cosineSimilarity(candidates[i], candidates[bestIdx]); if (sim > maxSimToSelected[i]) { @@ -147,6 +155,7 @@ final class Model2VecUtils { } } } + return selected; } @@ -166,14 +175,17 @@ final class Model2VecUtils { for (var i = 0; i < vector.length; i++) { normSq += vector[i] * vector[i]; } + if (normSq == 0.0) { - return Float32List.fromList(vector); + return .fromList(vector); } + final norm = math.sqrt(normSq); final result = Float32List(vector.length); for (var i = 0; i < vector.length; i++) { result[i] = vector[i] / norm; } + return result; } @@ -183,6 +195,7 @@ final class Model2VecUtils { if (vectors.isEmpty) { throw ArgumentError('Cannot pool an empty list of vectors'); } + final dim = vectors.first.length; final result = Float32List(dim); @@ -190,6 +203,7 @@ final class Model2VecUtils { if (v.length != dim) { throw ArgumentError('All vectors must have the same dimension'); } + for (var i = 0; i < dim; i++) { result[i] += v[i]; } @@ -229,6 +243,7 @@ final class Model2VecUtils { for (var i = 0; i < quantized.length; i++) { result[i] = quantized[i] / 127.0; } + return result; } @@ -241,7 +256,7 @@ final class Model2VecUtils { final bytes = base64Decode(base64String); final buffer = bytes.buffer; - return Float32List.view( + return .view( buffer, bytes.offsetInBytes, bytes.lengthInBytes ~/ Float32List.bytesPerElement, diff --git a/lib/src/worker_protocol.dart b/lib/src/worker_protocol.dart index 4f0d355..9977f74 100644 --- a/lib/src/worker_protocol.dart +++ b/lib/src/worker_protocol.dart @@ -38,6 +38,7 @@ class WorkerProtocol { final Channel _channel; late final StreamSubscription _subscription; final _queue = Queue<_Request>(); + var _inFlight = false; var _acceptingRequests = true; var _disposed = false; @@ -49,16 +50,15 @@ class WorkerProtocol { required int maxLength, }) { if (!_acceptingRequests) { - return Future>.error( - const Model2VecException( - Model2VecErrorKind.unknown, - 'embedding worker is closed', - ), + return .error( + const Model2VecException(.unknown, 'embedding worker is closed'), ); } + final request = _Request(batch, maxLength); _queue.add(request); _pump(); + return request.completer.future; } @@ -66,6 +66,7 @@ class WorkerProtocol { if (_inFlight || _queue.isEmpty) { return; } + _inFlight = true; final head = _queue.first; _channel.send((batch: head.batch, maxLength: head.maxLength)); @@ -75,6 +76,7 @@ class WorkerProtocol { if (!_inFlight || _queue.isEmpty) { return; // stray reply; nothing is awaiting it } + final request = _queue.removeFirst(); _inFlight = false; @@ -97,6 +99,7 @@ class WorkerProtocol { if (!_acceptingRequests) { return; } + _acceptingRequests = false; _failAll('embedding worker terminated unexpectedly'); } @@ -107,6 +110,7 @@ class WorkerProtocol { if (_disposed) { return; } + _disposed = true; _acceptingRequests = false; await _subscription.cancel(); @@ -115,7 +119,7 @@ class WorkerProtocol { } void _failAll(String message) { - final error = Model2VecException(Model2VecErrorKind.unknown, message); + final error = Model2VecException(.unknown, message); while (_queue.isNotEmpty) { _queue.removeFirst().completer.completeError(error); } diff --git a/test/exception_test.dart b/test/exception_test.dart index b7f4bb7..9144eff 100644 --- a/test/exception_test.dart +++ b/test/exception_test.dart @@ -4,15 +4,15 @@ import 'package:test/test.dart'; void main() { group('Model2VecException.fromNative', () { // Mirrors the CODE_* constants in native/src/lib.rs. - const cases = { - 1: Model2VecErrorKind.notInitialized, - 2: Model2VecErrorKind.modelLoadFailed, - 3: Model2VecErrorKind.initFromBytesFailed, - 4: Model2VecErrorKind.lockPoisoned, - 5: Model2VecErrorKind.nullArgument, - 6: Model2VecErrorKind.tokenizationFailed, - 7: Model2VecErrorKind.emptyResult, - 8: Model2VecErrorKind.panic, + const cases = { + 1: .notInitialized, + 2: .modelLoadFailed, + 3: .initFromBytesFailed, + 4: .lockPoisoned, + 5: .nullArgument, + 6: .tokenizationFailed, + 7: .emptyResult, + 8: .panic, }; for (final entry in cases.entries) { diff --git a/test/worker_protocol_test.dart b/test/worker_protocol_test.dart index bb35834..11511a1 100644 --- a/test/worker_protocol_test.dart +++ b/test/worker_protocol_test.dart @@ -29,7 +29,7 @@ class FakeChannel implements Channel { void deliver(Object? message) => _incoming.add(message); } -Float32List _vec(List xs) => Float32List.fromList(xs); +Float32List _vec(List xs) => .fromList(xs); void main() { group('WorkerProtocol', () { @@ -82,11 +82,7 @@ void main() { test('propagates a typed Model2VecException from the worker', () async { final future = protocol.embedBatch(['a'], maxLength: 8); channel.deliver( - const Model2VecException( - Model2VecErrorKind.notInitialized, - 'no model', - 1, - ), + const Model2VecException(.notInitialized, 'no model', 1), ); await expectLater( diff --git a/test/worker_support.dart b/test/worker_support.dart index 6e6f4dc..3242c51 100644 --- a/test/worker_support.dart +++ b/test/worker_support.dart @@ -44,7 +44,7 @@ void errorEntryPoint(SendPort mainSendPort) { } mainSendPort.send( const Model2VecException( - Model2VecErrorKind.tokenizationFailed, + .tokenizationFailed, 'boom', 6, ), @@ -120,7 +120,7 @@ void failThenEchoEntryPoint(SendPort mainSendPort) { if (!failed) { failed = true; mainSendPort.send( - const Model2VecException(Model2VecErrorKind.panic, 'boom'), + const Model2VecException(.panic, 'boom'), ); return; }