static_embeddings 0.1.5 → 1.5.6

data/lib/static_embeddings/importers/support.rb

@@ -0,0 +1,111 @@
+require "digest"
+require "json"
+require "static_embeddings/errors"
+require "static_embeddings/format/constants"
+require "static_embeddings/safetensors"
+require "static_embeddings/row_prefix_payload"
+
+module StaticEmbeddings
+  module Importers
+    module Support
+      module_function
+
+      def load_json(path, optional: false)
+        return nil if optional && !File.file?(path)
+        raise InvalidSourceError, "missing #{File.basename(path)} in #{File.dirname(path)}" unless File.file?(path)
+
+        JSON.parse(File.binread(path))
+      rescue JSON::ParserError => e
+        raise InvalidSourceError, "invalid JSON in #{path}: #{e.message}"
+      end
+
+      def load_object(path, optional: false)
+        value = load_json(path, optional: optional)
+        return nil if value.nil?
+        raise InvalidSourceError, "#{path} must contain a JSON object" unless value.is_a?(Hash)
+
+        value
+      end
+
+      def digest_files(paths, root: nil)
+        paths.each_with_object({}) do |path, digests|
+          next unless File.file?(path)
+
+          digests[digest_key(path, root)] = Digest::SHA256.file(path).hexdigest
+        end
+      end
+
+      def extract_matrix(path, vocab_size)
+        raise InvalidSourceError, "missing model.safetensors in #{File.dirname(path)}" unless File.file?(path)
+
+        name, tensor = sole_matrix_tensor(Safetensors.describe(path).fetch(:tensors))
+        rows, dim = tensor.fetch(:shape)
+        unless 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}"
+        end
+
+        [Safetensors.f32_payload(path, tensor), dim]
+      end
+
+      def resolve_max_tokens(requested, default)
+        return default if requested.nil?
+        return 0 if requested == false || requested == :unlimited || requested == 0
+
+        value = Integer(requested)
+        unless value.between?(1, Format::UINT32_MAX)
+          raise InvalidOptionError, "max_tokens must be 1..#{Format::UINT32_MAX}, false, or :unlimited"
+        end
+        value
+      end
+
+      def resolve_dimensions(requested, native_dim)
+        return native_dim if requested.nil?
+
+        value = Integer(requested)
+        unless value.between?(1, native_dim)
+          raise InvalidOptionError, "dimensions must be between 1 and native_dim #{native_dim}, got #{value}"
+        end
+        value
+      end
+
+      def normalize_mrl_dims(value, native_dim:)
+        return nil if value.nil?
+
+        dims = value.is_a?(Array) ? value : value.to_s.split(",").map(&:strip)
+        dims = dims.map { |item| Integer(item) }
+        if dims.empty? || dims.any? { |dim| !dim.between?(1, native_dim) }
+          raise InvalidOptionError, "trained_mrl_dims must contain dimensions between 1 and #{native_dim}"
+        end
+        dims.uniq.freeze
+      end
+
+      def slice_matrix(payload, native_dim, output_dim)
+        return payload if output_dim == native_dim
+
+        RowPrefixPayload.new(payload, native_dim: native_dim, output_dim: output_dim)
+      end
+
+      def sole_matrix_tensor(tensors)
+        matrices = tensors.select { |_, tensor| tensor.fetch(: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}); expected exactly one"
+      end
+
+      def digest_key(path, root)
+        return File.basename(path) unless root
+
+        root_real = File.realpath(root)
+        real = File.realpath(path)
+        prefix = root_real.end_with?(File::SEPARATOR) ? root_real : "#{root_real}#{File::SEPARATOR}"
+        real.start_with?(prefix) ? real.delete_prefix(prefix) : File.basename(real)
+      end
+    end
+  end
+end

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

@@ -12,29 +12,25 @@ module StaticEmbeddings
 
     def model_path(model_id, env = ENV)
       id = model_id.to_s
-      raise ArgumentError, "model_id must not be empty" if id.empty?
-      raise ArgumentError, "model_id contains a NUL byte" if id.include?("\0")
+      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)
-      if normalized.start_with?("/") ||
-         parts.any? { |part| part.empty? || part == "." || part == ".." || part.include?(":") }
-        raise ArgumentError, "model_id must be a relative slash-separated identifier"
-      end
+      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}"
-      unless path.start_with?(prefix)
-        raise ArgumentError, "model_id escapes the model cache"
-      end
+      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,6 +1,5 @@
-require "json"
-require "static_embeddings/converter"
-require "static_embeddings/safetensors"
+require "static_embeddings/importers"
+require "static_embeddings/format/constants"
 
 module StaticEmbeddings
   class Reference
@@ -26,49 +25,52 @@ 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 = Safetensors.f32_bytes(tensor).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,
-        added_tokens: supported_added_tokens(tokenizer)
+        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.supported_added_tokens(tokenizer)
-      allowed = Converter::STANDARD_SPECIAL_TOKENS
-      (tokenizer["added_tokens"] || []).each_with_object({}) do |token, out|
-        content = token["content"]
-        next unless token["special"] && allowed.key?(content) && token["normalized"] == false
+    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?
 
-        out[content] = Integer(token.fetch("id"))
+        id = model.tokens.index(content)
+        tokens[content] = id unless id.nil?
       end
     end
 
@@ -91,14 +93,22 @@ module StaticEmbeddings
     end
 
     def embed(text, max_tokens: @meta[:max_tokens])
-      used = tokenize(text, max_tokens: false).reject { |id| id == unk_id }
-      used = used.first(max_tokens) if max_tokens && max_tokens.positive? && used.length > max_tokens
-      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)

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

