static_embeddings 0.1.4 → 1.5.6

data/lib/static_embeddings/safetensors.rb

@@ -3,61 +3,157 @@ require "json"
 module StaticEmbeddings
   module Safetensors
     MAX_HEADER_BYTES = 100 * 1024 * 1024
-    SUPPORTED_DTYPES = { "F32" => 4 }.freeze
+    SUPPORTED_DTYPES = { "F32" => 4, "F16" => 2, "BF16" => 2 }.freeze
+    CONVERT_CHUNK_ELEMENTS = 64 * 1024
+    IO_CHUNK_BYTES = 1024 * 1024
+
+    class F32Payload
+      attr_reader :bytesize
+
+      def initialize(path, tensor)
+        @path = path
+        @offset = tensor.fetch(:absolute_offset)
+        @source_bytes = tensor.fetch(:source_bytes)
+        @dtype = tensor.fetch(:dtype)
+        @bytesize = tensor.fetch(:shape).reduce(1) { |count, dim| count * dim } * 4
+      end
 
-    module_function
+      def each_chunk
+        return enum_for(__method__) unless block_given?
+
+        File.open(@path, "rb") do |io|
+          io.seek(@offset, IO::SEEK_SET)
+          remaining = @source_bytes
+          while remaining.positive?
+            read_size = [remaining, IO_CHUNK_BYTES].min
+            # F16/BF16 elements are two bytes; never split one between chunks.
+            read_size -= 1 if @dtype != "F32" && read_size.odd?
+            read_size = remaining if read_size.zero?
+            chunk = io.read(read_size)
+            unless chunk&.bytesize == read_size
+              raise ConversionError, "tensor data is truncated while streaming #{@path}"
+            end
+            chunk.force_encoding(Encoding::BINARY)
+            yield case @dtype
+                  when "F32" then chunk
+                  when "F16" then Safetensors.convert_f16_to_f32(chunk)
+                  when "BF16" then Safetensors.convert_bf16_to_f32(chunk)
+                  else
+                    raise ConversionError, "tensor dtype #{@dtype} is not supported"
+                  end
+            remaining -= read_size
+          end
+        end
+      end
+    end
 
-    def read(path)
-      data = File.binread(path)
-      header, body_offset = parse_header(data)
-      body_size = data.bytesize - body_offset
+    module_function
 
-      tensors = header.each_with_object({}) do |(name, spec), acc|
-        next if name == "__metadata__"
+    def describe(path)
+      File.open(path, "rb") do |io|
+        file_size = io.stat.size
+        header, body_offset = read_header(io, file_size)
+        body_size = file_size - body_offset
+
+        tensors = header.each_with_object({}) do |(name, spec), acc|
+          next if name == "__metadata__"
+
+          dtype = spec["dtype"]
+          shape = spec["shape"]
+          begin_off, end_off = checked_offsets(name, spec["data_offsets"], body_size)
+          expected = checked_tensor_bytes(name, dtype, shape)
+          actual = end_off - begin_off
+          if actual != expected
+            raise ConversionError,
+                  "tensor #{name}: shape #{shape.inspect} implies #{expected} bytes, offsets span #{actual}"
+          end
+
+          acc[name] = {
+            dtype: dtype,
+            shape: shape,
+            absolute_offset: body_offset + begin_off,
+            source_bytes: actual
+          }
+        end
+
+        { metadata: header["__metadata__"] || {}, tensors: tensors }
+      end
+    end
 
-        acc[name] = tensor_from(data, body_offset, body_size, name, spec)
+    def read(path)
+      description = describe(path)
+      tensors = description.fetch(:tensors).transform_values do |tensor|
+        File.open(path, "rb") do |io|
+          io.seek(tensor.fetch(:absolute_offset), IO::SEEK_SET)
+          bytes = io.read(tensor.fetch(:source_bytes))
+          unless bytes&.bytesize == tensor.fetch(:source_bytes)
+            raise ConversionError, "tensor data is truncated while reading #{path}"
+          end
+          tensor.merge(bytes: bytes.force_encoding(Encoding::BINARY))
+        end
       end
+      { metadata: description.fetch(:metadata), tensors: tensors }
+    end
 
-      { metadata: header["__metadata__"] || {}, tensors: tensors }
+    def f32_payload(path, tensor)
+      F32Payload.new(path, tensor)
     end
 
     def write(path, name, shape, floats)
       body = floats.pack("e*")
