static_embeddings 0.1.5 → 1.5.7

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: fb1d3530fba4fad467251429797fdcaae527b5c6c18413e8847002e0c3f3b5e4
-  data.tar.gz: f02bb4c129449803a57b08bf9b7186af08a6d21edee1e68d18178b248b855e0b
+  metadata.gz: a8e174dd93ba3b3e9ee25047fd28ea5ba980d096a4922c0b8325c366208bff65
+  data.tar.gz: 741e0cc35791d9a335c82a05aed18e3388232828824168ff59ed1751c281e8a8
 SHA512:
-  metadata.gz: cec0a3959cb3f093c7bb2b06213b4c7e983986466a115eb79e0b045a5b52fdda9797be41365db41a74d85bb52c384c147d68d3d84c503fef10c122bba7b72333
-  data.tar.gz: e99039606af3de0106e1ac26efaa78a06f09acbd9c2faaae9b4ca17339bdeb4830ba922bef84dc141317898556d09b1ec41cf1ebfebcce5e141d877c3fc34598
+  metadata.gz: 58462ed4c203fb657630d092e1a4aa85857e5d44106bb97feae26f12579fc4b6d88d20632e462624ad0cb93e1d0d0e735f91ee12aaec80b637d29b313e6dc53d
+  data.tar.gz: 80e9897969dfe6aace3d851f1a11ffb6610b628cddc35239301a676b2ce32d97d1a8a614ae6dca0526703e61d6d813af89965d0906c593392e66cdd25f7e16da

data/CHANGELOG.md

@@ -1,5 +1,65 @@
 # Changelog
 
+## 1.5.7
+
+### Fixed
+
+- Model mappings are no longer closed while embedding, warmup, or token-ID
+  operations are using them without the GVL. `Model#close` raises
+  `StaticEmbeddings::Error` while such an operation is active.
+- Tokenizer code-point writes now verify the capacity invariant after growth.
+- Resolved Clang conversion warnings and formatting issues in the native code.
+
+## 1.5.6
+
+Multi-source WordPiece release. The native `.semb` v3 runtime is unchanged; new
+source families are normalized offline into the same audited
+`BERT_WORDPIECE_V1 -> lookup -> mean -> NONE/L2` contract. Model identity is
+chosen at `load_model`, not per `embed` call.
+
+### Added
+
+- **Model2Vec and Sentence Transformers importers.** A Model2Vec directory and a
+  Sentence Transformers `StaticEmbedding` directory (`modules.json` + exactly
+  one module path) compile to the same canonical data. The ST importer is
+  fail-closed and confines the module path to the source root with `realpath`.
+- **Source-faithful policies.** Model2Vec uses `UNK_DROP`, normalization from its
+  config and a 512-token default. Sentence Transformers `StaticEmbedding` uses
+  `UNK_INCLUDE`, `NORMALIZATION_NONE`, unlimited tokens and
+  `add_special_tokens=false`.
+- **`--dimensions N`.** Matryoshka prefix slicing happens while streaming the
+  embedding matrix; the runtime still loads an ordinary fixed-dimension `.semb`.
+- **`--max-tokens unlimited`.** Convert-time `0` remains an alias; runtime
+  `max_tokens: 0` is still rejected and `false` means unlimited.
+- **Russian retrieval fixture.** `tools/eval_retrieval.rb` and the checked-in
+  Russian FAQ corpus provide a small domain sanity check for the multilingual
+  WordPiece model; it is not presented as a general benchmark.
+
+### Changed
+
+- **Offline Ruby conversion was simplified.** Stateful `Converter` and the
+  behavior-heavy `CanonicalModel` were replaced by import functions, immutable
+  canonical data, pure provenance/meta transforms and a streaming format writer.
+  Writer/hash/trie/verifier code is split from runtime format constants, so plain
+  `require "static_embeddings"` does not load conversion code.
+- The Ruby `Reference` twin follows canonical UNK/normalization/max-token
+  policies instead of re-deriving Model2Vec defaults.
+- Model2Vec `normalize` now defaults to `false` when the key is absent, matching
+  the pinned upstream behavior. Unknown source families are rejected rather than
+  guessed from `tokenizer.json + model.safetensors`.
+- Header integer writes validate their unsigned range instead of silently
+  wrapping oversized Ruby integers.
+- `Model#provenance` is parsed once and cached.
+
+Local WordPiece conversions requiring **no C change**: `potion-base-8M`,
+`potion-science-32M`, `static-retrieval-mrl-en-v1` at 1024/512 and
+`static-similarity-mrl-multilingual-v1` at 512/256. The retrieval model has a
+recorded `SentenceTransformer.encode` oracle (436/436). Records are in
+`docs/MODEL_AUDIT.md`.
+
+SentencePiece/Unigram is intentionally **not** part of 1.5.6; it remains a
+future runtime capability rather than shipping an unaudited v4 path.
+
 ## 0.1.5
 
 Correctness/parity release. The native runtime architecture is unchanged, but

