static_embeddings 0.1.4 → 1.5.6

data/tools/check_st_parity.rb

@@ -0,0 +1,125 @@
+require "json"
+require "optparse"
+require "static_embeddings"
+
+options = {
+  min_cosine: 1.0 - 1e-6,
+  max_abs: 1e-5
+}
+
+parser = OptionParser.new do |opts|
+  opts.on("--model PATH") { |value| options[:model] = value }
+  opts.on("--oracle PATH") { |value| options[:oracle] = value }
+  opts.on("--min-cosine N", Float) { |value| options[:min_cosine] = value }
+  opts.on("--max-abs N", Float) { |value| options[:max_abs] = value }
+end
+parser.parse!(ARGV)
+unless options[:model] && options[:oracle]
+  abort "usage: ruby -Ilib tools/check_st_parity.rb --model MODEL.semb --oracle oracle.json"
+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"] == 1
+
+reference = payload.fetch("reference")
+abort "ST oracle must use add_special_tokens=false" unless reference["add_special_tokens"] == false
+abort "ST oracle must use normalize_embeddings=false" unless reference["normalize_embeddings"] == false
+if model.normalized?
+  abort "ST source-faithful .semb must not L2-normalize; got normalized?=true"
+end
+unless model.max_tokens == false
+  abort "ST source-faithful .semb must bake unlimited tokens; got max_tokens=#{model.max_tokens.inspect}"
+end
+if model.provenance["unk_policy"] != "include"
+  abort "ST source-faithful .semb must use UNK_INCLUDE; got #{model.provenance["unk_policy"].inspect}"
+end
+
+def prefix_dims(vector, dim)
+  return vector if vector.length == dim
+  abort "oracle vector shorter than model dim (#{vector.length} < #{dim})" if vector.length < dim
+
+  vector.first(dim)
+end
+
+def dot(a, b)
+  a.zip(b).sum { |x, y| x * y }
+end
+
+def norm(a)
+  Math.sqrt(a.sum { |x| x * x })
+end
+
+def vector_metrics(reference, got)
+  max_abs = reference.zip(got).map { |a, b| (a - b).abs }.max || 0.0
+  ref_zero = reference.all?(&:zero?)
+  got_zero = got.all?(&:zero?)
+  cosine =
+    if ref_zero && got_zero
+      1.0
+    elsif ref_zero || got_zero
+      0.0
+    else
+      dot(reference, got) / (norm(reference) * norm(got))
+    end
+  [cosine, max_abs]
+end
+
+raw_failures = []
+vector_failures = []
+diag_failures = []
+min_cosine = 1.0
+max_abs_all = 0.0
+vectors_checked = 0
+
+payload.fetch("rows").each_with_index do |row, i|
+  text = row.fetch("text")
+  label = row.fetch("label", i.to_s)
+  problems = []
+
+  expected_raw = row.fetch("hf_raw_token_ids")
+  got_raw = model.tokenize(text, max_tokens: false)
+  if got_raw != expected_raw
+    raw_failures << i
+    first = got_raw.zip(expected_raw).index { |a, b| a != b } || [got_raw.length, expected_raw.length].min
+    problems << "raw ids differ at #{first} (got #{got_raw.length}, ref #{expected_raw.length})"
+  end
+
+  got_vector = model.embed_array(text)
+  reference_vector = prefix_dims(row.fetch("st_encode_vector"), got_vector.length)
+  cosine, max_abs = vector_metrics(reference_vector, got_vector)
+  vectors_checked += 1
+  min_cosine = [min_cosine, cosine].min
+  max_abs_all = [max_abs_all, max_abs].max
+  unless cosine >= options[:min_cosine] && max_abs <= options[:max_abs]
+    vector_failures << i
+    problems << format("encode vector out of tolerance cos=%.10f max_abs=%.8g", cosine, max_abs)
+  end
+
+  if row.key?("static_embedding_vector")
+    diag = prefix_dims(row.fetch("static_embedding_vector"), got_vector.length)
+    dcos, dabs = vector_metrics(diag, got_vector)
+    unless dcos >= options[:min_cosine] && dabs <= options[:max_abs]
+      diag_failures << i
+      problems << format("StaticEmbedding diagnostic out of tolerance cos=%.10f max_abs=%.8g", dcos, dabs)
+    end
+  end
+
+  status = problems.empty? ? "ok" : "FAIL"
+  puts "#{status} idx=#{format('%03d', i)} label=#{label.inspect} bytes=#{text.bytesize}" +
+       (problems.empty? ? "" : " [#{problems.join('; ')}]")
+end
+
+puts "rows=#{payload.fetch("rows").length}"
+puts "vectors_checked=#{vectors_checked}"
+puts "min_cosine=#{min_cosine}"
+puts "max_abs_all=#{max_abs_all}"
+puts "raw_token_id_failures=#{raw_failures.inspect}"
+puts "vector_failures=#{vector_failures.inspect}"
+puts "static_embedding_failures=#{diag_failures.inspect}"
+
+unless raw_failures.empty? && vector_failures.empty? && diag_failures.empty?
+  abort "ST parity failed"
+end
+
+puts "ST corpus parity OK (#{payload.fetch("rows").length}/#{payload.fetch("rows").length})"

