static_embeddings 0.1.4 → 1.5.6

data/lib/static_embeddings/importers.rb

@@ -0,0 +1,50 @@
+require "static_embeddings/errors"
+require "static_embeddings/bert_wordpiece"
+require "static_embeddings/canonical"
+require "static_embeddings/importers/support"
+require "static_embeddings/importers/model2vec"
+require "static_embeddings/importers/sentence_transformers_static"
+
+module StaticEmbeddings
+  module Importers
+    module_function
+
+    def import(source_dir, **options)
+      root = source_root(source_dir)
+      case detect(root)
+      when :model2vec then Model2Vec.call(root, **options)
+      when :sentence_transformers_static then SentenceTransformersStatic.call(root, **options)
+      end
+    end
+
+    def detect(root)
+      config = Support.load_object(File.join(root, "config.json"), optional: true) || {}
+      tokenizer = File.join(root, "tokenizer.json")
+      weights = File.join(root, "model.safetensors")
+      modules = File.join(root, "modules.json")
+
+      return :model2vec if model2vec_config?(config) && File.file?(tokenizer) && File.file?(weights)
+      return :sentence_transformers_static if File.file?(modules)
+
+      if File.file?(tokenizer) && File.file?(weights)
+        raise InvalidSourceError,
+              "cannot prove source family in #{root}: config.json is not Model2Vec and modules.json is absent"
+      end
+
+      raise InvalidSourceError,
+            "cannot detect a static embedding source in #{root}: need a Model2Vec config with " \
+            "tokenizer.json + model.safetensors, or modules.json for Sentence Transformers StaticEmbedding"
+    end
+
+    def source_root(source_dir)
+      root = File.expand_path(source_dir.to_s)
+      raise InvalidSourceError, "source directory #{root} does not exist" unless File.directory?(root)
+
+      root
+    end
+
+    def model2vec_config?(config)
+      config["model_type"].to_s == "model2vec" || Array(config["architectures"]).include?("StaticModel")
+    end
+  end
+end

data/lib/static_embeddings/model.rb

@@ -4,39 +4,41 @@ module StaticEmbeddings
   class Model
     attr_reader :path
 
-    def provenance
-      raw = provenance_json
-      return {} if raw.nil?
+    if method_defined?(:max_tokens) && !method_defined?(:max_tokens_limit)
+      alias_method :max_tokens_limit, :max_tokens
 
-      JSON.parse(raw)
-    rescue JSON::ParserError, EncodingError => e
-      raise InvalidModelError, "invalid provenance JSON: #{e.message}"
+      def max_tokens
+        limit = max_tokens_limit
+        limit.zero? ? false : limit
+      end
+    end
+
+    def provenance
+      @provenance ||= parse_provenance.freeze
     end
 
     def model_id
       provenance["source_model_id"]
     end
 
-    def embed_array(text, **opts)
-      format = opts.key?(:format) ? opts[:format] : :f32
-      StaticEmbeddings.unpack(embed(text, **opts), dim, format: format).first
+    def embed_array(text, **options)
+      format = options.fetch(:format, :f32)
+      StaticEmbeddings.unpack(embed(text, **options), dim, format: format).first
     end
 
-    def embed_batch_arrays(texts, **opts)
-      format = opts.key?(:format) ? opts[:format] : :f32
-      StaticEmbeddings.unpack(embed_batch(texts, **opts), dim, format: format)
+    def embed_batch_arrays(texts, **options)
+      format = options.fetch(:format, :f32)
+      StaticEmbeddings.unpack(embed_batch(texts, **options), dim, format: format)
     end
 
-    def cosine_top_k(query_blob, matrix_blob, k, **opts)
-      raise ArgumentError, "dim: is set by the model" if opts.key?(:dim)
-
-      StaticEmbeddings.cosine_top_k(query_blob, matrix_blob, k, **opts.merge(dim: dim))
+    def cosine_top_k(query_blob, matrix_blob, k, **options)
+      reject_runtime_dim!(options)
+      StaticEmbeddings.cosine_top_k(query_blob, matrix_blob, k, **options.merge(dim: dim))
     end
 