data/README.md

@@ -1,11 +1,11 @@
 # static_embeddings
 
-A Ruby runtime for converted Model2Vec / potion static embedding models. A model
+A Ruby runtime for converted Model2Vec and Sentence Transformers static embedding models. A model
 is converted once into a local `.semb` file and loaded through a small C
 extension. No ONNX Runtime, no Rust or Python at runtime, no network access, one
 mmap-able file, binary float32 output.
 
-The current production target is `minishlab/potion-retrieval-32M`.
+The reference production target remains `minishlab/potion-retrieval-32M`; 1.5.6 also supports audited WordPiece `StaticEmbedding` sources.
 
 ```ruby
 require "static_embeddings"
@@ -49,6 +49,21 @@ bundle exec ruby -Ilib exe/static_embeddings convert ./potion-retrieval-32M \
   --id potion-retrieval-32m
 ```
 
+Sentence Transformers static models (`modules.json` + one `StaticEmbedding`
+module) use the same command. Matryoshka prefix-slice is a convert-time flag,
+not a runtime `dim:`:
+
+```bash
+bundle exec ruby -Ilib exe/static_embeddings convert ./static-retrieval-mrl-en-v1 \
+  --id static-retrieval-mrl-en-v1-512 \
+  --dimensions 512 \
+  --trained-mrl-dims 1024,512,256,128,64,32
+```
+
+`--max-tokens unlimited` bakes no token cap (Sentence Transformers default).
+`--max-tokens 0` is a convert-time alias; `embed(max_tokens: 0)` is still
+rejected. Query and documents must use the same loaded `.semb`.
+
 The result lands in `~/.cache/static_embeddings/models/potion-retrieval-32m.semb`.
 Inspect or verify it with the `inspect` and `verify` subcommands, then load it:
 

data/docs/ARCHITECTURE.md

@@ -8,10 +8,10 @@ HuggingFace / Model2Vec files          offline, once, on your machine
   model.safetensors
         |
         v
-  Converter (pure Ruby, strict)   <--- all parsing, all validation, all
-        |                              Unicode table generation happens here
+  import -> canonical data -> writer   <--- pure Ruby, offline, strict
+        |                                    all source parsing/validation here
         v
-  model.semb                       <--- flat, versioned, mmap-able
+  model.semb                            <--- flat, versioned, mmap-able
         |
         v
   C runtime                        <--- mmap + bounds checks + tokenize +
@@ -24,6 +24,15 @@ Everything expensive, fragile or security-sensitive about reading third-party
 model files happens once, offline, in a language where it is easy to get
 right. What remains in C is a bounds-checked mmap and three loops.
 
