static_embeddings 0.1.3 → 0.1.5

data/lib/static_embeddings/reference.rb

@@ -1,14 +1,18 @@
+require "json"
+require "static_embeddings/converter"
+require "static_embeddings/safetensors"
+
 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/
 
@@ -29,7 +33,7 @@ module StaticEmbeddings
       tokens = tokens_from(tokenizer)
       tensor = Safetensors.read(File.join(dir, "model.safetensors"))[:tensors].values.first
       rows, dim = tensor[:shape]
-      floats = tensor[:bytes].unpack("e*")
+      floats = Safetensors.f32_bytes(tensor).unpack("e*")
       normalizer = tokenizer["normalizer"] || {}
 
       new(tokens: tokens,
@@ -53,10 +57,21 @@ module StaticEmbeddings
         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
+        max_tokens: max_tokens,
+        added_tokens: supported_added_tokens(tokenizer)
       }
     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
+
+        out[content] = Integer(token.fetch("id"))
+      end
+    end
+
     def normalize_text(text)
       chars = text.chars
       chars = clean_chars(chars) if @meta[:clean_text]
@@ -71,13 +86,13 @@ 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 }
+      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?
 
       vector = pooled(used)
@@ -86,6 +101,38 @@ module StaticEmbeddings
 
     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/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.3"
+  VERSION = "0.1.5"
 end

data/lib/static_embeddings.rb

@@ -18,10 +18,6 @@ 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/model"
 
 module StaticEmbeddings
@@ -64,7 +60,9 @@ module StaticEmbeddings
       load(path, verify: verify)
     end
 
-    def convert(source_dir, output_path:, model_id: nil, max_tokens: Converter::REFERENCE_MAX_TOKENS)
+    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

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,98 @@ 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 = full_raw.reject { |id| id == model.unk_id }.first(max_length)
+  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"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: static_embeddings
 version: !ruby/object:Gem::Version
-  version: 0.1.3
+  version: 0.1.5
 platform: ruby
 authors:
 - Roman Haydarov