static_embeddings 1.5.8 → 1.5.9

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: cbb2b16d03e01ba91c8c856890e98537ff81c4b4961a4d4172dc2e8a27fb46e3
-  data.tar.gz: d182451276da091e5410fa5dcbf93ec08e773bfa51cc6226f0c244db132e3656
+  metadata.gz: 2de25b005846ea72148d725e56b30e25e1f0a3a8aace72cdada2865e0ec625f2
+  data.tar.gz: aeae810542ca9734a9fc6e1edf700950e856e054c6f2b1595d5553d5756ee2e4
 SHA512:
-  metadata.gz: 54d96917d4a2fd738787ee12bd57a63f3793893a4abe3e59a9b860efe8508a895cb4d1cbd865f4f5692af1b2aa46eea21b17b55ed090b98a9af8df59946feadc
-  data.tar.gz: da73a4ead8792691b1e8433ecb7487f121e1dda88c61054e56f0f05b56c43f922a49643f1668da4c34be1258f654f134975cae5db875075f2a99e9bd2ea5256a
+  metadata.gz: 9747f27eeaaec7272099bf4789afacda8b76017850aa4cc4d16a596bcb3b43d9b23e2686cb9a6b3a9431bfc88468a605a903ad9f9fd0eac6e1019746da69be58
+  data.tar.gz: d112d3b2381929555dad88e0b85b16a759a0ba0ddcf1b64cd8c3c44b90c122a9f56b616d6af8d8560d2ec6408bc61ee8f83340af2bf00435fda9f14f9a9d70b6

data/CHANGELOG.md

@@ -1,5 +1,64 @@
 # Changelog
 
+## 1.5.9
+
+Pins the upstream contract to **model2vec 0.10.0 / tokenizers 0.23.3** and
+closes a fail-open hole in the converter. The `.semb` format stays v3; files
+written by earlier versions load and embed exactly as before.
+
+### Fixed
+
+- **Model2Vec `weights` were silently ignored.** The importer picked the only 2-D
+  tensor in `model.safetensors` and skipped every 1-D tensor, so a model with
+  per-token `weights` converted without error and produced wrong vectors
+  (cosine 0.71–0.98 against `StaticModel` on the test fixture). Vocabulary-
+  quantized models (`embeddings` K×d + `weights` + `mapping`) failed with a
+  misleading "rows != tokens" error. Both are now baked into ordinary float32
+  rows, `row[i] = embeddings[mapping[i]] * weights[i]`, which is exactly what
+  `StaticModel` pools; the C runtime is unchanged. Weighted, quantized and
+  F64/I64 side-tensor variants match model2vec 0.10.0 within `max_abs 1.2e-7`.
+  Provenance records `baked_token_weights` / `baked_vocabulary_quantization`.
+- **Unknown safetensors tensors are rejected** instead of ignored.
+- **`tools/eval_retrieval.rb` no longer reports trivial metrics as results.**
+  With `k >= documents` Hit@k is 1.0 for every model by construction (the
+  bundled Russian FAQ set has 10 documents and used k=10). The tool now prints
+  full-ranking MRR, nDCG/recall/hit at k, a random-ranking baseline, bootstrap
+  95% intervals over queries, and warnings for trivial `k`, small query sets and
+  zero-vector queries.
+
+### Changed
+
+- **Model2Vec truncation follows 0.10.0 by default.** model2vec 0.10.0 moved
+  truncation into the tokenizer: `max_length` now counts raw ids including
+  `[UNK]`, and `[UNK]` is dropped afterwards. 0.9.0 dropped `[UNK]` first. On a
+  text whose `[UNK]` tokens precede the cap the two versions disagree (cosine
+  0.59 on a synthetic case). New header value `truncation_policy = 3`
+  (`raw_ids`) is the Model2Vec default; `--truncation usable-ids` (`2`) keeps the
+  0.9.0 contract for indexes built with it. Files written before 1.5.9 carry
+  `2` and are unaffected. `Model#truncation` reports `:raw_ids` / `:usable_ids`.
+  Sentence Transformers conversions are unchanged (with `UNK_INCLUDE` both
+  policies select the same ids).
+- **`config.json` `max_length` sets the default token cap**, as in 0.10.0
+  (`null` means unlimited; absent still means 512).
+- Oracle schema 3 (`tools/model2vec_oracle.py`) records the truncation policy
+  and supports both upstream versions; the checker refuses to compare a model
+  against an oracle of the other policy.
+- CI: `upstream_parity` and `potion_audit` pin 0.10.0 / 0.23.3 / numpy 2.5.4; new
+  `upstream_parity_usable_ids` keeps the 0.9.0 contract tested; new
+  non-blocking `upstream_canary` runs the corpus against unpinned PyPI releases
+  so the next upstream drift is reported instead of discovered in production.
+- README "Language safety" no longer presents the `[UNK]` ratio as sufficient:
+  text spelled out letter by letter has almost no `[UNK]` and is just as broken.
+
+### Upgrade notes
+
+Reconverting a Model2Vec source with 1.5.9 changes vectors only for texts with
+`[UNK]` before the token cap. If you keep an existing index, either keep the
+old `.semb` or reconvert with `--truncation usable-ids`; do not mix files of
+both policies in one index. Real-model audits in `docs/MODEL_AUDIT.md` were
+recorded under 0.9.0 and must be re-run (the `potion_audit` CI job does this
+for potion-retrieval-32M).
+
 ## 1.5.8
 
 ### Fixed

data/README.md

@@ -33,7 +33,7 @@ BERT WordPiece tokenizer profile, an mmap-ready vocabulary lookup, float32
 embedding rows, and provenance plus a checksum.
 
 The only supported tokenizer profile is `BERT_WORDPIECE_V1`, pinned for
-compatibility to `tokenizers 0.23.1`. The converter rejects unsupported
+compatibility to `tokenizers 0.23.3`. The converter rejects unsupported
 tokenizer features rather than approximating them. `.semb` format v3 also
 records which of the five standard BERT added tokens are active. Files produced
 by 0.1.4 and earlier use format v2 and must be reconverted.