+Only the offline import layer knows HuggingFace layouts. Model2Vec and Sentence
+Transformers `StaticEmbedding` sources are reduced to immutable canonical data:
+tokens, streamed matrix, tokenizer metadata, runtime policies and provenance.
+Pure transforms then build the v3 header/provenance and the existing WordPiece
+writer emits the artifact. A new WordPiece + mean-pooling source therefore does
+not require a C change. New tokenizer families remain explicit future
+capabilities and require a separately audited format/runtime change rather than
+source-name special cases.
+
 The C side is `se_format.c` (mmap and validation), `se_tokenizer.c`,
 `se_embed.c`, `se_unicode.c`, `se_f16.c` (half-precision codec and kernels),
 `se_topk.c`, `se_alloc_stats.c` (optional allocation counters), and
@@ -336,8 +345,8 @@ chunks -> .semb vectors --/
 1. **Which potion models actually match `BERT_WORDPIECE_V1`.** The runtime now
    implements the exact `normalized: false` extraction semantics needed by the
    five standard BERT added tokens, but still rejects arbitrary AddedVocabulary,
-   BPE and Unigram profiles. Each source model still needs an audit before it is
-   treated as compatible.
+   BPE and Unigram profiles. Type A (WordPiece + mean) models convert without a
+   C change; each source still needs an audit before it is treated as compatible.
 2. **Russian.** Distilling a multilingual teacher into a WordPiece vocabulary
    we control keeps the pure-C path, but skips the Tokenlearn pre-training
    that gives the published potion models their quality. That gap has to be

data/docs/MODEL_AUDIT.md

@@ -119,3 +119,113 @@ synthetic `tiny-wordpiece` fixture: `434/434`, all four failure lists empty,
 `min_cosine=0.9999999999999988`, `max_abs_all=5.960464477539063e-08`. That job
 proves the fixture loads in `StaticModel.from_pretrained` and that the checker
 contracts hold; it is not a potion audit.
