static_embeddings 0.1.4 → 0.1.5

data/ext/static_embeddings/se_format.c

@@ -102,7 +102,7 @@ int se_vocab_lookup_piece(const se_model_t *model, const uint8_t *prefix, size_t
     uint32_t pos = h & mask;
     uint32_t total_len = (uint32_t)(prefix_len + len);
 
-    for (uint32_t probe = 0; probe <= mask; probe++) {
+    for (uint32_t probe = 0; probe < model->meta.max_probe; probe++) {
         const se_vocab_slot_t *slot = &model->vocab_hash[pos];
         if (slot->token_id == SE_SLOT_EMPTY)
             return 0;
@@ -556,6 +556,31 @@ static se_status_t validate_vocab_hash(const se_model_t *model, struct section v
         se_error_set(err, SE_ERR_INVALID_FORMAT,
                      "[UNK] is not reachable through the vocabulary hash");
         rc = SE_ERR_INVALID_FORMAT;
+        goto done;
+    }
+
+    static const struct {
+        uint32_t bit;
+        const char *text;
+        uint32_t len;
+    } added[] = {
+        {SE_ADDED_PAD, "[PAD]", 5u}, {SE_ADDED_UNK, "[UNK]", 5u},   {SE_ADDED_CLS, "[CLS]", 5u},
+        {SE_ADDED_SEP, "[SEP]", 5u}, {SE_ADDED_MASK, "[MASK]", 6u},
+    };
+    const uint32_t ids[] = {m->pad_id, m->unk_id, m->cls_id, m->sep_id, m->mask_id};
+
+    for (size_t i = 0; i < SE_ARRAY_LEN(added); i++) {
+        if ((m->added_token_mask & added[i].bit) == 0)
+            continue;
+        uint32_t got = 0;
+        if (ids[i] >= m->vocab_size ||
+            !se_vocab_lookup(model, (const uint8_t *)added[i].text, added[i].len, &got) ||
+            got != ids[i]) {
+            se_error_set(err, SE_ERR_INVALID_FORMAT,
+                         "added token %s is not reachable at its recorded id", added[i].text);
+            rc = SE_ERR_INVALID_FORMAT;
+            goto done;
+        }
     }
 
 done:
@@ -646,6 +671,7 @@ static se_status_t validate(se_model_t *model, se_error_t *err) {
     m->subword_prefix_len = read_u32(base, SE_OFF_SUBWORD_PREFIX_LEN);
     m->max_token_chars = read_u32(base, SE_OFF_MAX_TOKEN_CHARS);
     m->max_probe = read_u32(base, SE_OFF_MAX_PROBE);
+    m->added_token_mask = read_u32(base, SE_OFF_ADDED_TOKEN_MASK);
     memcpy(m->subword_prefix, base + SE_OFF_SUBWORD_PREFIX, 8);
 
     if (m->tokenizer_type != SE_TOKENIZER_BERT_WORDPIECE_V1) {
@@ -669,7 +695,7 @@ static se_status_t validate(se_model_t *model, se_error_t *err) {
                      m->normalization_type);
         return SE_ERR_INVALID_FORMAT;
     }
-    if (m->truncation_policy != SE_TRUNCATE_IDS_BEFORE_POOLING) {
+    if (m->truncation_policy != SE_TRUNCATE_USABLE_IDS_BEFORE_POOLING) {
         se_error_set(err, SE_ERR_INVALID_FORMAT, "truncation policy %u is not supported",
                      m->truncation_policy);
         return SE_ERR_INVALID_FORMAT;
@@ -682,6 +708,11 @@ static se_status_t validate(se_model_t *model, se_error_t *err) {
         se_error_set(err, SE_ERR_INVALID_FORMAT, "unknown UNK policy %u", m->unk_policy);
         return SE_ERR_INVALID_FORMAT;
     }
+    if ((m->added_token_mask & ~SE_ADDED_TOKEN_MASK_ALL) != 0) {
+        se_error_set(err, SE_ERR_INVALID_FORMAT, "unknown added-token mask bits 0x%x",
+                     m->added_token_mask & ~SE_ADDED_TOKEN_MASK_ALL);
+        return SE_ERR_INVALID_FORMAT;
+    }
     if (m->empty_policy != SE_EMPTY_ZERO_VECTOR && m->empty_policy != SE_EMPTY_RAISE) {
         se_error_set(err, SE_ERR_INVALID_FORMAT, "unknown empty-input policy %u", m->empty_policy);
         return SE_ERR_INVALID_FORMAT;
@@ -704,8 +735,14 @@ static se_status_t validate(se_model_t *model, se_error_t *err) {
                      m->hash_table_size);
         return SE_ERR_INVALID_FORMAT;
     }
-    if (m->hash_table_size < m->vocab_size) {
-        se_error_set(err, SE_ERR_INVALID_FORMAT, "hash table smaller than vocabulary");
+    if (m->hash_table_size <= m->vocab_size ||
+        (uint64_t)m->vocab_size * 100u > (uint64_t)m->hash_table_size * 70u) {
+        se_error_set(err, SE_ERR_INVALID_FORMAT,
+                     "hash table load factor exceeds the supported 0.70 maximum");
+        return SE_ERR_INVALID_FORMAT;
+    }
+    if (m->max_probe == 0 || m->max_probe > m->hash_table_size) {
+        se_error_set(err, SE_ERR_INVALID_FORMAT, "invalid recorded max probe %u", m->max_probe);
         return SE_ERR_INVALID_FORMAT;
     }
     if (m->subword_prefix_len > 8) {

data/ext/static_embeddings/se_internal.h

@@ -8,7 +8,7 @@
 
 #define SE_MAGIC          "SEMBv1\0\0"
 #define SE_MAGIC_LEN      8
-#define SE_FORMAT_VERSION 2u
+#define SE_FORMAT_VERSION 3u
 #define SE_HEADER_SIZE    320u
 #define SE_ALIGNMENT      64u
 
@@ -18,7 +18,7 @@
 #define SE_NORMALIZATION_NONE          0u
 #define SE_NORMALIZATION_L2            1u
 
-#define SE_TRUNCATE_IDS_BEFORE_POOLING 1u
+#define SE_TRUNCATE_USABLE_IDS_BEFORE_POOLING 2u
 
 #define SE_UNK_INCLUDE 0u
 #define SE_UNK_DROP    1u
@@ -26,6 +26,14 @@
 #define SE_EMPTY_ZERO_VECTOR 0u
 #define SE_EMPTY_RAISE       1u
 
+#define SE_ADDED_PAD  (1u << 0)
+#define SE_ADDED_UNK  (1u << 1)
+#define SE_ADDED_CLS  (1u << 2)
+#define SE_ADDED_SEP  (1u << 3)
+#define SE_ADDED_MASK (1u << 4)
+#define SE_ADDED_TOKEN_MASK_ALL \
+    (SE_ADDED_PAD | SE_ADDED_UNK | SE_ADDED_CLS | SE_ADDED_SEP | SE_ADDED_MASK)
+
 #define SE_SLOT_EMPTY 0xFFFFFFFFu
 #define SE_SIZE_MAX   ((size_t)-1)
 
@@ -83,7 +91,8 @@ enum {
     SE_OFF_SEC_ROOT_TRIE = 208,
     SE_OFF_SEC_CONT_TRIE = 224,
     SE_OFF_CHECKSUM = 240,
-    SE_OFF_MAX_PROBE = 304
+    SE_OFF_MAX_PROBE = 304,
+    SE_OFF_ADDED_TOKEN_MASK = 308
 };
 
 typedef struct {
@@ -163,6 +172,7 @@ typedef struct {
     uint32_t max_input_chars_per_word;
     uint32_t max_token_chars;
     uint32_t max_probe;
+    uint32_t added_token_mask;
     uint32_t pad_id;
     uint32_t unk_id;
     uint32_t cls_id;
@@ -401,9 +411,11 @@ typedef struct {
     uint32_t truncated;
 } se_token_stats_t;
 
+typedef enum { SE_TOKEN_LIMIT_RAW = 0, SE_TOKEN_LIMIT_USABLE = 1 } se_token_limit_t;
+
 se_status_t se_tokenize(const se_model_t *model, se_scratch_t *scratch, const uint8_t *input,
-                        size_t input_len, uint32_t max_tokens, se_token_stats_t *stats,
-                        se_error_t *err, volatile sig_atomic_t *cancelled);
+                        size_t input_len, uint32_t max_tokens, se_token_limit_t limit_mode,
+                        se_token_stats_t *stats, se_error_t *err, volatile sig_atomic_t *cancelled);
 
 se_status_t se_embed_one(const se_model_t *model, se_scratch_t *scratch, const uint8_t *input,
                          size_t input_len, uint32_t max_tokens, float *out, se_token_stats_t *stats,

data/ext/static_embeddings/se_tokenizer.c

@@ -87,16 +87,17 @@ static int push_id(se_scratch_t *sc, size_t *n_ids, uint32_t id) {
     return 1;
 }
 
-#define SE_SCRATCH_CPS_KEEP   256
-#define SE_SCRATCH_IDS_KEEP   512
-#define SE_SCRATCH_BYTES_KEEP 1024
+#define SE_SCRATCH_CPS_KEEP       256
+#define SE_SCRATCH_IDS_RESERVE    512
+#define SE_SCRATCH_IDS_RETAIN_MAX 8192
+#define SE_SCRATCH_BYTES_KEEP     1024
 
 int se_scratch_reserve(se_scratch_t *s, uint32_t dim) {
     if (!grow_u32(&s->cps, &s->cps_cap, SE_SCRATCH_CPS_KEEP))
         return 0;
     if (!grow_u32(&s->cps2, &s->cps2_cap, SE_SCRATCH_CPS_KEEP))
         return 0;
-    if (!grow_u32(&s->ids, &s->ids_cap, SE_SCRATCH_IDS_KEEP))
+    if (!grow_u32(&s->ids, &s->ids_cap, SE_SCRATCH_IDS_RESERVE))
         return 0;
     if (!grow_bytes(&s->bytes, &s->bytes_cap, SE_SCRATCH_BYTES_KEEP))
         return 0;
@@ -149,7 +150,8 @@ static int shrink_bytes(uint8_t **buf, size_t *cap, size_t keep) {
 static void se_scratch_trim(se_scratch_t *s) {
     (void)shrink_u32(&s->cps, &s->cps_cap, SE_SCRATCH_CPS_KEEP);
     (void)shrink_u32(&s->cps2, &s->cps2_cap, SE_SCRATCH_CPS_KEEP);
-    (void)shrink_u32(&s->ids, &s->ids_cap, SE_SCRATCH_IDS_KEEP);
+    if (s->ids_cap > SE_SCRATCH_IDS_RETAIN_MAX)
+        (void)shrink_u32(&s->ids, &s->ids_cap, SE_SCRATCH_IDS_RETAIN_MAX);
     (void)shrink_bytes(&s->bytes, &s->bytes_cap, SE_SCRATCH_BYTES_KEEP);
 }
 
@@ -393,6 +395,37 @@ static int cp_is_cjk_segment(const se_model_t *m, uint32_t cp) {
     return m->meta.handle_chinese_chars && cp >= 0x3400 && se_is_cjk(cp);
 }
 
+typedef struct {
+    const char *text;
+    uint32_t len;
+    uint32_t bit;
+} se_added_token_spec_t;
+
+static const se_added_token_spec_t SE_ADDED_TOKENS[] = {
+    {"[MASK]", 6u, SE_ADDED_MASK}, {"[PAD]", 5u, SE_ADDED_PAD}, {"[UNK]", 5u, SE_ADDED_UNK},
+    {"[CLS]", 5u, SE_ADDED_CLS},   {"[SEP]", 5u, SE_ADDED_SEP},
+};
+
+static int boundary_splits_added_token(const se_model_t *model, const uint8_t *input,
+                                       size_t input_len, size_t boundary) {
+    if (model->meta.added_token_mask == 0 || boundary == 0 || boundary >= input_len)
+        return 0;
+
+    for (size_t k = 0; k < SE_ARRAY_LEN(SE_ADDED_TOKENS); k++) {
+        const se_added_token_spec_t *token = &SE_ADDED_TOKENS[k];
+        if ((model->meta.added_token_mask & token->bit) == 0)
+            continue;
+        for (size_t back = 1; back < token->len && back <= boundary; back++) {
+            size_t start = boundary - back;
+            if (start + token->len > input_len || input[start] != '[')
+                continue;
+            if (memcmp(input + start, token->text, token->len) == 0)
+                return 1;
+        }
+    }
+    return 0;
+}
+
 size_t se_prefix_boundary_len(const se_model_t *model, const uint8_t *input, size_t input_len,
                               size_t target, size_t backscan) {
     if (target >= input_len)
@@ -420,14 +453,15 @@ size_t se_prefix_boundary_len(const se_model_t *model, const uint8_t *input, siz
             continue;
 
         if (cp_is_cjk_segment(model, cp)) {
-            if (after <= target)
+            if (after <= target && !boundary_splits_added_token(model, input, input_len, after))
                 return after;
-            if (pos > 0)
+            if (pos > 0 && !boundary_splits_added_token(model, input, input_len, pos))
                 return pos;
-            return 0;
+            continue;
         }
 
-        if ((is_whitespace(model, cp) || is_punct(model, cp)) && after <= target)
+        if ((is_whitespace(model, cp) || is_punct(model, cp)) && after <= target &&
+            !boundary_splits_added_token(model, input, input_len, after))
             return after;
     }
 
@@ -454,6 +488,7 @@ typedef struct {
     const se_model_t *model;
     se_scratch_t *scratch;
     uint32_t max_tokens;
+    se_token_limit_t limit_mode;
     size_t n_ids;
     size_t n_unk;
     size_t segment_len;
@@ -481,17 +516,44 @@ static int append_segment_cp(token_state_t *st, uint32_t cp) {
     return 1;
 }
 
+static size_t limited_token_count(const token_state_t *st) {
+    if (st->limit_mode == SE_TOKEN_LIMIT_USABLE && st->model->meta.unk_policy == SE_UNK_DROP)
+        return st->n_ids - st->n_unk;
+    return st->n_ids;
+}
+
 static int cap_after_append(token_state_t *st) {
-    if (st->max_tokens == 0 || st->n_ids <= (size_t)st->max_tokens)
+    if (st->max_tokens == 0 || limited_token_count(st) <= (size_t)st->max_tokens)
         return 0;
 
-    size_t dropped_unk = 0;
-    for (size_t k = st->max_tokens; k < st->n_ids; k++) {
-        if (st->scratch->ids[k] == st->model->meta.unk_id)
-            dropped_unk++;
+    if (st->limit_mode == SE_TOKEN_LIMIT_RAW || st->model->meta.unk_policy != SE_UNK_DROP) {
+        size_t dropped_unk = 0;
+        for (size_t k = st->max_tokens; k < st->n_ids; k++) {
+            if (st->scratch->ids[k] == st->model->meta.unk_id)
+                dropped_unk++;
+        }
+        st->n_unk -= dropped_unk;
+        st->n_ids = st->max_tokens;
+    } else {
+        size_t usable = 0;
+        size_t kept_unk = 0;
+        size_t keep = 0;
+        for (size_t k = 0; k < st->n_ids; k++) {
+            uint32_t id = st->scratch->ids[k];
+            if (id == st->model->meta.unk_id) {
+                kept_unk++;
+                continue;
+            }
+            usable++;
+            if (usable == (size_t)st->max_tokens) {
+                keep = k + 1;
+                break;
+            }
+        }
+        st->n_ids = keep;
+        st->n_unk = kept_unk;
     }
-    st->n_unk -= dropped_unk;
-    st->n_ids = st->max_tokens;
+
     st->stats->truncated = 1;
     return 1;
 }
@@ -602,6 +664,54 @@ static wordpiece_status_t wordpiece(const se_model_t *m, se_scratch_t *sc, const
     return 1;
 }
 
+static uint32_t added_token_id(const se_model_t *model, uint32_t bit) {
+    switch (bit) {
+    case SE_ADDED_PAD:
+        return model->meta.pad_id;
+    case SE_ADDED_UNK:
+        return model->meta.unk_id;
+    case SE_ADDED_CLS:
+        return model->meta.cls_id;
+    case SE_ADDED_SEP:
+        return model->meta.sep_id;
+    case SE_ADDED_MASK:
+        return model->meta.mask_id;
+    default:
+        return SE_SLOT_EMPTY;
+    }
+}
+
+static int match_added_token(const se_model_t *model, const uint8_t *input, size_t input_len,
+                             size_t pos, uint32_t *id_out, size_t *len_out) {
+    if (model->meta.added_token_mask == 0 || pos >= input_len || input[pos] != '[')
+        return 0;
+
+    for (size_t k = 0; k < SE_ARRAY_LEN(SE_ADDED_TOKENS); k++) {
+        const se_added_token_spec_t *token = &SE_ADDED_TOKENS[k];
+        if ((model->meta.added_token_mask & token->bit) == 0)
+            continue;
+        if (token->len > input_len - pos)
+            continue;
+        if (memcmp(input + pos, token->text, token->len) != 0)
+            continue;
+
+        *id_out = added_token_id(model, token->bit);
+        *len_out = token->len;
+        return 1;
+    }
+    return 0;
+}
+
+static se_status_t append_direct_id(token_state_t *st, uint32_t id, int *stop) {
+    if (!push_id(st->scratch, &st->n_ids, id))
+        return oom(st, "adding a special token");
+    if (id == st->model->meta.unk_id)
+        st->n_unk++;
+    if (cap_after_append(st))
+        *stop = 1;
+    return SE_OK;
+}
+
 static se_status_t append_wordpiece(token_state_t *st, const uint32_t *word, size_t word_len,
                                     int *stop) {
     wordpiece_status_t wp =
@@ -752,6 +862,31 @@ static se_status_t tokenize_ascii_run(token_state_t *st, const uint8_t *input, s
         if (b >= 0x80)
             break;
 
+        if (b == '[' && st->model->meta.added_token_mask != 0) {
+            uint32_t added_id = 0;
+            size_t added_len = 0;
+            if (match_added_token(st->model, input, input_len, i, &added_id, &added_len)) {
+                se_status_t rc = flush_segment(st, stop);
+                if (rc != SE_OK) {
+                    *ip = i;
+                    return rc;
+                }
+                if (*stop) {
+                    *ip = i;
+                    return SE_OK;
+                }
+                rc = append_direct_id(st, added_id, stop);
+                if (rc != SE_OK) {
+                    *ip = i;
+                    return rc;
+                }
+                i += added_len;
+                if (*stop)
+                    break;
+                continue;
+            }
+        }
+
         if (i >= next_cancel_check) {
             if (token_cancelled(st)) {
                 *ip = i;
@@ -811,8 +946,9 @@ static se_status_t tokenize_ascii_run(token_state_t *st, const uint8_t *input, s
 }
 
 se_status_t se_tokenize(const se_model_t *model, se_scratch_t *sc, const uint8_t *input,
-                        size_t input_len, uint32_t max_tokens, se_token_stats_t *stats,
-                        se_error_t *err, volatile sig_atomic_t *cancelled) {
+                        size_t input_len, uint32_t max_tokens, se_token_limit_t limit_mode,
+                        se_token_stats_t *stats, se_error_t *err,
+                        volatile sig_atomic_t *cancelled) {
     memset(stats, 0, sizeof(*stats));
 
     token_state_t st;
@@ -820,6 +956,7 @@ se_status_t se_tokenize(const se_model_t *model, se_scratch_t *sc, const uint8_t
     st.model = model;
     st.scratch = sc;
     st.max_tokens = max_tokens;
+    st.limit_mode = limit_mode;
     st.stats = stats;
     st.err = err;
     st.cancelled = cancelled;

data/ext/static_embeddings/se_unicode.c

@@ -115,6 +115,6 @@ const se_map_entry_t *se_map_lookup(const se_map_entry_t *entries, uint32_t coun
 int se_is_cjk(uint32_t cp) {
     return (cp >= 0x4E00 && cp <= 0x9FFF) || (cp >= 0x3400 && cp <= 0x4DBF) ||
            (cp >= 0x20000 && cp <= 0x2A6DF) || (cp >= 0x2A700 && cp <= 0x2B73F) ||
-           (cp >= 0x2B740 && cp <= 0x2B81F) || (cp >= 0x2B820 && cp <= 0x2CEAF) ||
+           (cp >= 0x2B740 && cp <= 0x2B81F) || (cp >= 0x2B920 && cp <= 0x2CEAF) ||
            (cp >= 0xF900 && cp <= 0xFAFF) || (cp >= 0x2F800 && cp <= 0x2FA1F);
 }

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;
@@ -1030,6 +1031,10 @@ 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;
 }
@@ -1049,7 +1054,7 @@ typedef struct {
 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,
-                          &job->stats, &job->err, NULL);
+                          SE_TOKEN_LIMIT_RAW, &job->stats, &job->err, NULL);
     if (job->rc != SE_OK)
         return Qnil;
 
@@ -1237,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;
     }
@@ -1297,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;
 }
@@ -1607,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,