@@ -60,6 +60,20 @@ bundle exec ruby -Ilib exe/static_embeddings convert ./static-retrieval-mrl-en-v
   --trained-mrl-dims 1024,512,256,128,64,32
 ```
 
+Model2Vec conversions follow the pinned `model2vec 0.10.0` contract by default:
+the token cap counts raw tokenizer ids (with `[UNK]`), then `[UNK]` is dropped.
+An index built with model2vec 0.9.0, or with this gem before 1.5.9, used the
+opposite order; keep it consistent with `--truncation usable-ids`. The policy is
+recorded in the file (`model.truncation` is `:raw_ids` or `:usable_ids`), so old
+`.semb` files keep their behaviour without reconversion. A `max_length` in the
+source `config.json` (0.10.0 writes it; `null` means unlimited) becomes the
+default `max_tokens`.
+
+Model2Vec side tensors are baked into the rows at convert time: per-token
+`weights` and the vocabulary-quantization `mapping` produce
+`row[i] = embeddings[mapping[i]] * weights[i]`, exactly what `StaticModel`
+pools. Any other tensor in `model.safetensors` is rejected rather than ignored.
+
 `--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`.
@@ -139,18 +153,20 @@ Output is a binary `String` of little-endian float32 values in row-major order.
 exist for debugging and application code, not for the hot path.
 
 `tokenize` exposes raw WordPiece ids, including `[UNK]`, and its own
-`max_tokens` cap is therefore a raw-id cap. Embedding has a different, upstream
-pooling rule: for `UNK_DROP`, `[UNK]` is removed first and `max_tokens` is then
-applied to usable ids. To cache tokenization without changing the embedding,
-cache it unbounded:
+`max_tokens` cap is therefore a raw-id cap. Embedding applies the model's
+recorded truncation policy: with `:raw_ids` (model2vec 0.10.0, the default for
+new Model2Vec conversions) the cap counts raw ids and `[UNK]` is dropped
+afterwards; with `:usable_ids` (model2vec 0.9.0) `[UNK]` is dropped first and the
+cap counts the remaining ids. To cache tokenization without changing the
+embedding, cache it unbounded:
 
 ```ruby
 ids = model.tokenize(text, max_tokens: false)
 model.embed_token_ids(ids) == model.embed(text)
 ```
 
-`embed_token_ids` applies that same usable-id cap. This distinction matters only
-when `[UNK]` appears around the truncation boundary.
+`embed_token_ids` applies the same policy. The two policies differ only when
+`[UNK]` appears before the truncation boundary.
 
 ### `format: :f16`
 
@@ -207,11 +223,11 @@ model.dot_top_k(query, MATRIX, 10)          # lock-free, concurrent, GVL release
 
 When truncation is active the runtime starts with a leading slice sized from
 `max_tokens`, cut on a word boundary, and grows the budget until it has reached
-`max_tokens` **usable** ids. For ordinary in-vocabulary text this keeps the
-tokenized prefix to a few kilobytes even when the document is megabytes long.
-For OOV-heavy input it may scan farther — or the whole string — because `[UNK]`
-no longer consumes the usable-token budget. That is required for the corrected
-UNK-before-truncate contract.
+`max_tokens` ids as counted by the model's truncation policy. For ordinary
+in-vocabulary text this keeps the tokenized prefix to a few kilobytes even when
+the document is megabytes long. Under `:usable_ids`, OOV-heavy input may scan
+farther — or the whole string — because `[UNK]` does not consume the budget;
+under `:raw_ids` it does, so the scan stays bounded.
 
 That bounds the tokenizer, not the whole call. `embed` also has to establish
 that the Ruby `String` is valid UTF-8, and when Ruby has not computed the
@@ -286,35 +302,49 @@ end
 
 ## Language safety
 
-`potion-retrieval-32M` is an English retrieval model. Russian text can produce a
-high `[UNK]` ratio while still returning a valid vector, which is the most
-likely way to end up with a silently bad index. Check the ratio before trusting
-a corpus:
+`potion-retrieval-32M` is an English retrieval model. Foreign-language text still
+returns a valid vector, which is the most likely way to end up with a silently
+bad index. It fails in two different ways, and the `[UNK]` ratio sees only one:
+
+- words the vocabulary cannot spell become `[UNK]` (high `unk_count`);
+- words it *can* spell letter by letter (English BERT-derived vocabularies
+  usually contain single Cyrillic letters) become chains of one-character
+  pieces: `[UNK]` stays near zero and the vector is still close to meaningless.
+
+Check both before trusting a corpus:
 
 ```ruby
 stats = model.embed_with_stats(text)
 ratio = stats[:token_count].zero? ? 0.0 : stats[:unk_count].to_f / stats[:token_count]
+words = text.split.size
+pieces_per_word = words.zero? ? 0.0 : stats[:token_count].to_f / words
 
 warn "high [UNK] ratio: #{ratio.round(3)}" if ratio > 0.3
+warn "text is spelled out in fragments: #{pieces_per_word.round(1)} pieces/word" if pieces_per_word > 2.5
 ```
 
+In-vocabulary English is typically 1.1–1.4 pieces per word. `text.split` is a
+rough word count (it ignores punctuation splitting), which is enough for this
+check.
+
 Out-of-vocabulary text is also slower, because every word falls through to
 subword splitting instead of hitting the vocabulary directly.
 
 ## Correctness contract
 
-0.1.5 separates contracts that 0.1.4 accidentally mixed together:
+Pinned upstream: `model2vec 0.10.0`, `tokenizers 0.23.3`.
 
-- `tokenize` matches the pinned Hugging Face `tokenizers 0.23.1` raw
+- `tokenize` matches the pinned Hugging Face `tokenizers` raw
   `BertNormalizer + BertPreTokenizer + WordPiece` ids, including the supported
   `normalized: false` standard added tokens and `[UNK]`.
-- Embedding drops `[UNK]` when the model records `UNK_DROP`, then applies
-  `max_tokens` to the remaining usable ids, mean-pools, and performs source-model
-  L2 normalization.
-- The gem intentionally does **not** reproduce Model2Vec 0.9.0's character
-  pre-cut (`max_length * median_token_length`). Here `max_tokens` means actual
-  usable tokenizer ids. The oracle reports rows affected by this difference
-  separately instead of calling them parity.
+- Embedding applies the recorded truncation policy (`:raw_ids` caps raw ids and
+  then drops `[UNK]`, as model2vec 0.10.0 does; `:usable_ids` drops `[UNK]` and
+  then caps, as 0.9.0 did), mean-pools, and performs source-model L2
+  normalization.
+- The gem intentionally does **not** reproduce Model2Vec's character pre-cut
+  (`max_length * median_token_length`, still present in 0.10.0). Here
+  `max_tokens` means actual tokenizer ids. The oracle reports rows affected by
+  this difference separately instead of calling them parity.
 
 Where the execution contracts are the same, vectors are compared to
 `model2vec.StaticModel` with:

data/Rakefile

@@ -44,7 +44,7 @@ task parity: :compile do
       Generate it first (needs Python and the model2vec package):
 
         python3 -m venv .venv-model2vec && . .venv-model2vec/bin/activate
-        pip install "model2vec==0.9.0" "tokenizers==0.23.1" numpy
+        pip install "model2vec==0.10.0" "tokenizers==0.23.3" "numpy==2.5.4"
         python tools/model2vec_oracle.py /path/to/source-model --out #{oracle}
     MSG
   end

data/docs/ARCHITECTURE.md

@@ -89,7 +89,7 @@ already has a slot for I8; there is no dead I8 path in C.
 ### Why the semantics live in the header
 
 `max_tokens_default`, `unk_policy`, `empty_policy` and the normalizer flags
-are properties of the *model*, not options the caller guesses. The embedding contract caps the first 512 usable token ids by default;
+are properties of the *model*, not options the caller guesses. So is `truncation_policy`: `3` caps raw ids and then drops `[UNK]` (model2vec 0.10.0, the default for new Model2Vec conversions), `2` drops `[UNK]` first and caps the rest (model2vec 0.9.0, every file written before 1.5.9). The embedding contract caps the first 512 token ids by default;
 a model card's `model_max_length` of 1,000,000 is a different number entirely, and
 confusing them silently changes every vector for long chunks. Model2Vec 0.9.0
 also applies a character pre-cut before tokenization; this runtime intentionally
@@ -159,7 +159,7 @@ picked the wrong one would produce plausible, slightly wrong vectors forever.
 The Ruby build that generated the tables is stamped into provenance, so a
 model file and its embedded normalizer tables can never drift apart unnoticed.
 That stamp is not, by itself, proof of Hugging Face parity: 0.1.5 therefore
-checks a systematic boundary corpus against pinned `tokenizers 0.23.1` in CI.
+checks a systematic boundary corpus against pinned `tokenizers` (0.23.3 since 1.5.9) in CI.
 One subtle example is `Cn`: `bert.rs` says `is_other()` includes unassigned
 codepoints, while the `unicode_categories 0.1.1` implementation actually tests
 Cc/Cf/Co. The runtime follows the executable upstream behavior, not the comment.
@@ -168,9 +168,10 @@ Cc/Cf/Co. The runtime follows the executable upstream behavior, not the comment.
 
 ```
 tokenize -> raw ids (including [UNK])