+
+## 1.5.6 Patch 1 proof
+
+Architecture: Model2Vec / Sentence Transformers import functions → immutable canonical data →
+existing WordPiece `.semb` v3 writer. C runtime unchanged.
+
+### potion-retrieval-32M reconvert
+
+Official snapshot `6fc8051fab2a1e0ee76689cf08c853792ac285e7` ships a Sentence
+Transformers `modules.json` (`StaticEmbedding` + `Normalize`). Detection uses
+`config.json` `model_type=model2vec` first, so this stays UNK_DROP / L2 / 512.
+
+New file `potion-retrieval-32m-reconvert.semb` vs existing
+`potion-retrieval-32m-v3.semb`: vocab/hash/embeddings/norm_tables/tries
+byte-identical. 10-text corpus `max_abs=0`. Provenance JSON differs (new keys).
+
+### static-retrieval-mrl-en-v1 vs SentenceTransformer.encode
+
+Oracle: `tools/st_oracle.py` + `tools/check_st_parity.rb`,
+`sentence-transformers 6.0.1`, `add_special_tokens=false`,
+`normalize_embeddings=false`, 436 rows including `all-unk:private-use` and
+`long:known-words-600`.
+
+```text
+1024: 436/436  min_cosine=0.999999999999999  max_abs_all=7.62939453125e-06
+512:  436/436  min_cosine=0.9999999999999989 max_abs_all=7.62939453125e-06
+```
+
+512 compared as prefix of the 1024-d encode vectors. Ruby Reference matched C
+separately. Diagnostic `StaticEmbedding.forward` matched encode on this model.
+
+### UNK and max_tokens on the real models
+
+- potion-retrieval-32M `🧬 🧬 🧬` → UNK ids dropped → zero vector; 600×`hello`
+  → `truncated=true pooled=512`.
+- static-retrieval-mrl-en-v1-1024 same emoji string → three UNK ids pooled
+  (non-zero, mean of UNK rows); 600×`hello` → `truncated=false pooled=600`.
+- Private-use codepoints are BertNormalizer `Co` and become empty, not UNK.
+
+## 1.5.6 Type A conversions
+
+These records prove the offline import layer maps extra WordPiece sources onto
+the existing runtime. They are **not** a replacement for the
+potion-retrieval-32m Python oracle above. `static-retrieval-mrl-en-v1` has the
+separate `SentenceTransformer.encode` oracle recorded above; the multilingual
+similarity artifact was checked against the Ruby `.semb` Reference twin and the
+small Russian retrieval fixture, not claimed as a general upstream retrieval
+benchmark.
+
+All conversions used `static_embeddings 1.5.6` on arm64-darwin24.
+
+### minishlab/potion-base-8M
+
+- Snapshot: `bf8b056651a2c21b8d2565580b8569da283cab23`
+- Layout: Model2Vec (`UNK_DROP`, L2, `max_tokens` 512)
+- `.semb`: `potion-base-8m.semb` `dim=256` `vocab=29528` `33346200` bytes
+  SHA256 `8abd8d1f26511959e14ca26ec4dcacad1d3db9b4b66b2e67683fb2d06074f490`
+- Native C matched Ruby Reference on short English/OOV/empty texts (`max_abs < 1e-5`)
+- `model2vec 0.9.0` / `tokenizers 0.23.1` oracle: 434/434, vectors_checked=432,
+  intentional character-pretruncate deviations=2, min_cosine=0.9999999999999989,
+  max_abs_all=2.98e-08
+
+### minishlab/potion-science-32M
+
+- Snapshot: `7366079845507de14a4330007cdfa01bb92bca52`
+- Layout: Model2Vec (`UNK_DROP`, L2, `max_tokens` 512)
+- `.semb`: `potion-science-32m.semb` `dim=256` `vocab=124428` `140353240` bytes
+  SHA256 `e968f93b89d59d9c63c6a8152495283d1138144e61ed5c08d306c21c00b7cdb1`
+- Native C matched Ruby Reference on English and Russian snippets (`max_abs < 1e-5`)
+- `model2vec 0.9.0` / `tokenizers 0.23.1` oracle: 434/434, vectors_checked=432,
+  intentional character-pretruncate deviations=2, min_cosine=0.9999999999999986,
+  max_abs_all=2.98e-08
+
+### sentence-transformers/static-retrieval-mrl-en-v1
+
+- Snapshot: `f60985c706f192d45d218078e49e5a8b6f15283a`
+- Layout: Sentence Transformers StaticEmbedding (`UNK_INCLUDE`, no L2, unlimited)
+- `static-retrieval-mrl-en-v1-1024.semb` `dim=1024` `vocab=30522` `128186264` bytes
+  SHA256 `f316bb07348503418ed0f6d897ddbf7b44422d2077ca4860518f31f328610946`
+- `static-retrieval-mrl-en-v1-512.semb` `--dimensions 512` `65677208` bytes
+  SHA256 `efbb43c14def793d0bbb42ef591a0951caed23a49cdd94dcf72312569d1b7cff`
+- Prefix of a 1024-d vector matched the 512-d artifact on in-vocabulary English
+  (valid because this source does not L2-normalize)
+- Native C matched Ruby Reference on short texts
+
+### sentence-transformers/static-similarity-mrl-multilingual-v1
+
+- Snapshot: `b68f4122911bcffcd6e1f695f2d99cd6788972d8`
+- Layout: Sentence Transformers StaticEmbedding (`UNK_INCLUDE`, no L2, unlimited)
+- `static-similarity-mrl-multilingual-v1-512.semb` `--dimensions 512` `vocab=105879`
+  `227983824` bytes SHA256 `8d9a63d20ee23468ae2fd2584ef3f0ec75e720e68410e416157aaf2b873e110e`
+- `static-similarity-mrl-multilingual-v1-256.semb` `--dimensions 256` `119563728` bytes
+  SHA256 `2e9e689bd3afd18ed50a77f0f8842b4b5a5be39bfad75e3fa24ad99b478de068`
+- Native C matched Ruby Reference on Russian text (`max_abs 9.1e-7`)
+
+This model is published as similarity, not retrieval. `tools/eval_retrieval.rb`
+on `test/fixtures/russian_faq_eval.json` (10 queries / 10 docs, cosine@10,
+arm64-darwin24, 1.5.6):
+
+| model | dim | MRR | nDCG@10 |
+|---|---:|---:|---:|
+| static-similarity-mrl-multilingual-v1-256 | 256 | 0.875 | 0.906 |
+| static-similarity-mrl-multilingual-v1-512 | 512 | 0.850 | 0.889 |
+| static-retrieval-mrl-en-v1-512 | 512 | 0.792 | 0.842 |
+| potion-retrieval-32m | 512 | 0.733 | 0.799 |
+| potion-base-8m | 256 | 0.712 | 0.780 |
+
+Hit@10 was 1.0 for every model on this tiny labeled set. That is a domain sanity
+check, not a public retrieval benchmark, and it does not make the similarity
+model a hard-coded Russian default.