-    def dot_top_k(query_blob, matrix_blob, k, **opts)
-      raise ArgumentError, "dim: is set by the model" if opts.key?(:dim)
-
-      StaticEmbeddings.dot_top_k(query_blob, matrix_blob, k, **opts.merge(dim: dim))
+    def dot_top_k(query_blob, matrix_blob, k, **options)
+      reject_runtime_dim!(options)
+      StaticEmbeddings.dot_top_k(query_blob, matrix_blob, k, **options.merge(dim: dim))
     end
 
     def to_s
@@ -44,5 +46,18 @@ module StaticEmbeddings
     end
 
     alias inspect to_s
+
+    private
+
+    def parse_provenance
+      raw = provenance_json
+      raw.nil? ? {} : JSON.parse(raw)
+    rescue JSON::ParserError, EncodingError => e
+      raise InvalidModelError, "invalid provenance JSON: #{e.message}"
+    end
+
+    def reject_runtime_dim!(options)
+      raise InvalidOptionError, "dim: is set by the model" if options.key?(:dim)
+    end
   end
 end

data/lib/static_embeddings/paths.rb

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

data/lib/static_embeddings/provenance.rb

@@ -0,0 +1,58 @@
+require "json"
+require "static_embeddings/format/constants"
+require "static_embeddings/bert_wordpiece"
+
+module StaticEmbeddings
+  module Provenance
+    MODEL2VEC_VERSION = "0.9.0"
+    TOKENIZERS_VERSION = "0.23.1"
+    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."
+
+    module_function
+
+    def json(model, unicode_source:)
+      JSON.generate(payload(model, unicode_source: unicode_source).sort.to_h)
+    end
+
+    def payload(model, unicode_source:)
+      dimensions = model.dimensions
+      runtime = model.runtime
+      source = model.source
+      tokenizer = model.tokenizer
+
+      data = {
+        "add_special_tokens" => runtime.add_special_tokens,
+        "config_seq_length" => source.config_seq_length,
+        "converter_version" => StaticEmbeddings::VERSION,
+        "dim" => dimensions.output,
+        "format_version" => Format::VERSION,
+        "native_dim" => dimensions.native,
+        "normalize" => runtime.normalization == Format::NORMALIZATION_L2,
+        "notes" => NOTES,
+        "oracle" => source.oracle,
+        "output_dim" => dimensions.output,
+        "reference_impl" => source.oracle,
+        "reference_max_tokens" => runtime.max_tokens,
+        "reference_tokenizers_version" => TOKENIZERS_VERSION,
+        "source_family" => source.family,
+        "source_files_sha256" => source.files_sha256,
+        "source_model_id" => source.model,
+        "source_revision" => source.revision,
+        "tokenizer_class" => source.tokenizer_class,
+        "tokenizer_profile" => tokenizer.fetch(:tokenizer_profile, BertWordPiece::TOKENIZER_PROFILE),
+        "unicode_source" => unicode_source,
+        "unk_policy" => runtime.unk_policy == Format::UNK_DROP ? "drop" : "include",
+        "vocab_size" => model.tokens.length
+      }
+      data["trained_mrl_dims"] = dimensions.trained if dimensions.trained
+      add_model2vec_versions(data) if source.family == "model2vec"
+      data
+    end
+
+    def add_model2vec_versions(data)
+      data["reference_model2vec_version"] = MODEL2VEC_VERSION
+      data["reference_unicode_categories_version"] = UNICODE_CATEGORIES_VERSION
+    end
+  end
+end

data/lib/static_embeddings/reference.rb