-      header = JSON.generate(name => { "dtype" => "F32", "shape" => shape, "data_offsets" => [0, body.bytesize] })
+      write_raw(path, name, shape, "F32", body)
+    end
+
+    def write_raw(path, name, shape, dtype, body)
+      checked_tensor_bytes(name, dtype, shape).tap do |expected|
+        raise ArgumentError, "body size does not match #{dtype} shape" unless body.bytesize == expected
+      end
+
+      header = JSON.generate(name => { "dtype" => dtype, "shape" => shape, "data_offsets" => [0, body.bytesize] })
       padded = header << (" " * ((8 - (header.bytesize % 8)) % 8))
-      File.binwrite(path, [padded.bytesize].pack("Q<") << padded << body)
+      File.open(path, "wb") do |io|
+        io.write([padded.bytesize].pack("Q<"))
+        io.write(padded)
+        io.write(body)
+      end
+    end
+
+    def f32_bytes(tensor)
+      case tensor.fetch(:dtype)
+      when "F32"
+        tensor.fetch(:bytes)
+      when "F16"
+        convert_f16_to_f32(tensor.fetch(:bytes))
+      when "BF16"
+        convert_bf16_to_f32(tensor.fetch(:bytes))
+      else
+        raise ConversionError, "tensor dtype #{tensor[:dtype]} is not supported"
+      end
     end
 
-    def parse_header(data)
-      raise ConversionError, "safetensors file is shorter than its length prefix" if data.bytesize < 8
+    def read_header(io, file_size)
+      raise ConversionError, "safetensors file is shorter than its length prefix" if file_size < 8
+
+      prefix = io.read(8)
+      raise ConversionError, "safetensors file is shorter than its length prefix" unless prefix&.bytesize == 8
 
-      header_len = data.byteslice(0, 8).unpack1("Q<")
+      header_len = prefix.unpack1("Q<")
       unless header_len.positive? && header_len <= MAX_HEADER_BYTES
         raise ConversionError, "implausible safetensors header length #{header_len}"
       end
 
       body_offset = 8 + header_len
-      raise ConversionError, "safetensors header runs past end of file" if body_offset > data.bytesize
+      raise ConversionError, "safetensors header runs past end of file" if body_offset > file_size
 
-      raw_header = data.byteslice(8, header_len)
+      raw_header = io.read(header_len)
+      unless raw_header&.bytesize == header_len
+        raise ConversionError, "safetensors header runs past end of file"
+      end
       raise ConversionError, "safetensors header is not a JSON object" unless raw_header.lstrip.start_with?("{")
 
       [JSON.parse(raw_header), body_offset]
-    end
-
-    def tensor_from(data, body_offset, body_size, name, spec)
-      dtype = spec["dtype"]
-      shape = spec["shape"]
-      begin_off, end_off = checked_offsets(name, spec["data_offsets"], body_size)
-      expected = checked_tensor_bytes(name, dtype, shape)
-      actual = end_off - begin_off
-
-      if actual != expected
-        raise ConversionError,
-              "tensor #{name}: shape #{shape.inspect} implies #{expected} bytes, offsets span #{actual}"
-      end
-
-      { dtype: dtype, shape: shape, bytes: data.byteslice(body_offset + begin_off, actual) }
+    rescue JSON::ParserError => e
+      raise ConversionError, "invalid safetensors JSON header: #{e.message}"
     end
 
     def checked_offsets(name, offsets, body_size)
@@ -76,12 +172,60 @@ module StaticEmbeddings
 
     def checked_tensor_bytes(name, dtype, shape)
       element_size = SUPPORTED_DTYPES[dtype]
-      raise ConversionError, "tensor #{name}: dtype #{dtype} is not supported (F32 only)" unless element_size
+      unless element_size
+        raise ConversionError, "tensor #{name}: dtype #{dtype} is not supported (expected F32, F16, or BF16)"
+      end
       unless shape.is_a?(Array) && shape.all? { |dim| dim.is_a?(Integer) && dim >= 0 }
         raise ConversionError, "tensor #{name}: malformed shape #{shape.inspect}"
       end
 