data/lib/static_embeddings/version.rb

@@ -1,3 +1,3 @@
 module StaticEmbeddings
-  VERSION = "0.1.5"
+  VERSION = "1.5.6"
 end

data/lib/static_embeddings.rb

@@ -1,23 +1,20 @@
-require "json"
 require "static_embeddings/version"
 
 begin
   require "static_embeddings/static_embeddings"
 rescue LoadError
   ext_dir = File.expand_path("static_embeddings", __dir__)
-  so_path = %w[.so .bundle].lazy.map { |ext| File.join(ext_dir, "static_embeddings#{ext}") }
-                                .find { |path| File.file?(path) }
-  unless so_path
-    raise LoadError, "Could not find the compiled StaticEmbeddings extension. " \
-                     "Run: bundle exec rake compile"
+  extension = %w[.so .bundle].lazy.map { |suffix| File.join(ext_dir, "static_embeddings#{suffix}") }
+                              .find { |path| File.file?(path) }
+  unless extension
+    raise LoadError, "Could not find the compiled StaticEmbeddings extension. Run: bundle exec rake compile"
   end
-
-  require so_path
+  require extension
 end
 
 require "static_embeddings/errors"
 require "static_embeddings/paths"
-require "static_embeddings/format"
+require "static_embeddings/format/constants"
 require "static_embeddings/model"
 
 module StaticEmbeddings
@@ -25,11 +22,7 @@ module StaticEmbeddings
     def load(path, verify: false)
       expanded = File.expand_path(path.to_s)
       raise ModelNotFound, "no model at #{expanded}" unless File.file?(expanded)
-
-      if verify
-        result = Format.verify(expanded)
-        raise InvalidModelError, "checksum mismatch for #{expanded}" unless result[:ok]
-      end
+      raise InvalidModelError, "checksum mismatch for #{expanded}" if verify && !self.verify(expanded)[:ok]
 
       Model.new(expanded)
     end
@@ -54,63 +47,44 @@ module StaticEmbeddings
       path = model_path(model_id)
       unless File.file?(path)
         raise ModelNotFound,
-              "model #{model_id.inspect} is not installed. " \
-              "Convert it first: static_embeddings convert <hf-dir> --id #{model_id}"
+              "model #{model_id.inspect} is not installed. Convert it first: " \
+              "static_embeddings convert <hf-dir> --id #{model_id}"
       end
       load(path, verify: verify)
     end
 
-    def convert(source_dir, output_path:, model_id: nil, max_tokens: nil)
-      require "static_embeddings/converter"
-      max_tokens = Converter::REFERENCE_MAX_TOKENS if max_tokens.nil?
-      converter = Converter.new(source_dir)
-      converter.convert(output_path: output_path, model_id: model_id, max_tokens: max_tokens)
-      converter.report
+    def convert(source_dir, output_path:, model_id: nil, max_tokens: nil, dimensions: nil,
+                source_revision: nil, trained_mrl_dims: nil)
+      require "static_embeddings/conversion"
+      Conversion.call(
+        source_dir,
+        output_path: output_path,
+        model_id: model_id,
+        max_tokens: max_tokens,
+        dimensions: dimensions,
+        source_revision: source_revision,
+        trained_mrl_dims: trained_mrl_dims
+      )
     end
 
     def verify(path)
-      Format.verify(File.expand_path(path.to_s))
+      require "static_embeddings/format/verifier"
+      Format::Verifier.call(File.expand_path(path.to_s))
     end
 
     def unpack(blob, dim, format: :f32)
-      raise ArgumentError, "dim must be positive" unless dim.to_i.positive?
-
-      floats =
-        case normalize_format(format)
-        when :f32
-          raise ArgumentError, "f32 blob byte size must be a multiple of 4" unless (blob.bytesize % 4).zero?
-
-          blob.unpack("e*")
-        when :f16
-          raise ArgumentError, "f16 blob byte size must be a multiple of 2" unless (blob.bytesize % 2).zero?
-
-          decode_f16(blob)
-        end
-
-      raise ArgumentError, "blob is not a multiple of dim" unless (floats.length % dim).zero?
-
-      floats.each_slice(dim).to_a
+      require "static_embeddings/codec"
+      Codec.unpack(blob, dim, format: format)
     end
 
     def pack(rows, format: :f32)
-      flat = rows.first.is_a?(Array) ? rows.flatten(1) : rows
-
-      case normalize_format(format)
-      when :f32 then flat.map(&:to_f).pack("e*")
-      when :f16 then encode_f16(flat.map(&:to_f))
-      end
+      require "static_embeddings/codec"
+      Codec.pack(rows, format: format)
     end
 
     def normalize_format(format)
-      case format&.to_sym
-      when nil, :f32, :float32
-        :f32
-      when :f16, :float16
-        :f16
-      else
-        raise ArgumentError, "unsupported embedding format #{format.inspect} (expected :f32 or :f16)"
-      end
+      require "static_embeddings/codec"
+      Codec.normalize_format(format)
     end
-
   end
 end

data/static_embeddings.gemspec

@@ -7,8 +7,8 @@ Gem::Specification.new do |spec|
   spec.email   = ["romnhajdarov@gmail.com"]
 
   spec.summary     = "Fast local text embeddings for Ruby — no ONNX, no Rust, no network"
-  spec.description = "A small C-extension runtime for Model2Vec-style static embedding " \
-                     "models. Models are converted offline into a flat mmap-able .semb " \
+  spec.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."