@@ -1,14 +1,17 @@
+require "static_embeddings/importers"
+require "static_embeddings/format/constants"
+
 module StaticEmbeddings
   class Reference
     CJK_RANGES = [
       0x4E00..0x9FFF, 0x3400..0x4DBF, 0x20000..0x2A6DF, 0x2A700..0x2B73F,
-      0x2B740..0x2B81F, 0x2B820..0x2CEAF, 0xF900..0xFAFF, 0x2F800..0x2FA1F
+      0x2B740..0x2B81F, 0x2B920..0x2CEAF, 0xF900..0xFAFF, 0x2F800..0x2FA1F
     ].freeze
 
     ASCII_PUNCT = [33..47, 58..64, 91..96, 123..126].freeze
     ASCII_SPACES = [" ", "\t", "\n", "\r"].freeze
     RE_PUNCT = /\A\p{P}\z/
-    RE_CONTROL = /\A(?:\p{Cc}|\p{Cf}|\p{Co}|\p{Cs})\z/
+    RE_CONTROL = /\A(?:\p{Cc}|\p{Cf}|\p{Co})\z/
     RE_MN = /\A\p{Mn}\z/
     RE_WHITESPACE = /\A(?:\p{Zs}|[\u0085\u2028\u2029])\z/
 
@@ -22,41 +25,55 @@ module StaticEmbeddings
       @dim = meta.fetch(:dim)
     end
 
-    def self.from_source_dir(dir, max_tokens: Converter::REFERENCE_MAX_TOKENS)
-      tokenizer = JSON.parse(File.binread(File.join(dir, "tokenizer.json")))
-      config_path = File.join(dir, "config.json")
-      config = File.file?(config_path) ? JSON.parse(File.binread(config_path)) : {}
-      tokens = tokens_from(tokenizer)
-      tensor = Safetensors.read(File.join(dir, "model.safetensors"))[:tensors].values.first
-      rows, dim = tensor[:shape]
-      floats = tensor[:bytes].unpack("e*")
-      normalizer = tokenizer["normalizer"] || {}
+    def self.from_source_dir(dir, max_tokens: nil, dimensions: nil)
+      from_canonical(Importers.import(dir, max_tokens: max_tokens, dimensions: dimensions))
+    end
 