-      shape.reduce(1, :*) * element_size
+      shape.reduce(1) { |count, dim| count * dim } * element_size
+    end
+
+    def convert_f16_to_f32(bytes)
+      out = String.new(capacity: (bytes.bytesize / 2) * 4, encoding: Encoding::BINARY)
+      offset = 0
+      chunk_bytes = CONVERT_CHUNK_ELEMENTS * 2
+      while offset < bytes.bytesize
+        chunk = bytes.byteslice(offset, [chunk_bytes, bytes.bytesize - offset].min)
+        floats = if StaticEmbeddings.respond_to?(:decode_f16)
+                   StaticEmbeddings.decode_f16(chunk)
+                 else
+                   chunk.unpack("v*").map { |bits| half_to_float(bits) }
+                 end
+        out << floats.pack("e*")
+        offset += chunk.bytesize
+      end
+      out
+    end
+
+    def convert_bf16_to_f32(bytes)
+      out = String.new(capacity: (bytes.bytesize / 2) * 4, encoding: Encoding::BINARY)
+      offset = 0
+      chunk_bytes = CONVERT_CHUNK_ELEMENTS * 2
+      while offset < bytes.bytesize
+        chunk = bytes.byteslice(offset, [chunk_bytes, bytes.bytesize - offset].min)
+        words = chunk.unpack("v*")
+        out << words.map { |bits| bits << 16 }.pack("V*")
+        offset += chunk.bytesize
+      end
+      out
+    end
+
+    def half_to_float(bits)
+      sign = (bits >> 15) & 1
+      exponent = (bits >> 10) & 0x1F
+      fraction = bits & 0x3FF
+
+      value =
+        if exponent.zero?
+          fraction.zero? ? 0.0 : Math.ldexp(fraction.to_f, -24)
+        elsif exponent == 0x1F
+          fraction.zero? ? Float::INFINITY : Float::NAN
+        else
+          Math.ldexp(1.0 + fraction.to_f / 1024.0, exponent - 15)
+        end
+      sign.zero? ? value : -value
     end
   end
 end

data/lib/static_embeddings/version.rb

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

data/lib/static_embeddings.rb

@@ -1,27 +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/unicode_tables"
-require "static_embeddings/safetensors"
-require "static_embeddings/converter"
-require "static_embeddings/reference"
+require "static_embeddings/format/constants"
 require "static_embeddings/model"
 
 module StaticEmbeddings
@@ -29,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
@@ -58,61 +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: Converter::REFERENCE_MAX_TOKENS)
-      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."

data/tools/check_model2vec_parity.rb

@@ -4,8 +4,7 @@ require "static_embeddings"
 
 options = {
   min_cosine: 1.0 - 1e-6,
-  max_abs: 1e-5,
-  ids: true
+  max_abs: 1e-5
 }
 
 parser = OptionParser.new do |opts|
@@ -13,7 +12,6 @@ parser = OptionParser.new do |opts|
   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 }
-  opts.on("--[no-]ids") { |value| options[:ids] = value }
 end
 parser.parse!(ARGV)
 unless options[:model] && options[:oracle]
@@ -22,12 +20,18 @@ end
 
 model = StaticEmbeddings.load(options[:model], verify: true)
 payload = JSON.parse(File.read(options[:oracle], encoding: "UTF-8"))
-rows = payload.is_a?(Hash) ? payload.fetch("rows") : payload
-oracle_max_length = payload.is_a?(Hash) ? payload["max_length"] : nil
+abort "unsupported oracle schema #{payload["schema_version"].inspect}" unless payload["schema_version"] == 2
 
-if oracle_max_length && oracle_max_length != model.max_tokens
-  warn "WARNING oracle max_length=#{oracle_max_length} but model.max_tokens=#{model.max_tokens}"
+reference = payload.fetch("reference")
+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}"
 end
+if reference["unk_token_id"] && Integer(reference["unk_token_id"]) != model.unk_id
+  abort "oracle unk_token_id=#{reference["unk_token_id"]} but model.unk_id=#{model.unk_id}"
+end
+
+rows = payload.fetch("rows")
 
 def dot(a, b)
   a.zip(b).sum { |x, y| x * y }
@@ -37,71 +41,102 @@ 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 = []
+usable_failures = []
+vector_failures = []
+invariant_failures = []
+intentional_deviations = []
 min_cosine = 1.0
 max_abs_all = 0.0
-vector_failures = []
-id_failures = []
-id_rows_checked = 0
+vectors_checked = 0
 
 rows.each_with_index do |row, i|
   text = row.fetch("text")
-  ref = row.fetch("vector")
-  ref_ids = row["token_ids"]
-
+  label = row.fetch("label", i.to_s)
   problems = []
 
