static_embeddings 0.1.4 → 0.1.5

data/lib/static_embeddings/format.rb

@@ -3,7 +3,7 @@ require "digest"
 module StaticEmbeddings
   module Format
     MAGIC = "SEMBv1\0\0"
-    VERSION = 2
+    VERSION = 3
     HEADER_SIZE = 320
     ALIGNMENT = 64
 
@@ -12,7 +12,7 @@ module StaticEmbeddings
     POOLING_MEAN = 1
     NORMALIZATION_NONE = 0
     NORMALIZATION_L2 = 1
-    TRUNCATE_IDS_BEFORE_POOLING = 1
+    TRUNCATE_USABLE_IDS_BEFORE_POOLING = 2
     UNK_INCLUDE = 0
     UNK_DROP = 1
     EMPTY_ZERO_VECTOR = 0
@@ -26,6 +26,14 @@ module StaticEmbeddings
     CHECKSUM_OFFSET = 240
     CHECKSUM_SIZE = 32
     MAX_PROBE_OFFSET = 304
+    ADDED_TOKEN_MASK_OFFSET = 308
+
+    ADDED_PAD  = 1 << 0
+    ADDED_UNK  = 1 << 1
+    ADDED_CLS  = 1 << 2
+    ADDED_SEP  = 1 << 3
+    ADDED_MASK = 1 << 4
+    ADDED_TOKEN_MASK_ALL = ADDED_PAD | ADDED_UNK | ADDED_CLS | ADDED_SEP | ADDED_MASK
 
     SECTION_FIELDS = {
       vocab_strings: 128,
@@ -44,7 +52,7 @@ module StaticEmbeddings
       28 => TOKENIZER_BERT_WORDPIECE_V1,
       32 => DTYPE_F32,
       36 => POOLING_MEAN,
-      48 => TRUNCATE_IDS_BEFORE_POOLING,
+      48 => TRUNCATE_USABLE_IDS_BEFORE_POOLING,
       108 => HASH_SEED
     }.freeze
 
@@ -102,7 +110,7 @@ module StaticEmbeddings
     def write(path:, meta:, tokens:, matrix:, norm_tables:, provenance:)
       hash_size, strings, hash_blob, max_probe = build_hash_table(tokens)
       root_trie, continuation_trie = build_wordpiece_tries(tokens, meta.fetch(:subword_prefix))
-      sections, body = build_body(
+      payloads = {
         vocab_strings: strings,
         vocab_hash: hash_blob,
         embeddings: matrix,
@@ -110,15 +118,48 @@ module StaticEmbeddings
         provenance: provenance,
         root_trie: root_trie,
         continuation_trie: continuation_trie
-      )
-
+      }
+      sections, file_size = layout_sections(payloads)
       header = build_header(meta, tokens.length, hash_size, max_probe, sections)
-      file = header << body
-      digest = Digest::SHA256.digest(file)
-      file[CHECKSUM_OFFSET, CHECKSUM_SIZE] = digest
 
-      File.binwrite(path, file)
-      { bytes: file.bytesize, sha256: digest.unpack1("H*"), hash_table_size: hash_size }
+      digest = Digest::SHA256.new
+      File.open(path, "wb") do |io|
+        io.write(header)
+        digest << header
+        offset = HEADER_SIZE
+
+        payloads.each do |name, payload|
+          target = sections.fetch(name).first
+          padding = target - offset
+          if padding.positive?
+            zeros = "\0".b * padding
+            io.write(zeros)
+            digest << zeros
+          end
+          write_payload(io, digest, payload)
+          offset = target + payload.bytesize
+        end
+      end
+
+      checksum = digest.digest
+      File.open(path, "r+b") do |io|
+        io.seek(CHECKSUM_OFFSET, IO::SEEK_SET)
+        io.write(checksum)
+      end
+
+      { bytes: file_size, sha256: checksum.unpack1("H*"), hash_table_size: hash_size }
+    end
+
+    def write_payload(io, digest, payload)
+      if payload.respond_to?(:each_chunk)
+        payload.each_chunk do |chunk|
+          io.write(chunk)
+          digest << chunk
+        end
+      else
+        io.write(payload)
+        digest << payload
+      end
     end
 
     def verify(path)
@@ -240,20 +281,16 @@ module StaticEmbeddings
       end
     end
 
-    def build_body(payloads)
-      body = binary_string
+    def layout_sections(payloads)
+      offset = HEADER_SIZE
       sections = {}
       payloads.each do |name, payload|
-        align_body!(body)
-        sections[name] = [HEADER_SIZE + body.bytesize, payload.bytesize]
-        body << payload
+        padding = (ALIGNMENT - (offset % ALIGNMENT)) % ALIGNMENT
+        offset += padding
+        sections[name] = [offset, payload.bytesize]
+        offset += payload.bytesize
       end
-      [sections, body]
-    end
-
-    def align_body!(body)
-      padding = (ALIGNMENT - ((HEADER_SIZE + body.bytesize) % ALIGNMENT)) % ALIGNMENT
-      body << "\0".b * padding if padding.positive?
+      [sections, offset]
     end
 
     def build_header(meta, vocab_size, hash_size, max_probe, sections)
@@ -268,6 +305,7 @@ module StaticEmbeddings
       put_u32(header, 104, hash_size)
       put_u32(header, MAX_TOKEN_CHARS_OFFSET, meta.fetch(:max_token_chars))
       put_u32(header, MAX_PROBE_OFFSET, max_probe)
+      put_u32(header, ADDED_TOKEN_MASK_OFFSET, meta.fetch(:added_token_mask, 0))
       put_prefix(header, meta.fetch(:subword_prefix))
       SECTION_FIELDS.each { |name, field| put_section(header, field, sections.fetch(name)) }
 

data/lib/static_embeddings/paths.rb

@@ -11,7 +11,24 @@ 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 ArgumentError, "model_id must not be empty" if id.empty?
+      raise ArgumentError, "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
+
+      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
+      path
     end
 
     def builtin_path(name)

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.4"
+  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