-drop [UNK] if unk_policy = DROP
-truncate usable ids to max_tokens
-acc[j] += row[j] for each usable id, in token order
+policy 3 (raw_ids):    keep the first max_tokens raw ids, then drop [UNK]*
+policy 2 (usable_ids): drop [UNK]* first, then keep the first max_tokens ids
+                       (* only when unk_policy = DROP; both coincide otherwise)
+acc[j] += row[j] for each kept id, in token order
 divide by the number of pooled rows
 L2 normalize if the model says so
 ```
@@ -226,10 +227,11 @@ therefore serialised the process on work that the tokenizer would discard after
 `embed` and `embed_batch` instead copy a leading slice. The budget starts at
 `max_tokens * 16` bytes clamped to `[4096, 65536]`, and the slice is cut at a
 position where the tokenizer would have started a fresh word. A result is accepted only when that slice actually reached `max_tokens`
-usable ids; `[UNK]` does not consume the embedding budget. Otherwise the budget
-doubles and the text is tokenized again. Growth is geometric. OOV-heavy input
-may therefore reach the full document: deciding that fewer than `max_tokens`
-usable ids exist requires reading it.
+ids as counted by the truncation policy. Under `usable_ids`, `[UNK]` does not
+consume the embedding budget, so OOV-heavy input may reach the full document:
+deciding that fewer than `max_tokens` usable ids exist requires reading it.
+Under `raw_ids` every id counts and the window stays bounded. Otherwise the
+budget doubles and the text is tokenized again. Growth is geometric.
 
 Three properties make this exact rather than approximate:
 

data/docs/LIMITATIONS.md

@@ -23,6 +23,14 @@ the entry says whether that is a decision or just unfinished work.
   context; a word contributes the same row wherever it appears. If you need
   contextual embeddings, this is the wrong tool.
 
+- **Two upstream truncation contracts, chosen at convert time.** `.semb` files
+  record `raw_ids` (model2vec 0.10.0) or `usable_ids` (0.9.0); there is no
+  runtime switch, because mixing them inside one index silently changes
+  vectors. Model2Vec's character pre-cut is reproduced by neither.
+- **Model2Vec `weights`/`mapping` are expanded, not stored.** A
+  vocabulary-quantized source becomes a full-vocabulary `.semb`; the file does
+  not get smaller. int8 embedding tensors are rejected.
+
 ## Runtime
 
 - **No internal parallelism.** `threads:` is rejected rather than ignored. Split

data/docs/MODEL_AUDIT.md

@@ -229,3 +229,44 @@ arm64-darwin24, 1.5.6):
 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.
+
+## 1.5.9 record: model2vec 0.10.0 contract
+
+Recorded locally (x86_64-linux, Ruby 3.2.3, Python 3.12.3), deterministic
+`tiny-wordpiece` fixture, 434-row corpus from `tools/parity_cases.py`, oracle
+schema 3.
+
+`model2vec 0.10.0` / `tokenizers 0.23.3` / `numpy 2.5.4`, default conversion
+(`truncation=raw_ids`):
+
+```text
+rows=434
+vectors_checked=432
+intentional_character_pretruncate_deviations=2
+min_cosine=0.9999999999999982
+max_abs_all=1.1920928955078125e-07
+raw_token_id_failures=[] usable_token_id_failures=[] embed_invariant_failures=[] vector_failures=[]
+corpus parity OK (434/434)
+```
+
+`model2vec 0.9.0` / `tokenizers 0.23.3`, `--truncation usable-ids`
+(`--allow-version-mismatch`, since the oracle now pins 0.10.0):
+
+```text
+rows=434 vectors_checked=429 intentional_character_pretruncate_deviations=5
+min_cosine=0.9999999999999988 max_abs_all=5.960464477539063e-08
+corpus parity OK (434/434)
+```
+
+Negative control: the same `usable_ids` file against the 0.10.0 oracle, with the
+checker's policy guard disabled, fails on rows 426–429 and 432, so the corpus
+does detect the truncation-order change.
+
+Side tensors (`weights`, `weights + mapping` from `quantize_model(...,
+vocabulary_quantization=32)`, and an F64/I64 variant) on the same fixture:
+`max_abs` 6.0e-8, 6.0e-8 and 1.2e-7 against `StaticModel.encode`.
+
+**Not yet re-run under 0.10.0:** potion-retrieval-32M, potion-base-8M,
+potion-science-32M. Their records above are 0.9.0 / `usable_ids` results. The
+`potion_audit` CI job regenerates the potion-retrieval-32M record; append it
+here before trusting a 1.5.9 conversion of that model.

data/ext/static_embeddings/se_embed.c

@@ -197,8 +197,8 @@ se_status_t se_embed_ids(const se_model_t *model, se_scratch_t *sc, const uint32
 se_status_t se_embed_one(const se_model_t *model, se_scratch_t *sc, const uint8_t *input,
                          size_t input_len, uint32_t max_tokens, float *out, se_token_stats_t *stats,
                          se_error_t *err, se_cancel_flag_t *cancelled) {
-    se_status_t rc = se_tokenize(model, sc, input, input_len, max_tokens, SE_TOKEN_LIMIT_USABLE,
-                                 stats, err, cancelled);
+    se_status_t rc = se_tokenize(model, sc, input, input_len, max_tokens,
+                                 se_embed_limit_mode(model), stats, err, cancelled);
     if (rc != SE_OK)
         return rc;
     return embed_ids_core(model, sc, sc->ids, stats->token_count, out, err, cancelled);

data/ext/static_embeddings/se_format.c

@@ -695,7 +695,8 @@ 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_USABLE_IDS_BEFORE_POOLING) {
+    if (m->truncation_policy != SE_TRUNCATE_USABLE_IDS_BEFORE_POOLING &&
+        m->truncation_policy != SE_TRUNCATE_RAW_IDS_BEFORE_UNK_DROP) {
         se_error_set(err, SE_ERR_INVALID_FORMAT, "truncation policy %u is not supported",
                      m->truncation_policy);
         return SE_ERR_INVALID_FORMAT;

data/ext/static_embeddings/se_internal.h

@@ -28,6 +28,7 @@ static inline void se_cancel_request(se_cancel_flag_t *flag) {
 #define SE_NORMALIZATION_L2            1u
 
 #define SE_TRUNCATE_USABLE_IDS_BEFORE_POOLING 2u
+#define SE_TRUNCATE_RAW_IDS_BEFORE_UNK_DROP   3u
 
 #define SE_UNK_INCLUDE 0u
 #define SE_UNK_DROP    1u
@@ -422,6 +423,12 @@ typedef struct {
 
 typedef enum { SE_TOKEN_LIMIT_RAW = 0, SE_TOKEN_LIMIT_USABLE = 1 } se_token_limit_t;
 
+static inline se_token_limit_t se_embed_limit_mode(const se_model_t *model) {
+    return model->meta.truncation_policy == SE_TRUNCATE_RAW_IDS_BEFORE_UNK_DROP
+               ? SE_TOKEN_LIMIT_RAW
+               : SE_TOKEN_LIMIT_USABLE;
+}
+
 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_limit_t limit_mode,
                         se_token_stats_t *stats, se_error_t *err, se_cancel_flag_t *cancelled);

data/ext/static_embeddings/static_embeddings.c

@@ -463,6 +463,12 @@ static VALUE model_max_tokens(VALUE self) {
     return UINT2NUM(get_model(self)->model.meta.max_tokens_default);
 }
 
+static VALUE model_truncation(VALUE self) {
+    return se_embed_limit_mode(&get_model(self)->model) == SE_TOKEN_LIMIT_RAW
+               ? ID2SYM(rb_intern("raw_ids"))
+               : ID2SYM(rb_intern("usable_ids"));
+}
+
 static VALUE model_normalized_p(VALUE self) {
     return get_model(self)->model.meta.normalization_type == SE_NORMALIZATION_L2 ? Qtrue : Qfalse;
 }
@@ -1290,7 +1296,8 @@ 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 && w->model.meta.unk_policy == SE_UNK_DROP) {
+    if (max_tokens != 0 && w->model.meta.unk_policy == SE_UNK_DROP &&
+        se_embed_limit_mode(&w->model) == SE_TOKEN_LIMIT_USABLE) {
         size_t usable = 0;
         size_t cutoff = n;
         for (size_t i = 0; i < n; i++) {
@@ -1710,6 +1717,7 @@ RUBY_FUNC_EXPORTED void Init_static_embeddings(void) {
     rb_define_method(cModel, "dim", model_dim, 0);
     rb_define_method(cModel, "vocab_size", model_vocab_size, 0);
     rb_define_method(cModel, "max_tokens", model_max_tokens, 0);
+    rb_define_method(cModel, "truncation", model_truncation, 0);
     rb_define_method(cModel, "normalized?", model_normalized_p, 0);
     rb_define_method(cModel, "lowercase?", model_lowercase_p, 0);
     rb_define_method(cModel, "unk_id", model_unk_id, 0);

data/lib/static_embeddings/canonical.rb

@@ -2,9 +2,9 @@ module StaticEmbeddings
   module Canonical
     Dimensions = Struct.new(:native, :output, :trained, keyword_init: true)
     Runtime = Struct.new(:normalization, :unk_policy, :empty_policy, :max_tokens, :add_special_tokens,
-                         keyword_init: true)
+                         :truncation, keyword_init: true)
     Source = Struct.new(:family, :model, :revision, :oracle, :files_sha256, :tokenizer_class,
-                        :config_seq_length, keyword_init: true)
+                        :config_seq_length, :baked, keyword_init: true)
     Model = Struct.new(:tokens, :matrix, :dimensions, :runtime, :tokenizer, :source, keyword_init: true)
 
     module_function
@@ -13,18 +13,20 @@ module StaticEmbeddings
       Dimensions.new(native: native, output: output, trained: trained&.freeze).freeze
     end
 
-    def runtime(normalization:, unk_policy:, empty_policy:, max_tokens:, add_special_tokens: false)
+    def runtime(normalization:, unk_policy:, empty_policy:, max_tokens:, add_special_tokens: false,
+                truncation: Format::TRUNCATE_USABLE_IDS_BEFORE_POOLING)
       Runtime.new(
         normalization: normalization,
         unk_policy: unk_policy,
         empty_policy: empty_policy,
         max_tokens: max_tokens,
-        add_special_tokens: add_special_tokens
+        add_special_tokens: add_special_tokens,
+        truncation: truncation
       ).freeze
     end
 
     def source(family:, model:, oracle:, files_sha256:, revision: nil, tokenizer_class: nil,
-               config_seq_length: nil)
+               config_seq_length: nil, baked: nil)
       Source.new(
         family: family,
         model: model,
@@ -32,7 +34,8 @@ module StaticEmbeddings
         oracle: oracle,
         files_sha256: files_sha256.freeze,
         tokenizer_class: tokenizer_class,
-        config_seq_length: config_seq_length
+        config_seq_length: config_seq_length,
+        baked: baked&.freeze
       ).freeze
     end
 

data/lib/static_embeddings/cli.rb

@@ -26,6 +26,9 @@ module StaticEmbeddings
           --max-tokens N      Model2Vec default: 512; Sentence Transformers default: unlimited
           --max-tokens unlimited
           --dimensions N      Keep the first N dims (Matryoshka prefix slice)
+          --truncation raw-ids|usable-ids
+                              Model2Vec default raw-ids (model2vec 0.10.0);
+                              usable-ids reproduces model2vec 0.9.0
           --revision SHA      Source revision recorded in provenance
           --trained-mrl-dims  1024,512,256,... training dims recorded in provenance
 
@@ -85,12 +88,13 @@ module StaticEmbeddings
           parser.on("--dimensions N", Integer) { |value| options[:dimensions] = value }
           parser.on("--revision SHA") { |value| options[:source_revision] = value }
           parser.on("--trained-mrl-dims LIST") { |value| options[:trained_mrl_dims] = value }
+          parser.on("--truncation POLICY", %w[raw-ids usable-ids]) { |value| options[:truncation] = value }
         end.parse!(argv)
       end
     end
 
     def conversion_options(options)
-      options.select { |key, _| %i[max_tokens dimensions source_revision trained_mrl_dims].include?(key) }
+      options.select { |key, _| %i[max_tokens dimensions source_revision trained_mrl_dims truncation].include?(key) }
     end
 
     def parse_max_tokens(value)
@@ -111,6 +115,7 @@ module StaticEmbeddings
         "  bytes      #{report[:bytes]}",
         "  sha256     #{report[:sha256]}",
         "  max_tokens #{max_tokens}",
+        "  truncation #{report[:truncation]}",
         "",
         "Record this conversion in docs/MODEL_AUDIT.md before trusting the vectors."
       ].join("\n")
@@ -143,6 +148,7 @@ module StaticEmbeddings
         "dim" => model.dim,
         "vocab_size" => model.vocab_size,
         "max_tokens" => model.max_tokens,
+        "truncation" => model.truncation,
         "normalized" => model.normalized?,
         "lowercase" => model.lowercase?,
         "unk_id" => model.unk_id,

data/lib/static_embeddings/conversion.rb

@@ -25,6 +25,7 @@ module StaticEmbeddings
         dim: model.dimensions.output,
         native_dim: model.dimensions.native,
         max_tokens: model.runtime.max_tokens,
+        truncation: Format::TRUNCATION_POLICIES.key(model.runtime.truncation),
         provenance: JSON.parse(provenance)
       )
     end
@@ -36,6 +37,7 @@ module StaticEmbeddings
         dim: model.dimensions.output,
         normalization_type: runtime.normalization,
         max_tokens_default: runtime.max_tokens,
+        truncation_policy: runtime.truncation,
         add_special_tokens: runtime.add_special_tokens,
         unk_policy: runtime.unk_policy,
         empty_policy: runtime.empty_policy,

data/lib/static_embeddings/format/constants.rb

@@ -12,6 +12,11 @@ module StaticEmbeddings
     NORMALIZATION_NONE = 0
     NORMALIZATION_L2 = 1
     TRUNCATE_USABLE_IDS_BEFORE_POOLING = 2
+    TRUNCATE_RAW_IDS_BEFORE_UNK_DROP = 3
+    TRUNCATION_POLICIES = {
+      raw_ids: TRUNCATE_RAW_IDS_BEFORE_UNK_DROP,
+      usable_ids: TRUNCATE_USABLE_IDS_BEFORE_POOLING
+    }.freeze
     UNK_INCLUDE = 0
     UNK_DROP = 1
     EMPTY_ZERO_VECTOR = 0
@@ -53,7 +58,6 @@ module StaticEmbeddings
       28 => TOKENIZER_BERT_WORDPIECE_V1,
       32 => DTYPE_F32,
       36 => POOLING_MEAN,
-      48 => TRUNCATE_USABLE_IDS_BEFORE_POOLING,
       108 => HASH_SEED
     }.freeze
 
@@ -61,6 +65,7 @@ module StaticEmbeddings
       20 => :dim,
       40 => :normalization_type,
       44 => :max_tokens_default,
+      48 => :truncation_policy,
       56 => :unk_policy,
       60 => :empty_policy,
       80 => :max_input_chars_per_word,

data/lib/static_embeddings/importers/model2vec.rb

@@ -3,18 +3,20 @@ module StaticEmbeddings
     module Model2Vec
       SOURCE_FILES = %w[tokenizer.json config.json tokenizer_config.json model.safetensors].freeze
       DEFAULT_MAX_TOKENS = 512
+      DEFAULT_TRUNCATION = :raw_ids
       FAMILY = "model2vec"
       ORACLE = "model2vec.StaticModel"
 
       module_function
 
       def call(root, model_id: nil, max_tokens: nil, dimensions: nil,
-               source_revision: nil, trained_mrl_dims: nil)
+               source_revision: nil, trained_mrl_dims: nil, truncation: nil)
         tokenizer = Support.load_object(File.join(root, "tokenizer.json"))
         config = Support.load_object(File.join(root, "config.json"), optional: true) || {}
         tokenizer_config = Support.load_object(File.join(root, "tokenizer_config.json"), optional: true) || {}
         profile, tokens = BertWordPiece.compile(tokenizer, tokenizer_config)
-        payload, native_dim = Support.extract_matrix(File.join(root, "model.safetensors"), tokens.length)
+        payload, native_dim, baked = Support.extract_matrix(File.join(root, "model.safetensors"), tokens.length,
+                                                            side_tensors: Support::MODEL2VEC_SIDE_TENSORS)
         output_dim = Support.resolve_dimensions(dimensions, native_dim)
 
         Canonical.model(
@@ -29,7 +31,8 @@ module StaticEmbeddings
             normalization: normalization(config),
             unk_policy: Format::UNK_DROP,
             empty_policy: Format::EMPTY_ZERO_VECTOR,
-            max_tokens: Support.resolve_max_tokens(max_tokens, DEFAULT_MAX_TOKENS)
+            max_tokens: Support.resolve_max_tokens(max_tokens, default_max_tokens(config)),
+            truncation: Support.resolve_truncation(truncation, DEFAULT_TRUNCATION)
           ),
           tokenizer: BertWordPiece.runtime_meta(profile, tokens),
           source: Canonical.source(
@@ -39,11 +42,24 @@ module StaticEmbeddings
             oracle: ORACLE,
             files_sha256: Support.digest_files(SOURCE_FILES.map { |name| File.join(root, name) }),
             tokenizer_class: profile[:tokenizer_class],
-            config_seq_length: config["seq_length"]
+            config_seq_length: config["seq_length"],
+            baked: baked
           )
         )
       end
 
+      def default_max_tokens(config)
+        return DEFAULT_MAX_TOKENS unless config.key?("max_length")
+
+        value = config["max_length"]
+        return 0 if value.nil?
+        unless value.is_a?(Integer) && value.between?(1, Format::UINT32_MAX)
+          raise InvalidSourceError,
+                "config.json max_length must be an integer in 1..#{Format::UINT32_MAX} or null, got #{value.inspect}"
+        end
+        value
+      end
+
       def normalization(config)
         config.fetch("normalize", false) ? Format::NORMALIZATION_L2 : Format::NORMALIZATION_NONE
       end

data/lib/static_embeddings/importers/sentence_transformers_static.rb

@@ -11,7 +11,7 @@ module StaticEmbeddings
       module_function
 
       def call(root, model_id: nil, max_tokens: nil, dimensions: nil,
-               source_revision: nil, trained_mrl_dims: nil)
+               source_revision: nil, trained_mrl_dims: nil, truncation: nil)
         module_dir = module_directory(root)
         tokenizer = Support.load_object(File.join(module_dir, "tokenizer.json"))
         tokenizer_config = Support.load_object(File.join(module_dir, "tokenizer_config.json"), optional: true) || {}
@@ -32,7 +32,8 @@ module StaticEmbeddings
             normalization: Format::NORMALIZATION_NONE,
             unk_policy: Format::UNK_INCLUDE,
             empty_policy: Format::EMPTY_ZERO_VECTOR,
-            max_tokens: Support.resolve_max_tokens(max_tokens, DEFAULT_MAX_TOKENS)
+            max_tokens: Support.resolve_max_tokens(max_tokens, DEFAULT_MAX_TOKENS),
+            truncation: Support.resolve_truncation(truncation, :usable_ids)
           ),
           tokenizer: BertWordPiece.runtime_meta(profile, tokens),
           source: Canonical.source(

data/lib/static_embeddings/importers/support.rb

@@ -4,6 +4,7 @@ require "static_embeddings/errors"
 require "static_embeddings/format/constants"
 require "static_embeddings/safetensors"
 require "static_embeddings/row_prefix_payload"
+require "static_embeddings/weighted_payload"
 
 module StaticEmbeddings
   module Importers
@@ -35,20 +36,65 @@ module StaticEmbeddings
         end
       end
 
-      def extract_matrix(path, vocab_size)
+      MODEL2VEC_SIDE_TENSORS = %w[weights mapping].freeze
+
+      def extract_matrix(path, vocab_size, side_tensors: [])
         raise InvalidSourceError, "missing model.safetensors in #{File.dirname(path)}" unless File.file?(path)
 
-        name, tensor = sole_matrix_tensor(Safetensors.describe(path).fetch(:tensors))
+        tensors = Safetensors.describe(path).fetch(:tensors)
+        name, tensor = sole_matrix_tensor(tensors)
+        reject_unknown_tensors(tensors, name, side_tensors)
         rows, dim = tensor.fetch(:shape)
-        unless rows == vocab_size
+        unless dim.positive? && dim <= Format::UINT32_MAX
+          raise ConversionError, "embedding matrix #{name} has unsupported dimension #{dim}"
+        end
+
+        payload = Safetensors.f32_payload(path, tensor)
+        weights = side_vector(path, tensors, "weights", vocab_size, side_tensors)
+        mapping = side_vector(path, tensors, "mapping", vocab_size, side_tensors)
+
+        if mapping
+          validate_mapping(mapping, rows)
+        elsif rows != vocab_size
           raise ConversionError,
                 "embedding matrix #{name} has #{rows} rows but the tokenizer has #{vocab_size} tokens"
         end
-        unless dim.positive? && dim <= Format::UINT32_MAX
-          raise ConversionError, "embedding matrix #{name} has unsupported dimension #{dim}"
+        validate_weights(weights) if weights
+
+        if weights || mapping
+          payload = WeightedPayload.new(payload, dim: dim, source_rows: rows, output_rows: vocab_size,
+                                                 weights: weights, mapping: mapping)
+        end
+        [payload, dim, { weights: !weights.nil?, mapping_clusters: mapping ? rows : nil }]
+      end
+
+      def reject_unknown_tensors(tensors, matrix_name, side_tensors)
+        unknown = tensors.keys - [matrix_name] - side_tensors
+        return if unknown.empty?
+
+        raise ConversionError,
+              "model.safetensors contains tensors this converter does not understand: #{unknown.inspect}; " \
+              "refusing to ignore them because that would silently change vectors"
+      end
+
+      def side_vector(path, tensors, key, vocab_size, side_tensors)
+        return nil unless side_tensors.include?(key) && tensors.key?(key)
+
+        tensor = tensors.fetch(key)
+        unless tensor.fetch(:shape) == [vocab_size]
+          raise ConversionError, "#{key} tensor has shape #{tensor[:shape].inspect}, expected [#{vocab_size}]"
+        end
+        Safetensors.read_vector(path, tensor)
+      end
+
+      def validate_mapping(mapping, rows)
+        unless mapping.all? { |v| v.is_a?(Integer) && v >= 0 && v < rows }
+          raise ConversionError, "mapping tensor must contain integer row indices in 0...#{rows}"
         end
+      end
 
-        [Safetensors.f32_payload(path, tensor), dim]
+      def validate_weights(weights)
+        raise ConversionError, "weights tensor contains non-finite values" unless weights.all? { |w| w.to_f.finite? }
       end
 
       def resolve_max_tokens(requested, default)
@@ -62,6 +108,14 @@ module StaticEmbeddings
         value
       end
 
+      def resolve_truncation(requested, default)
+        key = (requested || default).to_s.tr("-", "_").to_sym
+        Format::TRUNCATION_POLICIES.fetch(key) do
+          raise InvalidOptionError,
+                "truncation must be one of #{Format::TRUNCATION_POLICIES.keys.inspect}, got #{requested.inspect}"
+        end
+      end
+
       def resolve_dimensions(requested, native_dim)
         return native_dim if requested.nil?
 

data/lib/static_embeddings/provenance.rb

@@ -4,8 +4,9 @@ require "static_embeddings/bert_wordpiece"
 
 module StaticEmbeddings
   module Provenance
-    MODEL2VEC_VERSION = "0.9.0"
-    TOKENIZERS_VERSION = "0.23.1"
+    MODEL2VEC_VERSION = "0.10.0"
+    MODEL2VEC_USABLE_IDS_VERSION = "0.9.0"
+    TOKENIZERS_VERSION = "0.23.3"
     UNICODE_CATEGORIES_VERSION = "0.1.1"
     NOTES = "Vectors are only reference-compatible if docs/MODEL_AUDIT.md records a passing oracle run for this source revision."
 
@@ -35,6 +36,7 @@ module StaticEmbeddings
         "reference_impl" => source.oracle,
         "reference_max_tokens" => runtime.max_tokens,
         "reference_tokenizers_version" => TOKENIZERS_VERSION,
+        "truncation" => Format::TRUNCATION_POLICIES.key(runtime.truncation).to_s,
         "source_family" => source.family,
         "source_files_sha256" => source.files_sha256,
         "source_model_id" => source.model,
@@ -46,12 +48,21 @@ module StaticEmbeddings
         "vocab_size" => model.tokens.length
       }
       data["trained_mrl_dims"] = dimensions.trained if dimensions.trained
-      add_model2vec_versions(data) if source.family == "model2vec"
+      add_baked_side_tensors(data, source.baked)
+      add_model2vec_versions(data, runtime) if source.family == "model2vec"
       data
     end
 
-    def add_model2vec_versions(data)
-      data["reference_model2vec_version"] = MODEL2VEC_VERSION
+    def add_baked_side_tensors(data, baked)
+      return unless baked && (baked[:weights] || baked[:mapping_clusters])
+
+      data["baked_token_weights"] = baked[:weights]
+      data["baked_vocabulary_quantization"] = baked[:mapping_clusters] if baked[:mapping_clusters]
+    end
+
+    def add_model2vec_versions(data, runtime)
+      data["reference_model2vec_version"] =
+        runtime.truncation == Format::TRUNCATE_RAW_IDS_BEFORE_UNK_DROP ? MODEL2VEC_VERSION : MODEL2VEC_USABLE_IDS_VERSION
       data["reference_unicode_categories_version"] = UNICODE_CATEGORIES_VERSION
     end
   end

data/lib/static_embeddings/reference.rb

@@ -25,52 +25,55 @@ module StaticEmbeddings
       @dim = meta.fetch(:dim)
     end
 
-    def self.from_source_dir(dir, max_tokens: nil, dimensions: nil)
-      from_canonical(Importers.import(dir, max_tokens: max_tokens, dimensions: dimensions))
-    end
-
-    def self.from_canonical(model)
-      dim = model.dimensions.output
-      floats = matrix_bytes(model.matrix).unpack("e*")
-      rows = model.tokens.length
-      new(
-        tokens: model.tokens,
-        matrix: Array.new(rows) { |index| floats[index * dim, dim] },
-        meta: meta_from_canonical(model)
-      )
-    end
-
-    def self.matrix_bytes(matrix)
-      return matrix unless matrix.respond_to?(:each_chunk)
-
-      matrix.each_chunk.each_with_object(Format.binary_string) { |chunk, packed| packed << chunk }
-    end
-
-    def self.meta_from_canonical(model)
-      tokenizer = model.tokenizer
-      runtime = model.runtime
-      {
-        dim: model.dimensions.output,
-        lowercase: tokenizer.fetch(:do_lower_case),
-        strip_accents: tokenizer.fetch(:strip_accents),
-        clean_text: tokenizer.fetch(:clean_text),
-        handle_chinese_chars: tokenizer.fetch(:handle_chinese_chars),
-        max_input_chars_per_word: tokenizer.fetch(:max_input_chars_per_word),
-        unk_token: model.tokens.fetch(tokenizer.fetch(:unk_id)),
-        normalize: runtime.normalization == Format::NORMALIZATION_L2,
-        max_tokens: runtime.max_tokens,
-        unk_policy: runtime.unk_policy,
-        added_tokens: added_tokens(model)
-      }
-    end
-
-    def self.added_tokens(model)
-      mask = model.tokenizer.fetch(:added_token_mask)
-      BertWordPiece::STANDARD_SPECIAL_TOKENS.each_with_object({}) do |(content, bit), tokens|
-        next if (mask & bit).zero?
-
-        id = model.tokens.index(content)
-        tokens[content] = id unless id.nil?
+    class << self
+      def from_source_dir(dir, max_tokens: nil, dimensions: nil, truncation: nil)
+        from_canonical(Importers.import(dir, max_tokens: max_tokens, dimensions: dimensions, truncation: truncation))
+      end
+
+      def from_canonical(model)
+        dim = model.dimensions.output
+        floats = matrix_bytes(model.matrix).unpack("e*")
+        rows = model.tokens.length
+        new(
+          tokens: model.tokens,
+          matrix: Array.new(rows) { |index| floats[index * dim, dim] },
+          meta: meta_from_canonical(model)
+        )
+      end
+
+      def matrix_bytes(matrix)
+        return matrix unless matrix.respond_to?(:each_chunk)
+
+        matrix.each_chunk.each_with_object(Format.binary_string) { |chunk, packed| packed << chunk }
+      end
+
+      def meta_from_canonical(model)
+        tokenizer = model.tokenizer
+        runtime = model.runtime
+        {
+          dim: model.dimensions.output,
+          lowercase: tokenizer.fetch(:do_lower_case),
+          strip_accents: tokenizer.fetch(:strip_accents),
+          clean_text: tokenizer.fetch(:clean_text),
+          handle_chinese_chars: tokenizer.fetch(:handle_chinese_chars),
+          max_input_chars_per_word: tokenizer.fetch(:max_input_chars_per_word),
+          unk_token: model.tokens.fetch(tokenizer.fetch(:unk_id)),
+          normalize: runtime.normalization == Format::NORMALIZATION_L2,
+          max_tokens: runtime.max_tokens,
+          unk_policy: runtime.unk_policy,
+          truncation: runtime.truncation,
+          added_tokens: added_tokens(model)
+        }
+      end
+
+      def added_tokens(model)
+        mask = model.tokenizer.fetch(:added_token_mask)
+        BertWordPiece::STANDARD_SPECIAL_TOKENS.each_with_object({}) do |(content, bit), tokens|
+          next if (mask & bit).zero?
+
+          id = model.tokens.index(content)
+          tokens[content] = id unless id.nil?
+        end
       end
     end
 
@@ -93,8 +96,12 @@ module StaticEmbeddings
     end
 
     def embed(text, max_tokens: @meta[:max_tokens])
-      ids = apply_unk_policy(tokenize(text, max_tokens: false))
-      ids = ids.first(max_tokens) if max_tokens && max_tokens.positive? && ids.length > max_tokens
+      if @meta.fetch(:truncation, Format::TRUNCATE_USABLE_IDS_BEFORE_POOLING) == Format::TRUNCATE_RAW_IDS_BEFORE_UNK_DROP
+        ids = apply_unk_policy(tokenize(text, max_tokens: max_tokens))
+      else
+        ids = apply_unk_policy(tokenize(text, max_tokens: false))
+        ids = ids.first(max_tokens) if max_tokens && max_tokens.positive? && ids.length > max_tokens
+      end
       return Array.new(@dim, 0.0) if ids.empty?
 
       vector = pooled(ids)

data/lib/static_embeddings/safetensors.rb

@@ -4,6 +4,8 @@ module StaticEmbeddings
   module Safetensors
     MAX_HEADER_BYTES = 100 * 1024 * 1024
     SUPPORTED_DTYPES = { "F32" => 4, "F16" => 2, "BF16" => 2 }.freeze
+    AUX_DTYPES = { "F64" => 8, "I64" => 8, "I32" => 4 }.freeze
+    ELEMENT_SIZES = SUPPORTED_DTYPES.merge(AUX_DTYPES).freeze
     CONVERT_CHUNK_ELEMENTS = 64 * 1024
     IO_CHUNK_BYTES = 1024 * 1024
 
@@ -96,9 +98,35 @@ module StaticEmbeddings
     end
 
     def f32_payload(path, tensor)
+      unless SUPPORTED_DTYPES.key?(tensor.fetch(:dtype))
+        raise ConversionError,
+              "embedding matrix dtype #{tensor.fetch(:dtype)} is not supported (expected F32, F16, or BF16)"
+      end
       F32Payload.new(path, tensor)
     end
 
+    def read_vector(path, tensor)
+      bytes = File.open(path, "rb") do |io|
+        io.seek(tensor.fetch(:absolute_offset), IO::SEEK_SET)
+        io.read(tensor.fetch(:source_bytes))
+      end
+      unless bytes&.bytesize == tensor.fetch(:source_bytes)
+        raise ConversionError, "tensor data is truncated while reading #{path}"
+      end
+      bytes.force_encoding(Encoding::BINARY)
+
+      case tensor.fetch(:dtype)
+      when "F32" then bytes.unpack("e*")
+      when "F64" then bytes.unpack("E*")
+      when "F16" then convert_f16_to_f32(bytes).unpack("e*")
+      when "BF16" then convert_bf16_to_f32(bytes).unpack("e*")
+      when "I64" then bytes.unpack("q<*")
+      when "I32" then bytes.unpack("l<*")
+      else
+        raise ConversionError, "tensor dtype #{tensor[:dtype]} is not supported"
+      end
+    end
+
     def write(path, name, shape, floats)
       body = floats.pack("e*")
       write_raw(path, name, shape, "F32", body)
@@ -171,9 +199,10 @@ module StaticEmbeddings
     end
 
     def checked_tensor_bytes(name, dtype, shape)
-      element_size = SUPPORTED_DTYPES[dtype]
+      element_size = ELEMENT_SIZES[dtype]
       unless element_size
-        raise ConversionError, "tensor #{name}: dtype #{dtype} is not supported (expected F32, F16, or BF16)"
+        raise ConversionError,
+              "tensor #{name}: dtype #{dtype} is not supported (expected #{ELEMENT_SIZES.keys.join(', ')})"
       end
       unless shape.is_a?(Array) && shape.all? { |dim| dim.is_a?(Integer) && dim >= 0 }
         raise ConversionError, "tensor #{name}: malformed shape #{shape.inspect}"

data/lib/static_embeddings/version.rb

@@ -1,3 +1,3 @@
 module StaticEmbeddings
-  VERSION = "1.5.8"
+  VERSION = "1.5.9"
 end

data/lib/static_embeddings/weighted_payload.rb

@@ -0,0 +1,73 @@
+require "static_embeddings/errors"
+require "static_embeddings/format/constants"
+
+module StaticEmbeddings
+  class WeightedPayload
+    FLOAT_BYTES = 4
+    ROWS_PER_CHUNK = 1024
+
+    attr_reader :bytesize
+
+    def initialize(inner, dim:, source_rows:, output_rows:, weights: nil, mapping: nil)
+      raise InvalidOptionError, "weights or mapping is required" if weights.nil? && mapping.nil?
+      unless inner.bytesize == source_rows * dim * FLOAT_BYTES
+        raise ConversionError, "embedding payload does not hold #{source_rows} rows of #{dim} floats"
+      end
+
+      @inner = inner
+      @row_bytes = dim * FLOAT_BYTES
+      @output_rows = output_rows
+      @weights = weights
+      @mapping = mapping
+      @bytesize = output_rows * @row_bytes
+    end
+
+    def each_chunk(&block)
+      return enum_for(__method__) unless block_given?
+
+      @mapping ? each_mapped_chunk(&block) : each_streamed_chunk(&block)
+    end
+
+    private
+
+    def each_mapped_chunk
+      table = Format.binary_string(@inner.bytesize)
+      @inner.each_chunk { |chunk| table << chunk }
+
+      index = 0
+      while index < @output_rows
+        count = [ROWS_PER_CHUNK, @output_rows - index].min
+        out = Format.binary_string(count * @row_bytes)
+        count.times do
+          row = table.byteslice(@mapping[index] * @row_bytes, @row_bytes)
+          out << (@weights ? scale(row, @weights[index]) : row)
+          index += 1
+        end
+        yield out
+      end
+    end
+
+    def each_streamed_chunk
+      remainder = Format.binary_string
+      index = 0
+      @inner.each_chunk do |chunk|
+        data = remainder.empty? ? chunk : remainder + chunk
+        whole = (data.bytesize / @row_bytes) * @row_bytes
+        remainder = (data.byteslice(whole, data.bytesize - whole) || Format.binary_string).force_encoding(Encoding::BINARY)
+        next if whole.zero?
+
+        out = Format.binary_string(whole)
+        (whole / @row_bytes).times do |r|
+          out << scale(data.byteslice(r * @row_bytes, @row_bytes), @weights[index])
+          index += 1
+        end
+        yield out
+      end
+      raise ConversionError, "truncated embedding matrix while applying weights" unless remainder.empty?
+    end
+
+    def scale(row, weight)
+      row.unpack("e*").map! { |x| x * weight }.pack("e*")
+    end
+  end
+end

data/lib/static_embeddings.rb

@@ -54,7 +54,7 @@ module StaticEmbeddings
     end
 
     def convert(source_dir, output_path:, model_id: nil, max_tokens: nil, dimensions: nil,
-                source_revision: nil, trained_mrl_dims: nil)
+                source_revision: nil, trained_mrl_dims: nil, truncation: nil)
       require "static_embeddings/conversion"
       Conversion.call(
         source_dir,
@@ -63,7 +63,8 @@ module StaticEmbeddings
         max_tokens: max_tokens,
         dimensions: dimensions,
         source_revision: source_revision,
-        trained_mrl_dims: trained_mrl_dims
+        trained_mrl_dims: trained_mrl_dims,
+        truncation: truncation
       )
     end
 

data/tools/check_model2vec_parity.rb

@@ -20,9 +20,15 @@ end
 
 model = StaticEmbeddings.load(options[:model], verify: true)
 payload = JSON.parse(File.read(options[:oracle], encoding: "UTF-8"))
-abort "unsupported oracle schema #{payload["schema_version"].inspect}" unless payload["schema_version"] == 2
+schema = payload["schema_version"]
+abort "unsupported oracle schema #{schema.inspect}" unless [2, 3].include?(schema)
 
 reference = payload.fetch("reference")
+oracle_truncation = (reference["truncation"] || "usable_ids").to_sym
+if oracle_truncation != model.truncation
+  abort "oracle truncation=#{oracle_truncation} (model2vec #{reference['model2vec']}) " \
+        "but model.truncation=#{model.truncation}; reconvert with --truncation #{oracle_truncation.to_s.tr('_', '-')}"
+end
 max_length = Integer(reference.fetch("max_length"))
 if max_length != model.max_tokens
   abort "oracle max_length=#{max_length} but model.max_tokens=#{model.max_tokens}"
@@ -79,16 +85,17 @@ rows.each_with_index do |row, i|
   end
 
   full_raw = model.tokenize(text, max_tokens: false)
-  expected_usable = row.fetch("static_usable_token_ids")
-  got_usable = if reference["unk_token_id"].nil?
-                 full_raw.first(max_length)
+  expected_usable = row["static_pooled_token_ids"] || row.fetch("static_usable_token_ids")
+  drop_unk = ->(ids) { reference["unk_token_id"].nil? ? ids : ids.reject { |id| id == model.unk_id } }
+  got_usable = if model.truncation == :raw_ids
+                 drop_unk.call(full_raw.first(max_length))
                else
-                 full_raw.reject { |id| id == model.unk_id }.first(max_length)
+                 drop_unk.call(full_raw).first(max_length)
                end
   if got_usable != expected_usable
     usable_failures << i
     first = got_usable.zip(expected_usable).index { |a, b| a != b } || [got_usable.length, expected_usable.length].min
-    problems << "usable ids differ at #{first} (got #{got_usable.length}, ref #{expected_usable.length})"
+    problems << "pooled ids differ at #{first} (got #{got_usable.length}, ref #{expected_usable.length})"
   end
 
   got_vector_blob = model.embed(text)

data/tools/eval_retrieval.rb

@@ -2,11 +2,14 @@ require "json"
 require "optparse"
 require "static_embeddings"
 
-options = { k: 10 }
+options = { k: 10, bootstrap: 2000, seed: 42, per_query: false }
 parser = OptionParser.new do |opts|
   opts.on("--model PATH") { |value| options[:model] = value }
   opts.on("--dataset PATH") { |value| options[:dataset] = value }
   opts.on("--k N", Integer) { |value| options[:k] = value }
+  opts.on("--bootstrap N", Integer, "bootstrap resamples (0 disables)") { |value| options[:bootstrap] = value }
+  opts.on("--seed N", Integer) { |value| options[:seed] = value }
+  opts.on("--per-query", "include each query's first relevant rank") { options[:per_query] = true }
 end
 parser.parse!(ARGV)
 abort "usage: ruby -Ilib tools/eval_retrieval.rb --model MODEL.semb --dataset dataset.json" unless options[:model] && options[:dataset]
@@ -15,44 +18,109 @@ payload = JSON.parse(File.read(options[:dataset], encoding: "UTF-8"))
 docs = payload.fetch("documents")
 queries = payload.fetch("queries")
 k = options[:k]
+abort "--k must be positive" unless k.positive?
+abort "--bootstrap must be nonnegative" if options[:bootstrap].negative?
+abort "dataset has no documents" if docs.empty?
+abort "dataset has no queries" if queries.empty?
 
-model = StaticEmbeddings.load(options[:model], verify: true)
-matrix = model.embed_batch(docs.map { |doc| doc.fetch("text") })
+doc_ids = docs.map { |doc| doc.fetch("id").to_s }
+unknown = queries.flat_map { |q| q.fetch("relevant").map(&:to_s) } - doc_ids
+abort "queries reference unknown document ids: #{unknown.uniq.inspect}" unless unknown.empty?
 
 def dcg(gains)
   gains.each_with_index.sum { |gain, i| gain / Math.log2(i + 2) }
 end
 
-mrr = 0.0
-ndcg = 0.0
-hits = 0
+def query_metrics(ranking, relevant, k)
+  first = ranking.index { |id| relevant.include?(id) }
+  top = ranking.first(k)
+  gains = top.map { |id| relevant.include?(id) ? 1.0 : 0.0 }
+  ideal = Array.new([relevant.length, top.length].min, 1.0)
+  ideal_dcg = dcg(ideal)
+  {
+    first_rank: first && first + 1,
+    rr: first ? 1.0 / (first + 1) : 0.0,
+    rr_at_k: first && first < k ? 1.0 / (first + 1) : 0.0,
+    ndcg_at_k: ideal_dcg.positive? ? dcg(gains) / ideal_dcg : 0.0,
+    recall_at_k: relevant.empty? ? 0.0 : gains.sum / relevant.length,
+    hit_at_k: first && first < k ? 1.0 : 0.0
+  }
+end
+
+METRICS = %i[rr rr_at_k ndcg_at_k recall_at_k hit_at_k].freeze
+NAMES = { rr: "mrr", rr_at_k: "mrr_at_k", ndcg_at_k: "ndcg_at_k", recall_at_k: "recall_at_k", hit_at_k: "hit_at_k" }.freeze
+
+def mean(values)
+  values.sum / values.length.to_f
+end
+
+def summarize(per_query)
+  METRICS.to_h { |m| [NAMES.fetch(m), mean(per_query.map { |q| q.fetch(m) })] }
+end
+
+def bootstrap(per_query, resamples, rng)
+  return nil if resamples.zero?
+
+  n = per_query.length
+  samples = METRICS.to_h { |m| [m, []] }
+  resamples.times do
+    draw = Array.new(n) { per_query[rng.rand(n)] }
+    METRICS.each { |m| samples[m] << mean(draw.map { |q| q.fetch(m) }) }
+  end
+  METRICS.to_h do |m|
+    sorted = samples[m].sort
+    [NAMES.fetch(m), [sorted[(0.025 * (resamples - 1)).round], sorted[(0.975 * (resamples - 1)).round]]]
+  end
+end
+
+model = StaticEmbeddings.load(options[:model], verify: true)
+matrix = model.embed_batch(docs.map { |doc| doc.fetch("text") }).freeze
+rng = Random.new(options[:seed])
 
-queries.each do |query|
+zero_queries = []
+per_query = queries.map do |query|
   relevant = query.fetch("relevant").map(&:to_s)
   blob = model.embed(query.fetch("text"))
-  ranked = model.cosine_top_k(blob, matrix, [k, docs.length].min)
-  ids = ranked.map { |index, _| docs[index].fetch("id").to_s }
-  rank = ids.index { |id| relevant.include?(id) }
-  mrr += rank ? 1.0 / (rank + 1) : 0.0
-  hits += 1 if rank && rank < k
-  gains = ids.map { |id| relevant.include?(id) ? 1.0 : 0.0 }
-  ideal = [1.0] * [relevant.length, k].min + [0.0] * [k - relevant.length, 0].max
-  ideal = ideal.first(gains.length)
-  denom = dcg(ideal)
-  ndcg += denom.positive? ? dcg(gains) / denom : 0.0
-end
-
-n = queries.length.to_f
+  ranking =
+    if blob.unpack("e*").all?(&:zero?)
+      zero_queries << query.fetch("id", query.fetch("text")).to_s
+      []
+    else
+      model.cosine_top_k(blob, matrix, docs.length).map { |index, _| doc_ids[index] }
+    end
+
+  query_metrics(ranking, relevant, k).merge(id: query.fetch("id", nil))
+end
+
+random_runs = 200
+random = Array.new(random_runs) do
+  summarize(queries.map { |query| query_metrics(doc_ids.shuffle(random: rng), query.fetch("relevant").map(&:to_s), k) })
+end
+random_baseline = random.first.keys.to_h { |key| [key, mean(random.map { |run| run.fetch(key) })] }
+
+warnings = []
+if k >= docs.length
+  warnings << "k=#{k} >= documents=#{docs.length}: hit_at_k and recall_at_k cannot distinguish rankings when queries have relevant documents"
+end
+if queries.length < 50
+  warnings << "only #{queries.length} queries: compare models by the bootstrap intervals, not the point estimates"
+end
+warnings << "#{zero_queries.length} queries embedded to a zero vector: #{zero_queries.inspect}" unless zero_queries.empty?
+
 result = {
   "model_id" => model.model_id,
   "dim" => model.dim,
   "normalized" => model.normalized?,
+  "truncation" => model.truncation,
   "queries" => queries.length,
   "documents" => docs.length,
   "k" => k,
-  "mrr" => mrr / n,
-  "ndcg_at_k" => ndcg / n,
-  "hit_at_k" => hits / n
+  "metrics" => summarize(per_query),
+  "bootstrap_95" => bootstrap(per_query, options[:bootstrap], rng),
+  "random_baseline" => random_baseline,
+  "warnings" => warnings
 }
+result["per_query"] = per_query.map { |q| { "id" => q[:id], "first_relevant_rank" => q[:first_rank] } } if options[:per_query]
+warnings.each { |warning| warn "warning: #{warning}" }
 puts JSON.pretty_generate(result)
 model.close

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: static_embeddings
 version: !ruby/object:Gem::Version
-  version: 1.5.8
+  version: 1.5.9
 platform: ruby
 authors:
 - Roman Haydarov
@@ -127,6 +127,7 @@ files:
 - lib/static_embeddings/safetensors.rb
 - lib/static_embeddings/unicode_tables.rb
 - lib/static_embeddings/version.rb
+- lib/static_embeddings/weighted_payload.rb
 - static_embeddings.gemspec
 - tools/benchmark.rb
 - tools/build_demo_model.rb