-  if options[:ids] && ref_ids
-    id_rows_checked += 1
-    got_ids = model.tokenize(text)
-    if got_ids != ref_ids
-      id_failures << i
-      first = got_ids.zip(ref_ids).index { |a, b| a != b } || [got_ids.length, ref_ids.length].min
-      problems << "ids differ at #{first} (got #{got_ids.length}, ref #{ref_ids.length})"
-    end
+  expected_raw = row.fetch("hf_raw_token_ids")
+  got_raw = model.tokenize(text)
+  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 = model.embed(text).unpack("e*")
-  max_abs = ref.zip(got).map { |a, b| (a - b).abs }.max || 0.0
-  ref_zero = ref.all?(&:zero?)
-  got_zero = got.all?(&:zero?)
-  cosine =
-    if ref_zero && got_zero
-      1.0
-    elsif ref_zero || got_zero
-      0.0
-    else
-      dot(ref, got) / (norm(ref) * norm(got))
-    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)
+               else
+                 full_raw.reject { |id| id == model.unk_id }.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})"
+  end
 
-  min_cosine = [min_cosine, cosine].min
-  max_abs_all = [max_abs_all, max_abs].max
+  got_vector_blob = model.embed(text)
+  pooled_blob = model.embed_token_ids(full_raw)
+  unless got_vector_blob == pooled_blob
+    invariant_failures << i
+    problems << "embed(text) != embed_token_ids(unbounded tokenize(text))"
+  end
 
-  unless cosine >= options[:min_cosine] && max_abs <= options[:max_abs]
-    vector_failures << i
-    problems << "vector out of tolerance"
+  model2vec_ids = row.fetch("model2vec_token_ids")
+  declared_deviation = row.fetch("model2vec_character_pretruncate_changes_ids")
+  actual_deviation = model2vec_ids != expected_usable
+  if actual_deviation != declared_deviation
+    problems << "oracle character-pretruncate flag is inconsistent"
+    usable_failures << i unless usable_failures.include?(i)
+  elsif actual_deviation
+    intentional_deviations << i
+  else
+    reference_vector = row.fetch("model2vec_vector")
+    got_vector = got_vector_blob.unpack("e*")
+    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("vector out of tolerance cos=%.10f max_abs=%.8g", cosine, max_abs)
+    end
   end
 
   status = problems.empty? ? "ok" : "FAIL"
-  label = text.bytesize > 64 ? "#{text[0, 32].inspect}...(#{text.bytesize}B)" : text.inspect
-  line = format("%s idx=%02d cos=%.10f max_abs=%.8g bytes=%d text=%s",
-                status, i, cosine, max_abs, text.bytesize, label)
-  line += "  [#{problems.join('; ')}]" unless problems.empty?
-  puts line
+  suffix = actual_deviation ? " intentional-character-pretruncate-deviation" : ""
+  puts "#{status} idx=#{format('%03d', i)} label=#{label.inspect} bytes=#{text.bytesize}#{suffix}" +
+       (problems.empty? ? "" : " [#{problems.join('; ')}]")
 end
 
 puts "rows=#{rows.length}"
-puts "id_rows_checked=#{id_rows_checked}"
-puts "min_cosine=#{min_cosine}"
-puts "max_abs_all=#{max_abs_all}"
-puts "token_id_failures=#{id_failures.inspect}"
+puts "vectors_checked=#{vectors_checked}"
+puts "intentional_character_pretruncate_deviations=#{intentional_deviations.length}"
+puts "min_cosine=#{min_cosine}" if vectors_checked.positive?
+puts "max_abs_all=#{max_abs_all}" if vectors_checked.positive?
+puts "raw_token_id_failures=#{raw_failures.inspect}"
+puts "usable_token_id_failures=#{usable_failures.inspect}"
+puts "embed_invariant_failures=#{invariant_failures.inspect}"
 puts "vector_failures=#{vector_failures.inspect}"
 
-if id_rows_checked.zero? && options[:ids]
-  warn "WARNING oracle has no token_ids; regenerate it with the current tools/model2vec_oracle.py"
-end
-
-unless id_failures.empty? && vector_failures.empty?
-  abort "parity failed: token ids #{id_failures.inspect}, vectors #{vector_failures.inspect}"
+unless raw_failures.empty? && usable_failures.empty? && invariant_failures.empty? && vector_failures.empty?
+  abort "parity failed"
 end
 
-puts "parity OK"
+puts "corpus parity OK (#{rows.length}/#{rows.length}); intentional Model2Vec character pre-truncation deviations are reported separately"