data/tools/eval_retrieval.rb

@@ -0,0 +1,58 @@
+require "json"
+require "optparse"
+require "static_embeddings"
+
+options = { k: 10 }
+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 }
+end
+parser.parse!(ARGV)
+abort "usage: ruby -Ilib tools/eval_retrieval.rb --model MODEL.semb --dataset dataset.json" unless options[:model] && options[:dataset]
+
+payload = JSON.parse(File.read(options[:dataset], encoding: "UTF-8"))
+docs = payload.fetch("documents")
+queries = payload.fetch("queries")
+k = options[:k]
+
+model = StaticEmbeddings.load(options[:model], verify: true)
+matrix = model.embed_batch(docs.map { |doc| doc.fetch("text") })
+
+def dcg(gains)
+  gains.each_with_index.sum { |gain, i| gain / Math.log2(i + 2) }
+end
+
+mrr = 0.0
+ndcg = 0.0
+hits = 0
+
+queries.each 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
+result = {
+  "model_id" => model.model_id,
+  "dim" => model.dim,
+  "normalized" => model.normalized?,
+  "queries" => queries.length,
+  "documents" => docs.length,
+  "k" => k,
+  "mrr" => mrr / n,
+  "ndcg_at_k" => ndcg / n,
+  "hit_at_k" => hits / n
+}
+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: 0.1.4
+  version: 1.5.6
 platform: ruby
 authors:
 - Roman Haydarov
@@ -65,10 +65,11 @@ dependencies:
     - - "~>"
       - !ruby/object:Gem::Version
         version: '1.2'
-description: A small C-extension runtime for Model2Vec-style static embedding models.
-  Models are converted offline into a flat mmap-able .semb file; at runtime the gem
-  tokenizes (BERT WordPiece), looks up rows and mean-pools them. Releases the GVL
-  on large native work, rejects internal thread fan-out, and links nothing but libc.
+description: A small C-extension runtime for Model2Vec and Sentence Transformers static
+  WordPiece embedding models. Models are converted offline into a flat mmap-able .semb
+  file; at runtime the gem tokenizes (BERT WordPiece), looks up rows and mean-pools
+  them. Releases the GVL on large native work, rejects internal thread fan-out, and
+  links nothing but libc.
 email:
 - romnhajdarov@gmail.com
 executables:
@@ -100,14 +101,29 @@ files:
 - ext/static_embeddings/se_topk.c
 - ext/static_embeddings/se_unicode.c
 - ext/static_embeddings/static_embeddings.c
+- lib/models/demo.semb
 - lib/static_embeddings.rb
+- lib/static_embeddings/bert_wordpiece.rb
+- lib/static_embeddings/canonical.rb
 - lib/static_embeddings/cli.rb