-      new(tokens: tokens,
-          matrix: Array.new(rows) { |i| floats[i * dim, dim] },
-          meta: meta_from(tokenizer, config, normalizer, dim, max_tokens))
+    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.tokens_from(tokenizer)
-      vocab = tokenizer.dig("model", "vocab")
-      vocab.each_with_object(Array.new(vocab.length)) { |(token, id), tokens| tokens[id] = token }
+    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(tokenizer, config, normalizer, dim, max_tokens)
-      lowercase = normalizer.fetch("lowercase", true)
+    def self.meta_from_canonical(model)
+      tokenizer = model.tokenizer
+      runtime = model.runtime
       {
-        dim: dim,
-        lowercase: lowercase,
-        strip_accents: normalizer["strip_accents"].nil? ? lowercase : normalizer["strip_accents"],
-        clean_text: normalizer.fetch("clean_text", true),
-        handle_chinese_chars: normalizer.fetch("handle_chinese_chars", true),
-        max_input_chars_per_word: tokenizer.dig("model", "max_input_chars_per_word") || 100,
-        unk_token: tokenizer.dig("model", "unk_token") || "[UNK]",
-        normalize: config.key?("normalize") ? config["normalize"] : true,
-        max_tokens: max_tokens
+        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?
+      end
+    end
+
     def normalize_text(text)
       chars = text.chars
       chars = clean_chars(chars) if @meta[:clean_text]
@@ -71,21 +88,61 @@ module StaticEmbeddings
     end
 
     def tokenize(text, max_tokens: @meta[:max_tokens])
-      ids = []
-      pre_tokenize(normalize_text(text)).each { |word| ids.concat(wordpiece(word)) }
+      ids = tokenize_with_added_tokens(text)
       max_tokens && max_tokens.positive? && ids.length > max_tokens ? ids.first(max_tokens) : ids
     end
 
     def embed(text, max_tokens: @meta[:max_tokens])
-      used = tokenize(text, max_tokens: max_tokens).reject { |id| id == unk_id }
-      return Array.new(@dim, 0.0) if used.empty?
+      ids = apply_unk_policy(tokenize(text, max_tokens: false))
+      ids = ids.first(max_tokens) if max_tokens && max_tokens.positive? && ids.length > max_tokens
+      return Array.new(@dim, 0.0) if ids.empty?
 
-      vector = pooled(used)
+      vector = pooled(ids)
       @meta[:normalize] ? l2_normalize(vector) : vector
     end
 
+    def apply_unk_policy(ids)
+      case @meta.fetch(:unk_policy)
+      when Format::UNK_DROP then ids.reject { |id| id == unk_id }
+      when Format::UNK_INCLUDE then ids
+      else raise InvalidModelError, "unknown unk_policy #{@meta[:unk_policy].inspect}"
+      end
+    end
+
     private
 
+    def tokenize_with_added_tokens(text)
+      added = @meta[:added_tokens]
+      return tokenize_plain(text) if added.empty?
+
+      ids = []
+      cursor = 0
+      binary = text.b
+      while cursor < text.bytesize
+        match = added.keys.filter_map do |literal|
+          index = binary.index(literal.b, cursor)
+          index && [index, -literal.bytesize, literal]
+        end.min
+
+        unless match
+          ids.concat(tokenize_plain(text.byteslice(cursor, text.bytesize - cursor)))
+          break
+        end
+
+        index, _, literal = match
+        ids.concat(tokenize_plain(text.byteslice(cursor, index - cursor))) if index > cursor
+        ids << added.fetch(literal)
+        cursor = index + literal.bytesize
+      end
+      ids
+    end
+
+    def tokenize_plain(text)
+      return [] if text.nil? || text.empty?
+
+      pre_tokenize(normalize_text(text)).flat_map { |word| wordpiece(word) }
+    end
+
     def clean_chars(chars)
       chars.filter_map do |ch|
         cp = ch.ord

data/lib/static_embeddings/row_prefix_payload.rb

@@ -0,0 +1,59 @@
+require "static_embeddings/errors"
+require "static_embeddings/format/constants"
+
+module StaticEmbeddings
+  class RowPrefixPayload
+    FLOAT_BYTES = 4
+
+    attr_reader :bytesize
+
+    def initialize(inner, native_dim:, output_dim:)
+      raise InvalidOptionError, "native_dim must be positive" unless native_dim.positive?
+      unless output_dim.between?(1, native_dim)
+        raise InvalidOptionError, "output_dim must be between 1 and #{native_dim}"
+      end
+
+      @inner = inner
+      @source_row_bytes = native_dim * FLOAT_BYTES
+      @output_row_bytes = output_dim * FLOAT_BYTES
+      unless (inner.bytesize % @source_row_bytes).zero?
+        raise ConversionError, "embedding payload is not a whole number of #{native_dim}-d rows"
+      end
+      @bytesize = (inner.bytesize / @source_row_bytes) * @output_row_bytes
+    end
+
+    def each_chunk
+      return enum_for(__method__) unless block_given?
+
+      remainder = Format.binary_string
+      @inner.each_chunk do |chunk|
+        data = remainder.empty? ? chunk : remainder + chunk
+        complete_rows, remainder = split_complete_rows(data)
+        yield slice_rows(complete_rows) unless complete_rows.empty?
+      end
+
+      raise ConversionError, "truncated embedding matrix while slicing dimensions" unless remainder.empty?
+    end
+
+    private
+
+    def split_complete_rows(data)
+      bytes = (data.bytesize / @source_row_bytes) * @source_row_bytes
+      complete = data.byteslice(0, bytes) || Format.binary_string
+      remainder = data.byteslice(bytes, data.bytesize - bytes) || Format.binary_string
+      remainder.force_encoding(Encoding::BINARY)
+      [complete, remainder]
+    end
+
+    def slice_rows(data)
+      rows = data.bytesize / @source_row_bytes
+      out = Format.binary_string(rows * @output_row_bytes)
+      offset = 0
+      rows.times do
+        out << data.byteslice(offset, @output_row_bytes)
+        offset += @source_row_bytes
+      end
+      out
+    end
+  end
+end