static_embeddings 0.1.3 → 0.1.5

data/ext/static_embeddings/static_embeddings.c

@@ -80,6 +80,7 @@ static ID id_blocking_p;
 static ID id_vector;
 static ID id_token_count;
 static ID id_unk_count;
+static ID id_pooled_count;
 static ID id_truncated;
 static ID id_dim;
 static ID id_allow_unfrozen;
@@ -274,10 +275,14 @@ static void batch_worker_run(batch_job_t *job, se_scratch_t *scratch) {
 
 static void *batch_execute(void *arg) {
     batch_job_t *job = (batch_job_t *)arg;
-    se_scratch_t scratch;
-    se_scratch_init(&scratch);
-    batch_worker_run(job, &scratch);
-    se_scratch_free(&scratch);
+    se_scratch_t *scratch = se_scratch_acquire(job->model->meta.dim);
+    if (!scratch) {
+        se_error_set(&job->error, SE_ERR_OOM, "out of memory while sizing scratch buffers");
+        job->failed = 1;
+        return NULL;
+    }
+    batch_worker_run(job, scratch);
+    se_scratch_release(scratch);
     return NULL;
 }
 
@@ -928,6 +933,32 @@ static VALUE embed_one_via_batch(VALUE self, VALUE text, VALUE max_tokens_opt,
     return embed_texts_internal(self, text, 0, 1, max_tokens_opt, format, validation, stats);
 }
 
+typedef struct {
+    se_scratch_t *scratch;
+    const se_model_t *model;
+    const uint8_t *input;
+    size_t input_len;
+    uint32_t max_tokens;
+    float *out;
+    se_token_stats_t *stats;
+    se_error_t err;
+    se_status_t rc;
+} embed_one_scratch_job_t;
+
+static VALUE embed_one_scratch_body(VALUE arg) {
+    embed_one_scratch_job_t *job = (embed_one_scratch_job_t *)(uintptr_t)arg;
+    job->rc = se_embed_one(job->model, job->scratch, job->input, job->input_len, job->max_tokens,
+                           job->out, job->stats, &job->err, NULL);
+    return Qnil;
+}
+
+static VALUE embed_one_scratch_ensure(VALUE arg) {
+    embed_one_scratch_job_t *job = (embed_one_scratch_job_t *)(uintptr_t)arg;
+    se_scratch_release(job->scratch);
+    job->scratch = NULL;
+    return Qnil;
+}
+
 static VALUE embed_one_value(VALUE self, VALUE text, VALUE max_tokens_opt,
                              se_vector_format_t format, se_encoding_validation_t validation,
                              se_token_stats_t *stats) {
@@ -947,25 +978,28 @@ static VALUE embed_one_value(VALUE self, VALUE text, VALUE max_tokens_opt,
     rb_enc_associate(result, binary_encoding);
     float *out = (float *)RSTRING_PTR(result);
 
-    se_scratch_t scratch;
-    se_scratch_init(&scratch);
-    if (!se_scratch_reserve(&scratch, dim)) {
-        se_scratch_free(&scratch);
+    se_token_stats_t local_stats;
+    embed_one_scratch_job_t job;
+    memset(&job, 0, sizeof(job));
+    job.model = &w->model;
+    job.input = (const uint8_t *)RSTRING_PTR(text);
+    job.input_len = (size_t)RSTRING_LEN(text);
+    job.max_tokens = resolve_max_tokens(&w->model, max_tokens_opt);
+    job.out = out;
+    job.stats = stats ? stats : &local_stats;
+    se_error_clear(&job.err);
+
+    job.scratch = se_scratch_acquire(dim);
+    if (!job.scratch)
         rb_raise(rb_eNoMemError, "out of memory");
-    }
 
-    se_error_t err;
-    se_error_clear(&err);
-    se_token_stats_t local_stats;
-    se_status_t rc =
-        se_embed_one(&w->model, &scratch, (const uint8_t *)RSTRING_PTR(text),
-                     (size_t)RSTRING_LEN(text), resolve_max_tokens(&w->model, max_tokens_opt), out,
-                     stats ? stats : &local_stats, &err, NULL);
-    se_scratch_free(&scratch);
+    rb_ensure(embed_one_scratch_body, (VALUE)(uintptr_t)&job, embed_one_scratch_ensure,
+              (VALUE)(uintptr_t)&job);
     RB_GC_GUARD(text);
+    RB_GC_GUARD(result);
 
-    if (rc != SE_OK)
-        raise_se(&err);
+    if (job.rc != SE_OK)
+        raise_se(&job.err);
 
     return result;
 }
@@ -997,10 +1031,46 @@ static VALUE model_embed_with_stats(int argc, VALUE *argv, VALUE self) {
     rb_hash_aset(hash, ID2SYM(id_vector), vector);
     rb_hash_aset(hash, ID2SYM(id_token_count), UINT2NUM(stats.token_count));
     rb_hash_aset(hash, ID2SYM(id_unk_count), UINT2NUM(stats.unk_count));
+    uint32_t pooled_count = get_model(self)->model.meta.unk_policy == SE_UNK_DROP
+                                ? stats.token_count - stats.unk_count
+                                : stats.token_count;
+    rb_hash_aset(hash, ID2SYM(id_pooled_count), UINT2NUM(pooled_count));
     rb_hash_aset(hash, ID2SYM(id_truncated), stats.truncated ? Qtrue : Qfalse);
     return hash;
 }
 
+typedef struct {
+    se_scratch_t *scratch;
+    const se_model_t *model;
+    const uint8_t *input;
+    size_t input_len;
+    uint32_t max_tokens;
+    se_token_stats_t stats;
+    se_error_t err;
+    se_status_t rc;
+    VALUE ids;
+} tokenize_scratch_job_t;
+
+static VALUE tokenize_scratch_body(VALUE arg) {
+    tokenize_scratch_job_t *job = (tokenize_scratch_job_t *)(uintptr_t)arg;
+    job->rc = se_tokenize(job->model, job->scratch, job->input, job->input_len, job->max_tokens,
+                          SE_TOKEN_LIMIT_RAW, &job->stats, &job->err, NULL);
+    if (job->rc != SE_OK)
+        return Qnil;
+
+    job->ids = rb_ary_new_capa((long)job->stats.token_count);
+    for (uint32_t i = 0; i < job->stats.token_count; i++)
+        rb_ary_push(job->ids, UINT2NUM(job->scratch->ids[i]));
+    return Qnil;
+}
+
+static VALUE tokenize_scratch_ensure(VALUE arg) {
+    tokenize_scratch_job_t *job = (tokenize_scratch_job_t *)(uintptr_t)arg;
+    se_scratch_release(job->scratch);
+    job->scratch = NULL;
+    return Qnil;
+}
+
 static VALUE model_tokenize(int argc, VALUE *argv, VALUE self) {
     VALUE text, opts;
     rb_scan_args(argc, argv, "1:", &text, &opts);
@@ -1011,32 +1081,27 @@ static VALUE model_tokenize(int argc, VALUE *argv, VALUE self) {
     check_text_encoding_mode(
         text, -1, resolve_encoding_validation(lookup_option(opts, id_validate_encoding)));
 
-    uint32_t max_tokens = resolve_max_tokens(&w->model, lookup_option(opts, id_max_tokens));
-
-    se_scratch_t scratch;
-    se_scratch_init(&scratch);
-    if (!se_scratch_reserve(&scratch, w->model.meta.dim)) {
-        se_scratch_free(&scratch);
+    tokenize_scratch_job_t job;
+    memset(&job, 0, sizeof(job));
+    job.model = &w->model;
+    job.input = (const uint8_t *)RSTRING_PTR(text);
+    job.input_len = (size_t)RSTRING_LEN(text);
+    job.max_tokens = resolve_max_tokens(&w->model, lookup_option(opts, id_max_tokens));
+    job.ids = Qnil;
+    se_error_clear(&job.err);
+
+    job.scratch = se_scratch_acquire(w->model.meta.dim);
+    if (!job.scratch)
         rb_raise(rb_eNoMemError, "out of memory");
-    }
 
-    se_token_stats_t stats;
-    se_error_t err;
-    se_error_clear(&err);
-    se_status_t rc = se_tokenize(&w->model, &scratch, (const uint8_t *)RSTRING_PTR(text),
-                                 (size_t)RSTRING_LEN(text), max_tokens, &stats, &err, NULL);
-    if (rc != SE_OK) {
-        se_scratch_free(&scratch);
-        raise_se(&err);
-    }
+    rb_ensure(tokenize_scratch_body, (VALUE)(uintptr_t)&job, tokenize_scratch_ensure,
+              (VALUE)(uintptr_t)&job);
+    RB_GC_GUARD(text);
 
-    VALUE ids = rb_ary_new_capa((long)stats.token_count);
-    for (uint32_t i = 0; i < stats.token_count; i++)
-        rb_ary_push(ids, UINT2NUM(scratch.ids[i]));
+    if (job.rc != SE_OK)
+        raise_se(&job.err);
 
-    se_scratch_free(&scratch);
-    RB_GC_GUARD(text);
-    return ids;
+    return job.ids;
 }
 
 typedef struct {
@@ -1077,17 +1142,15 @@ static VALUE num2ull_at_value(VALUE arg) {
 
 static void *ids_execute(void *arg) {
     ids_job_t *job = (ids_job_t *)arg;
-    se_scratch_t scratch;
-    se_scratch_init(&scratch);
-    if (!se_scratch_reserve(&scratch, job->model->meta.dim)) {
+    se_scratch_t *scratch = se_scratch_acquire(job->model->meta.dim);
+    if (!scratch) {
         se_error_set(&job->error, SE_ERR_OOM, "out of memory while sizing scratch buffers");
-        se_scratch_free(&scratch);
         return NULL;
     }
     se_error_clear(&job->error);
-    se_embed_ids(job->model, &scratch, job->ids, job->n_ids, job->out, &job->stats, &job->error,
+    se_embed_ids(job->model, scratch, job->ids, job->n_ids, job->out, &job->stats, &job->error,
                  &job->cancelled);
-    se_scratch_free(&scratch);
+    se_scratch_release(scratch);
     return NULL;
 }
 
@@ -1179,7 +1242,27 @@ static VALUE embed_token_ids_value(VALUE self, VALUE ids_value, VALUE max_tokens
     size_t n = (size_t)n_long;
     int truncated = 0;
 
-    if (max_tokens != 0 && n > (size_t)max_tokens) {
+    if (max_tokens != 0 && w->model.meta.unk_policy == SE_UNK_DROP) {
+        size_t usable = 0;
+        size_t cutoff = n;
+        for (size_t i = 0; i < n; i++) {
+            unsigned long long value = NUM2ULL(rb_ary_entry(ids_value, (long)i));
+            if (value >= w->model.meta.vocab_size)
+                rb_raise(rb_eArgError, "token id at index %zu is out of range", i);
+
+            if ((uint32_t)value == w->model.meta.unk_id)
+                continue;
+
+            usable++;
+            if (usable == (size_t)max_tokens) {
+                cutoff = i + 1;
+            } else if (usable > (size_t)max_tokens) {
+                n = cutoff;
+                truncated = 1;
+                break;
+            }
+        }
+    } else if (max_tokens != 0 && n > (size_t)max_tokens) {
         n = (size_t)max_tokens;
         truncated = 1;
     }
@@ -1239,6 +1322,10 @@ static VALUE model_embed_token_ids_with_stats(int argc, VALUE *argv, VALUE self)
     rb_hash_aset(hash, ID2SYM(id_vector), vector);
     rb_hash_aset(hash, ID2SYM(id_token_count), UINT2NUM(stats.token_count));
     rb_hash_aset(hash, ID2SYM(id_unk_count), UINT2NUM(stats.unk_count));
+    uint32_t pooled_count = get_model(self)->model.meta.unk_policy == SE_UNK_DROP
+                                ? stats.token_count - stats.unk_count
+                                : stats.token_count;
+    rb_hash_aset(hash, ID2SYM(id_pooled_count), UINT2NUM(pooled_count));
     rb_hash_aset(hash, ID2SYM(id_truncated), stats.truncated ? Qtrue : Qfalse);
     return hash;
 }
@@ -1549,6 +1636,7 @@ RUBY_FUNC_EXPORTED void Init_static_embeddings(void) {
     id_vector = rb_intern("vector");
     id_token_count = rb_intern("token_count");
     id_unk_count = rb_intern("unk_count");
+    id_pooled_count = rb_intern("pooled_count");
     id_truncated = rb_intern("truncated");
     id_dim = rb_intern("dim");
     id_allow_unfrozen = rb_intern("allow_unfrozen");

data/lib/static_embeddings/cli.rb

@@ -65,6 +65,7 @@ module StaticEmbeddings
     end
 
     def convert(argv)
+      require "static_embeddings/converter"
       options = parse_convert_options(argv)
       source = required_arg(argv, "usage: static_embeddings convert SOURCE_DIR [--out PATH]")
       model_id = options[:id] || File.basename(File.expand_path(source))

data/lib/static_embeddings/converter.rb

@@ -1,13 +1,25 @@
 require "json"
 require "digest"
+require "static_embeddings/format"
+require "static_embeddings/safetensors"
+require "static_embeddings/unicode_tables"
 
 module StaticEmbeddings
   class Converter
     REFERENCE_IMPL = "model2vec.StaticModel"
+    REFERENCE_MODEL2VEC_VERSION = "0.9.0"
+    REFERENCE_TOKENIZERS_VERSION = "0.23.1"
+    REFERENCE_UNICODE_CATEGORIES_VERSION = "0.1.1"
     REFERENCE_MAX_TOKENS = 512
 
     ALLOWED_NORMALIZER_KEYS = %w[type clean_text handle_chinese_chars strip_accents lowercase].freeze
-    STANDARD_SPECIAL_TOKENS = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"].freeze
+    STANDARD_SPECIAL_TOKENS = {
+      "[PAD]" => Format::ADDED_PAD,
+      "[UNK]" => Format::ADDED_UNK,
+      "[CLS]" => Format::ADDED_CLS,
+      "[SEP]" => Format::ADDED_SEP,
+      "[MASK]" => Format::ADDED_MASK
+    }.freeze
     SOURCE_FILES = %w[tokenizer.json config.json tokenizer_config.json model.safetensors].freeze
     TOKENIZER_PROFILE = "BERT_WORDPIECE_V1"
 
@@ -22,6 +34,7 @@ module StaticEmbeddings
       source = load_source
       profile = audit_tokenizer(source[:tokenizer], source[:tokenizer_config])
       tokens = extract_vocab(source[:tokenizer])
+      validate_added_token_ids!(profile[:added_tokens], tokens)
       matrix, dim = extract_matrix(tokens.length)
       meta = runtime_meta(source[:config], profile, tokens, max_tokens)
 
@@ -62,9 +75,13 @@ module StaticEmbeddings
       assert_wordpiece!(model)
       assert_normalizer!(normalizer)
       assert_pre_tokenizer!(pre_tokenizer)
-      audit_added_tokens(tokenizer)
+      added_tokens = audit_added_tokens(tokenizer)
 
       profile = profile_from(model, normalizer, tokenizer_config)
+      profile[:added_tokens] = added_tokens
+      profile[:added_token_mask] = added_tokens.reduce(0) do |mask, token|
+        mask | STANDARD_SPECIAL_TOKENS.fetch(token.fetch("content"))
+      end
       reject!("clean_text=false is not supported by the runtime") unless profile[:clean_text]
       profile
     end
@@ -123,17 +140,40 @@ module StaticEmbeddings
     def audit_added_tokens(tokenizer)
       added = tokenizer["added_tokens"] || []
       bad_content = added.reject { |token| standard_special?(token) }
-      reject!("tokenizer declares non-standard added_tokens #{bad_content.map { |t| t['content'] }.inspect}") unless bad_content.empty?
+      unless bad_content.empty?
+        reject!("tokenizer declares non-standard added_tokens #{bad_content.map { |t| t['content'] }.inspect}")
+      end
 
       whitespace = added.find { |token| token["content"].to_s.match?(/\s/) }
       reject!("added token #{whitespace['content'].inspect} contains whitespace") if whitespace
 
       flagged = added.find { |token| token["lstrip"] || token["rstrip"] || token["single_word"] }
       reject!("added token #{flagged['content'].inspect} uses lstrip/rstrip/single_word") if flagged
+
+      normalized = added.find { |token| token["normalized"] != false }
+      if normalized
+        reject!("added token #{normalized['content'].inspect} must use normalized=false")
+      end
+
+      duplicate = added.group_by { |token| token["content"] }.find { |_, rows| rows.length > 1 }
+      reject!("duplicate added token #{duplicate[0].inspect}") if duplicate
+
+      added
     end
 
     def standard_special?(token)
-      token["special"] && STANDARD_SPECIAL_TOKENS.include?(token["content"])
+      token["special"] && STANDARD_SPECIAL_TOKENS.key?(token["content"])
+    end
+
+    def validate_added_token_ids!(added_tokens, tokens)
+      added_tokens.each do |token|
+        content = token.fetch("content")
+        id = token["id"]
+        reject!("added token #{content.inspect} has non-integer id #{id.inspect}") unless id.is_a?(Integer)
+        unless id.between?(0, tokens.length - 1) && tokens[id] == content
+          reject!("added token #{content.inspect} id #{id} does not match model.vocab")
+        end
+      end
     end
 
     def extract_vocab(tokenizer)
@@ -159,7 +199,7 @@ module StaticEmbeddings
       path = File.join(source_dir, "model.safetensors")
       raise ConversionError, "missing model.safetensors in #{source_dir}" unless File.file?(path)
 
-      name, tensor = sole_matrix_tensor(Safetensors.read(path)[:tensors])
+      name, tensor = sole_matrix_tensor(Safetensors.describe(path)[:tensors])
       rows, dim = tensor[:shape]
       if rows != vocab_size
         raise ConversionError,
@@ -167,7 +207,7 @@ module StaticEmbeddings
               "#{vocab_size} tokens — refusing to guess the mapping"
       end
 
-      [tensor[:bytes], dim]
+      [Safetensors.f32_payload(path, tensor), dim]
     end
 
     def sole_matrix_tensor(tensors)
@@ -203,6 +243,7 @@ module StaticEmbeddings
         strip_accents: profile[:strip_accents],
         handle_chinese_chars: profile[:handle_chinese_chars],
         clean_text: profile[:clean_text],
+        added_token_mask: profile.fetch(:added_token_mask),
         max_input_chars_per_word: profile[:max_input_chars_per_word],
         max_token_chars: max_token_chars(tokens, profile[:continuing_subword_prefix]),
         subword_prefix: profile[:continuing_subword_prefix]
@@ -241,7 +282,7 @@ module StaticEmbeddings
     def source_digests
       SOURCE_FILES.each_with_object({}) do |name, acc|
         path = File.join(source_dir, name)
-        acc[name] = Digest::SHA256.hexdigest(File.binread(path)) if File.file?(path)
+        acc[name] = Digest::SHA256.file(path).hexdigest if File.file?(path)
       end
     end
 
@@ -252,6 +293,9 @@ module StaticEmbeddings
         "source_model_id" => model_id || File.basename(File.expand_path(source_dir)),
         "source_files_sha256" => source_digests,
         "reference_impl" => REFERENCE_IMPL,
+        "reference_model2vec_version" => REFERENCE_MODEL2VEC_VERSION,
+        "reference_tokenizers_version" => REFERENCE_TOKENIZERS_VERSION,
+        "reference_unicode_categories_version" => REFERENCE_UNICODE_CATEGORIES_VERSION,
         "reference_max_tokens" => meta[:max_tokens_default],
         "unicode_source" => UnicodeTables.source_stamp,
         "tokenizer_profile" => TOKENIZER_PROFILE,

data/lib/static_embeddings/format.rb

@@ -3,7 +3,7 @@ require "digest"
 module StaticEmbeddings
   module Format
     MAGIC = "SEMBv1\0\0"
-    VERSION = 2
+    VERSION = 3
     HEADER_SIZE = 320
     ALIGNMENT = 64
 
@@ -12,7 +12,7 @@ module StaticEmbeddings
     POOLING_MEAN = 1
     NORMALIZATION_NONE = 0
     NORMALIZATION_L2 = 1
-    TRUNCATE_IDS_BEFORE_POOLING = 1
+    TRUNCATE_USABLE_IDS_BEFORE_POOLING = 2
     UNK_INCLUDE = 0
     UNK_DROP = 1
     EMPTY_ZERO_VECTOR = 0
@@ -26,6 +26,14 @@ module StaticEmbeddings
     CHECKSUM_OFFSET = 240
     CHECKSUM_SIZE = 32
     MAX_PROBE_OFFSET = 304
+    ADDED_TOKEN_MASK_OFFSET = 308
+
+    ADDED_PAD  = 1 << 0
+    ADDED_UNK  = 1 << 1
+    ADDED_CLS  = 1 << 2
+    ADDED_SEP  = 1 << 3
+    ADDED_MASK = 1 << 4
+    ADDED_TOKEN_MASK_ALL = ADDED_PAD | ADDED_UNK | ADDED_CLS | ADDED_SEP | ADDED_MASK
 
     SECTION_FIELDS = {
       vocab_strings: 128,
@@ -44,7 +52,7 @@ module StaticEmbeddings
       28 => TOKENIZER_BERT_WORDPIECE_V1,
       32 => DTYPE_F32,
       36 => POOLING_MEAN,
-      48 => TRUNCATE_IDS_BEFORE_POOLING,
+      48 => TRUNCATE_USABLE_IDS_BEFORE_POOLING,
       108 => HASH_SEED
     }.freeze
 
@@ -102,7 +110,7 @@ module StaticEmbeddings
     def write(path:, meta:, tokens:, matrix:, norm_tables:, provenance:)
       hash_size, strings, hash_blob, max_probe = build_hash_table(tokens)
       root_trie, continuation_trie = build_wordpiece_tries(tokens, meta.fetch(:subword_prefix))
-      sections, body = build_body(
+      payloads = {
         vocab_strings: strings,
         vocab_hash: hash_blob,
         embeddings: matrix,
@@ -110,15 +118,48 @@ module StaticEmbeddings
         provenance: provenance,
         root_trie: root_trie,
         continuation_trie: continuation_trie
-      )
-
+      }
+      sections, file_size = layout_sections(payloads)
       header = build_header(meta, tokens.length, hash_size, max_probe, sections)
-      file = header << body
-      digest = Digest::SHA256.digest(file)
-      file[CHECKSUM_OFFSET, CHECKSUM_SIZE] = digest
 
-      File.binwrite(path, file)
-      { bytes: file.bytesize, sha256: digest.unpack1("H*"), hash_table_size: hash_size }
+      digest = Digest::SHA256.new
+      File.open(path, "wb") do |io|
+        io.write(header)
+        digest << header
+        offset = HEADER_SIZE
+
+        payloads.each do |name, payload|
+          target = sections.fetch(name).first
+          padding = target - offset
+          if padding.positive?
+            zeros = "\0".b * padding
+            io.write(zeros)
+            digest << zeros
+          end
+          write_payload(io, digest, payload)
+          offset = target + payload.bytesize
+        end
+      end
+
+      checksum = digest.digest
+      File.open(path, "r+b") do |io|
+        io.seek(CHECKSUM_OFFSET, IO::SEEK_SET)
+        io.write(checksum)
+      end
+
+      { bytes: file_size, sha256: checksum.unpack1("H*"), hash_table_size: hash_size }
+    end
+
+    def write_payload(io, digest, payload)
+      if payload.respond_to?(:each_chunk)
+        payload.each_chunk do |chunk|
+          io.write(chunk)
+          digest << chunk
+        end
+      else
+        io.write(payload)
+        digest << payload
+      end
     end
 
     def verify(path)
@@ -240,20 +281,16 @@ module StaticEmbeddings
       end
     end
 
-    def build_body(payloads)
-      body = binary_string
+    def layout_sections(payloads)
+      offset = HEADER_SIZE
       sections = {}
       payloads.each do |name, payload|
-        align_body!(body)
-        sections[name] = [HEADER_SIZE + body.bytesize, payload.bytesize]
-        body << payload
+        padding = (ALIGNMENT - (offset % ALIGNMENT)) % ALIGNMENT
+        offset += padding
+        sections[name] = [offset, payload.bytesize]
+        offset += payload.bytesize
       end
-      [sections, body]
-    end
-
-    def align_body!(body)
-      padding = (ALIGNMENT - ((HEADER_SIZE + body.bytesize) % ALIGNMENT)) % ALIGNMENT
-      body << "\0".b * padding if padding.positive?
+      [sections, offset]
     end
 
     def build_header(meta, vocab_size, hash_size, max_probe, sections)
@@ -268,6 +305,7 @@ module StaticEmbeddings
       put_u32(header, 104, hash_size)
       put_u32(header, MAX_TOKEN_CHARS_OFFSET, meta.fetch(:max_token_chars))
       put_u32(header, MAX_PROBE_OFFSET, max_probe)
+      put_u32(header, ADDED_TOKEN_MASK_OFFSET, meta.fetch(:added_token_mask, 0))
       put_prefix(header, meta.fetch(:subword_prefix))
       SECTION_FIELDS.each { |name, field| put_section(header, field, sections.fetch(name)) }
 

data/lib/static_embeddings/paths.rb

@@ -11,7 +11,24 @@ module StaticEmbeddings
     end
 
     def model_path(model_id, env = ENV)
-      File.join(cache_dir(env), "models", "#{model_id}.semb")
+      id = model_id.to_s
+      raise ArgumentError, "model_id must not be empty" if id.empty?
+      raise ArgumentError, "model_id contains a NUL byte" if id.include?("\0")
+
+      normalized = id.tr("\\", "/")
+      parts = normalized.split("/", -1)
+      if normalized.start_with?("/") ||
+         parts.any? { |part| part.empty? || part == "." || part == ".." || part.include?(":") }
+        raise ArgumentError, "model_id must be a relative slash-separated identifier"
+      end
+
+      base = File.expand_path(File.join(cache_dir(env), "models"))
+      path = File.expand_path(File.join(base, "#{normalized}.semb"))
+      prefix = base.end_with?(File::SEPARATOR) ? base : "#{base}#{File::SEPARATOR}"
+      unless path.start_with?(prefix)
+        raise ArgumentError, "model_id escapes the model cache"
+      end
+      path
     end
 
     def builtin_path(name)