-- lib/static_embeddings/converter.rb
+- lib/static_embeddings/codec.rb
+- lib/static_embeddings/conversion.rb
 - lib/static_embeddings/errors.rb
 - lib/static_embeddings/format.rb
+- lib/static_embeddings/format/constants.rb
+- lib/static_embeddings/format/hash_table.rb
+- lib/static_embeddings/format/trie.rb
+- lib/static_embeddings/format/verifier.rb
+- lib/static_embeddings/format/writer.rb
+- lib/static_embeddings/importers.rb
+- lib/static_embeddings/importers/model2vec.rb
+- lib/static_embeddings/importers/sentence_transformers_static.rb
+- lib/static_embeddings/importers/support.rb
 - lib/static_embeddings/model.rb
 - lib/static_embeddings/paths.rb
+- lib/static_embeddings/provenance.rb
 - lib/static_embeddings/reference.rb
+- lib/static_embeddings/row_prefix_payload.rb
 - lib/static_embeddings/safetensors.rb
 - lib/static_embeddings/unicode_tables.rb
 - lib/static_embeddings/version.rb
@@ -115,6 +131,8 @@ files:
 - tools/benchmark.rb
 - tools/build_demo_model.rb
 - tools/check_model2vec_parity.rb
+- tools/check_st_parity.rb
+- tools/eval_retrieval.rb
 - tools/make_fixture_model.rb
 homepage: https://github.com/roman-haidarov/static_embeddings
 licenses:

data/lib/static_embeddings/converter.rb