data/ext/static_embeddings/se_alloc_stats.c

@@ -5,9 +5,8 @@
 #include <string.h>
 
 static const char *const se_alloc_category_names[SE_ALLOC_CATEGORY_COUNT] = {
-    "unknown",    "scratch",     "batch_input",      "batch_output",
-    "batch_stats", "batch_index", "token_ids",       "topk_query",
-    "topk_best",  "topk_matrix_copy", "format_validate", "model_file"};
+    "unknown",   "scratch",    "batch_input", "batch_output",     "batch_stats",     "batch_index",
+    "token_ids", "topk_query", "topk_best",   "topk_matrix_copy", "format_validate", "model_file"};
 
 const char *se_alloc_category_name(se_alloc_category_t category) {
     if ((unsigned)category >= SE_ALLOC_CATEGORY_COUNT)
@@ -176,7 +175,8 @@ void *se_alloc_stats_realloc(se_alloc_category_t category, void *ptr, size_t byt
     if (old_header->h.magic != SE_ALLOC_STATS_MAGIC)
         abort();
 
-    se_alloc_category_t old_category = normalize_category((se_alloc_category_t)old_header->h.category);
+    se_alloc_category_t old_category =
+        normalize_category((se_alloc_category_t)old_header->h.category);
     size_t old_bytes = old_header->h.size;
     se_alloc_header_t *new_header = (se_alloc_header_t *)realloc(old_header, total);
     if (!new_header)
@@ -223,8 +223,7 @@ void se_alloc_stats_snapshot(se_alloc_stats_t out[SE_ALLOC_CATEGORY_COUNT]) {
     for (size_t i = 0; i < SE_ALLOC_CATEGORY_COUNT; i++) {
         out[i].current_bytes = atomic_load_size(&se_alloc_stats[i].current_bytes);
         out[i].peak_bytes = atomic_load_size(&se_alloc_stats[i].peak_bytes);
-        out[i].total_allocated_bytes =
-            atomic_load_size(&se_alloc_stats[i].total_allocated_bytes);
+        out[i].total_allocated_bytes = atomic_load_size(&se_alloc_stats[i].total_allocated_bytes);
         out[i].total_freed_bytes = atomic_load_size(&se_alloc_stats[i].total_freed_bytes);
         out[i].alloc_count = atomic_load_size(&se_alloc_stats[i].alloc_count);
         out[i].realloc_count = atomic_load_size(&se_alloc_stats[i].realloc_count);

data/ext/static_embeddings/se_f16.c

@@ -35,7 +35,7 @@ static void write_u16le(uint8_t *dst, uint16_t v) {
 }
 
 static uint16_t read_u16le(const uint8_t *src) {
-    return (uint16_t)src[0] | ((uint16_t)src[1] << 8);
+    return (uint16_t)((uint16_t)src[0] | ((uint16_t)src[1] << 8));
 }
 
 uint16_t se_float_to_f16_bits(float value) {

data/ext/static_embeddings/se_tokenizer.c

@@ -512,6 +512,8 @@ static int append_segment_cp(token_state_t *st, uint32_t cp) {
         st->segment_ascii = 0;
     if (!grow_u32(&st->scratch->cps, &st->scratch->cps_cap, st->segment_len + 1))
         return 0;
+    if (st->segment_len >= st->scratch->cps_cap)
+        return 0;
     st->scratch->cps[st->segment_len++] = cp;
     return 1;
 }
@@ -906,6 +908,10 @@ static se_status_t tokenize_ascii_run(token_state_t *st, const uint8_t *input, s
                     return oom(st, "building token segment");
                 }
             }
+            if (st->segment_len >= sc->cps_cap) {
+                *ip = i;
+                return oom(st, "building token segment");
+            }
             sc->cps[st->segment_len++] =
                 (lc && b >= 'A' && b <= 'Z') ? (uint32_t)(b + 32u) : (uint32_t)b;
             i++;

data/ext/static_embeddings/se_topk.c

@@ -73,8 +73,7 @@ float se_dot_product_f32(const float *q, const float *row, size_t dim) {
     return dot;
 }
 
-float se_dot_and_row_sq_f32(const float *q, const float *row, size_t dim,
-                            float *row_sq_out) {
+float se_dot_and_row_sq_f32(const float *q, const float *row, size_t dim, float *row_sq_out) {
     size_t j = 0;
 #if defined(SE_HAVE_NEON)
     float32x4_t d0 = vdupq_n_f32(0.0f);

data/ext/static_embeddings/static_embeddings.c

@@ -93,6 +93,7 @@ RUBY_FUNC_EXPORTED void Init_static_embeddings(void);
 typedef struct {
     se_model_t model;
     int open;
+    unsigned int active_ops;
 } model_wrapper_t;
 
 static void model_free(void *ptr) {
@@ -177,7 +178,7 @@ static void check_text_encoding_mode(VALUE str, long index, se_encoding_validati
                  rb_enc_name(enc));
     }
 
-    int cr = ENC_CODERANGE(str);
+    int cr = (int)ENC_CODERANGE(str);
     if (cr == ENC_CODERANGE_UNKNOWN) {
         if (mode == SE_VALIDATE_ENCODING_PREFIX)
             return;
@@ -216,6 +217,7 @@ typedef struct {
 
 typedef struct {
     const se_model_t *model;
+    model_wrapper_t *owner;
     batch_input_t input;
     uint32_t max_tokens;
     float *out;
@@ -228,6 +230,8 @@ typedef struct {
     se_error_t error;
 } batch_job_t;
 
+static void run_job(batch_job_t *job, size_t total_bytes);
+
 static void job_fail(batch_job_t *job, se_error_t err, size_t index) {
     if (job->failed)
         return;
@@ -361,6 +365,21 @@ static VALUE fiber_worker_ensure(VALUE wrapper) {
     return Qnil;
 }
 
+static VALUE run_job_body(VALUE arg) {
+    batch_job_t *job = (batch_job_t *)(uintptr_t)arg;
+    size_t total_bytes = 0;
+    for (size_t i = 0; i < job->input.count; i++)
+        total_bytes += job->input.lengths[i];
+    run_job(job, total_bytes);
+    return Qnil;
+}
+
+static VALUE run_job_ensure(VALUE arg) {
+    batch_job_t *job = (batch_job_t *)(uintptr_t)arg;
+    job->owner->active_ops--;
+    return Qnil;
+}
+
 static void run_job(batch_job_t *job, size_t total_bytes) {
     VALUE scheduler = current_fiber_scheduler();
 
@@ -386,11 +405,8 @@ static void run_job(batch_job_t *job, size_t total_bytes) {
 
 static VALUE protected_run_job(VALUE arg) {
     batch_job_t *job = (batch_job_t *)(uintptr_t)arg;
-    size_t total_bytes = 0;
-    for (size_t i = 0; i < job->input.count; i++)
-        total_bytes += job->input.lengths[i];
-    run_job(job, total_bytes);
-    return Qnil;
+    job->owner->active_ops++;
+    return rb_ensure(run_job_body, arg, run_job_ensure, arg);
 }
 
 static VALUE model_alloc(VALUE klass) {
@@ -422,6 +438,8 @@ static VALUE model_close(VALUE self) {
     model_wrapper_t *w;
     TypedData_Get_Struct(self, model_wrapper_t, &model_type, w);
     if (w->open) {
+        if (w->active_ops != 0)
+            rb_raise(eError, "model is busy with an active operation");
         se_model_close(&w->model);
         w->open = 0;
     }
@@ -474,18 +492,40 @@ typedef struct {
     size_t pages;
 } warmup_job_t;
 
+typedef struct {
+    VALUE self;
+    model_wrapper_t *owner;
+    warmup_job_t job;
+} warmup_run_t;
+
 static void *warmup_execute(void *arg) {
     warmup_job_t *job = (warmup_job_t *)arg;
     job->pages = se_model_warmup(job->model);
     return NULL;
 }
 
+static VALUE warmup_body(VALUE arg) {
+    warmup_run_t *run = (warmup_run_t *)(uintptr_t)arg;
+    rb_thread_call_without_gvl(warmup_execute, &run->job, NULL, NULL);
+    RB_GC_GUARD(run->self);
+    return Qnil;
+}
+
+static VALUE warmup_ensure(VALUE arg) {
+    warmup_run_t *run = (warmup_run_t *)(uintptr_t)arg;
+    run->owner->active_ops--;
+    return Qnil;
+}
+
 static VALUE model_warmup(VALUE self) {
     model_wrapper_t *w = get_model(self);
-    warmup_job_t job;
-    job.model = &w->model;
-    job.pages = 0;
-    rb_thread_call_without_gvl(warmup_execute, &job, NULL, NULL);
+    warmup_run_t run;
+    run.self = self;
+    run.owner = w;
+    run.job.model = &w->model;
+    run.job.pages = 0;
+    w->active_ops++;
+    rb_ensure(warmup_body, (VALUE)(uintptr_t)&run, warmup_ensure, (VALUE)(uintptr_t)&run);
     RB_GC_GUARD(self);
     return self;
 }
@@ -829,6 +869,7 @@ static VALUE embed_texts_internal(VALUE self, VALUE texts, int is_array, size_t
         batch_job_t job;
         memset(&job, 0, sizeof(job));
         job.model = model;
+        job.owner = w;
         job.max_tokens = max_tokens;
         job.out = run.out;
         job.stats = run.stats;
@@ -1116,6 +1157,7 @@ typedef struct {
 
 typedef struct {
     VALUE ids_value;
+    model_wrapper_t *owner;
     const se_model_t *model;
     uint32_t max_tokens;
     uint32_t *ids;
@@ -1126,6 +1168,7 @@ typedef struct {
     se_token_stats_t *stats_out;
     se_token_stats_t stats;
     int truncated;
+    int active;
 } ids_run_t;
 
 typedef struct {
@@ -1197,7 +1240,11 @@ static VALUE embed_token_ids_body(VALUE arg) {
     se_error_clear(&job.error);
 
     if (run->n >= SE_IDS_GVL_UNLOCK_THRESHOLD) {
+        run->owner->active_ops++;
+        run->active = 1;
         rb_thread_call_without_gvl(ids_execute, &job, ids_unblock_cancel, &job);
+        run->owner->active_ops--;
+        run->active = 0;
         rb_thread_check_ints();
         if (job.cancelled)
             rb_raise(rb_eInterrupt, "operation cancelled");
@@ -1225,6 +1272,8 @@ static VALUE embed_token_ids_body(VALUE arg) {
 
 static VALUE embed_token_ids_ensure(VALUE arg) {
     ids_run_t *run = (ids_run_t *)(uintptr_t)arg;
+    if (run->active)
+        run->owner->active_ops--;
     se_free(run->ids);
     se_free(run->out);
     run->ids = NULL;
@@ -1280,6 +1329,7 @@ static VALUE embed_token_ids_value(VALUE self, VALUE ids_value, VALUE max_tokens
     ids_run_t run;
     memset(&run, 0, sizeof(run));
     run.ids_value = ids_value;
+    run.owner = w;
     run.model = &w->model;
     run.max_tokens = max_tokens;
     run.n = n;