@@ -1,284 +0,0 @@
-require "json"
-require "digest"
-
-module StaticEmbeddings
-  class Converter
-    REFERENCE_IMPL = "model2vec.StaticModel"
-    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
-    SOURCE_FILES = %w[tokenizer.json config.json tokenizer_config.json model.safetensors].freeze
-    TOKENIZER_PROFILE = "BERT_WORDPIECE_V1"
-
-    attr_reader :source_dir, :report
-
-    def initialize(source_dir)
-      @source_dir = source_dir
-      @report = {}
-    end
-
-    def convert(output_path:, model_id: nil, max_tokens: REFERENCE_MAX_TOKENS)
-      source = load_source
-      profile = audit_tokenizer(source[:tokenizer], source[:tokenizer_config])
-      tokens = extract_vocab(source[:tokenizer])
-      matrix, dim = extract_matrix(tokens.length)
-      meta = runtime_meta(source[:config], profile, tokens, max_tokens)
-
-      result = write_model(output_path, meta, tokens, matrix,
-                           provenance_for(model_id, meta, profile, source[:config], tokens.length, dim))
-      @report = result.merge(vocab_size: tokens.length, dim: dim, provenance: JSON.parse(result[:provenance]))
-    end
-
-    private
-
-    def load_source
-      {
-        tokenizer: load_json("tokenizer.json"),
-        config: load_json("config.json", optional: true) || {},
-        tokenizer_config: load_json("tokenizer_config.json", optional: true) || {}
-      }
-    end
-
-    def load_json(name, optional: false)
-      path = File.join(source_dir, name)
-      return nil if optional && !File.file?(path)
-      raise ConversionError, "missing #{name} in #{source_dir}" unless File.file?(path)
-
-      JSON.parse(File.binread(path))
-    end
-
-    def reject!(reason)
-      raise UnsupportedModelError,
-            "#{reason}. This converter only accepts the #{TOKENIZER_PROFILE} profile; " \
-            "supporting this model would require implementing that behaviour in the runtime first."
-    end
-
-    def audit_tokenizer(tokenizer, tokenizer_config)
-      model = tokenizer["model"] or reject!("tokenizer.json has no model section")
-      normalizer = tokenizer["normalizer"] or reject!("tokenizer has no normalizer")
-      pre_tokenizer = tokenizer["pre_tokenizer"]
-
-      assert_wordpiece!(model)
-      assert_normalizer!(normalizer)
-      assert_pre_tokenizer!(pre_tokenizer)
-      audit_added_tokens(tokenizer)
-
-      profile = profile_from(model, normalizer, tokenizer_config)
-      reject!("clean_text=false is not supported by the runtime") unless profile[:clean_text]
-      profile
-    end
-
-    def assert_wordpiece!(model)
-      reject!("tokenizer model.type is #{model['type'].inspect}, expected WordPiece") unless model["type"] == "WordPiece"
-
-      prefix = model.fetch("continuing_subword_prefix", "##")
-      reject!("continuing_subword_prefix #{prefix.inspect} is not supported") unless prefix == "##"
-
-      unk = model.fetch("unk_token", "[UNK]")
-      reject!("unk_token #{unk.inspect} is not supported") unless unk == "[UNK]"
-    end
-
-    def assert_normalizer!(normalizer)
-      reject!("normalizer type #{normalizer['type'].inspect} is not BertNormalizer") unless normalizer["type"] == "BertNormalizer"
-
-      unknown = normalizer.keys - ALLOWED_NORMALIZER_KEYS
-      reject!("normalizer has unsupported keys #{unknown.inspect}") unless unknown.empty?
-    end
-
-    def assert_pre_tokenizer!(pre_tokenizer)
-      return if pre_tokenizer && pre_tokenizer["type"] == "BertPreTokenizer"
-
-      reject!("pre_tokenizer type #{pre_tokenizer && pre_tokenizer['type'].inspect} is not BertPreTokenizer")
-    end
-
-    def profile_from(model, normalizer, tokenizer_config)
-      lowercase = flag_value(normalizer, "lowercase", tokenizer_config["do_lower_case"], true)
-      {
-        lowercase: lowercase,
-        strip_accents: strip_accents?(normalizer, lowercase),
-        clean_text: flag_value(normalizer, "clean_text", nil, true),
-        handle_chinese_chars: flag_value(normalizer, "handle_chinese_chars",
-                                         tokenizer_config["tokenize_chinese_chars"], true),
-        continuing_subword_prefix: model.fetch("continuing_subword_prefix", "##"),
-        unk_token: model.fetch("unk_token", "[UNK]"),
-        max_input_chars_per_word: model.fetch("max_input_chars_per_word", 100).to_i,
-        tokenizer_class: tokenizer_config["tokenizer_class"]
-      }
-    end
-
-    def flag_value(hash, key, fallback, default)
-      return !!hash[key] if hash.key?(key) && !hash[key].nil?
-      return !!fallback unless fallback.nil?
-
-      default
-    end
-
-    def strip_accents?(normalizer, lowercase)
-      return !!normalizer["strip_accents"] if normalizer.key?("strip_accents") && !normalizer["strip_accents"].nil?
-
-      lowercase
-    end
-
-    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?
-
-      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
-    end
-
-    def standard_special?(token)
-      token["special"] && STANDARD_SPECIAL_TOKENS.include?(token["content"])
-    end
-
-    def extract_vocab(tokenizer)
-      vocab = tokenizer.dig("model", "vocab") or reject!("tokenizer.json has no model.vocab")
-      vocab.each_with_object(Array.new(vocab.length)) do |(token, id), tokens|
-        validate_vocab_id!(token, id, vocab.length, tokens)
-        tokens[id] = token
-      end.tap do |tokens|
-        missing = tokens.index(nil)
-        raise ConversionError, "vocab has a hole at id #{missing}" if missing
-      end
-    end
-
-    def validate_vocab_id!(token, id, size, tokens)
-      reject!("vocab id #{id.inspect} is not an integer") unless id.is_a?(Integer)
-      unless id.between?(0, size - 1)
-        raise ConversionError, "vocab id #{id} for #{token.inspect} is outside 0...#{size}"
-      end
-      raise ConversionError, "duplicate vocab id #{id}" unless tokens[id].nil?
-    end
-
-    def extract_matrix(vocab_size)
-      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])
-      rows, dim = tensor[:shape]
-      if rows != vocab_size
-        raise ConversionError,
-              "embedding matrix #{name} has #{rows} rows but the tokenizer has " \
-              "#{vocab_size} tokens — refusing to guess the mapping"
-      end
-
-      [tensor[:bytes], dim]
-    end
-
-    def sole_matrix_tensor(tensors)
-      matrices = tensors.select { |_, tensor| tensor[:shape].length == 2 }
-      raise ConversionError, "model.safetensors contains no 2-D tensor" if matrices.empty?
-      return matrices.first if matrices.length == 1
-
-      raise ConversionError,
-            "model.safetensors contains several 2-D tensors (#{matrices.keys.inspect}); " \
-            "a static embedding model must have exactly one"
-    end
-
-    def runtime_meta(config, profile, tokens, max_tokens)
-      base_meta(config, max_tokens)
-        .merge(tokenizer_meta(profile, tokens))
-        .merge(token_ids(tokens, profile))
-    end
-
-    def base_meta(config, max_tokens)
-      {
-        dim: nil,
-        normalization_type: normalization_from(config),
-        max_tokens_default: max_tokens.to_i,
-        add_special_tokens: false,
-        unk_policy: Format::UNK_DROP,
-        empty_policy: Format::EMPTY_ZERO_VECTOR
-      }
-    end
-
-    def tokenizer_meta(profile, tokens)
-      {
-        do_lower_case: profile[:lowercase],
-        strip_accents: profile[:strip_accents],
-        handle_chinese_chars: profile[:handle_chinese_chars],
-        clean_text: profile[:clean_text],
-        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]
-      }
-    end
-
-    def max_token_chars(tokens, prefix)
-      tokens.map do |token|
-        body = token.start_with?(prefix) ? token[prefix.length..] : token
-        body.each_char.count
-      end.max || 1
-    end
-
-    def token_ids(tokens, profile)
-      {
-        pad_id: token_id(tokens, "[PAD]", 0),
-        unk_id: token_id(tokens, profile[:unk_token]),
-        cls_id: token_id(tokens, "[CLS]", 0),
-        sep_id: token_id(tokens, "[SEP]", 0),
-        mask_id: token_id(tokens, "[MASK]", 0)
-      }
-    end
-
-    def normalization_from(config)
-      config.key?("normalize") && !config["normalize"] ? Format::NORMALIZATION_NONE : Format::NORMALIZATION_L2
-    end
-
-    def token_id(tokens, token, fallback = nil)
-      id = tokens.index(token)
-      return id unless id.nil?
-      return fallback unless fallback.nil?
-
-      raise ConversionError, "vocabulary has no #{token.inspect}"
-    end
-
-    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)
-      end
-    end
-
-    def provenance_for(model_id, meta, profile, config, vocab_size, dim)
-      ordered_json(
-        "format_version" => Format::VERSION,
-        "converter_version" => StaticEmbeddings::VERSION,
-        "source_model_id" => model_id || File.basename(File.expand_path(source_dir)),
-        "source_files_sha256" => source_digests,
-        "reference_impl" => REFERENCE_IMPL,
-        "reference_max_tokens" => meta[:max_tokens_default],
-        "unicode_source" => UnicodeTables.source_stamp,
-        "tokenizer_profile" => TOKENIZER_PROFILE,
-        "tokenizer_class" => profile[:tokenizer_class],
-        "vocab_size" => vocab_size,
-        "dim" => dim,
-        "normalize" => meta[:normalization_type] == Format::NORMALIZATION_L2,
-        "config_seq_length" => config["seq_length"],
-        "notes" => "Vectors are only reference-compatible if docs/MODEL_AUDIT.md records a passing oracle run for this source revision."
-      )
-    end
-
-    def ordered_json(hash)
-      JSON.generate(hash.sort_by { |key, _| key }.to_h)
-    end
-
-    def write_model(output_path, meta, tokens, matrix, provenance)
-      dim = matrix.bytesize / (tokens.length * 4)
-      result = Format.write(
-        path: output_path,
-        meta: meta.merge(dim: dim),
-        tokens: tokens,
-        matrix: matrix,
-        norm_tables: UnicodeTables.packed,
-        provenance: provenance
-      )
-      result.merge(provenance: provenance)
-    end
-  end
-end