static_embeddings 0.1.4 → 1.5.6

data/lib/static_embeddings/cli.rb

@@ -1,16 +1,16 @@
+require "fileutils"
 require "json"
 require "optparse"
-require "fileutils"
 
 module StaticEmbeddings
-  class CLI
+  module CLI
     COMMANDS = {
-      "convert" => :convert,
-      "verify" => :verify,
-      "inspect" => :inspect_model,
-      "tokenize" => :tokenize,
-      "embed" => :embed,
-      "cache-path" => :cache_path,
+      "convert" => :convert_command,
+      "verify" => :verify_command,
+      "inspect" => :inspect_command,
+      "tokenize" => :tokenize_command,
+      "embed" => :embed_command,
+      "cache-path" => :cache_path_command,
       "help" => :usage,
       "-h" => :usage,
       "--help" => :usage,
@@ -20,10 +20,14 @@ module StaticEmbeddings
     HELP = <<~TEXT
       static_embeddings <command> [options]
 
-        convert SOURCE_DIR    Convert a HuggingFace/Model2Vec directory to .semb
+        convert SOURCE_DIR    Convert Model2Vec or Sentence Transformers StaticEmbedding to .semb
           --out PATH          Output file (default: <cache>/models/<id>.semb)
           --id ID             Model id recorded in provenance
-          --max-tokens N      Truncation limit baked into the file (default 512)
+          --max-tokens N      Model2Vec default: 512; Sentence Transformers default: unlimited
+          --max-tokens unlimited
+          --dimensions N      Keep the first N dims (Matryoshka prefix slice)
+          --revision SHA      Source revision recorded in provenance
+          --trained-mrl-dims  1024,512,256,... training dims recorded in provenance
 
         verify PATH           Recompute the SHA-256 embedded in the header
         inspect PATH          Print header fields and provenance
@@ -32,66 +36,78 @@ module StaticEmbeddings
         cache-path            Print the model cache directory
     TEXT
 
-    def self.run(argv)
-      new.run(argv)
-    end
+    module_function
 
     def run(argv)
-      method = COMMANDS[argv.shift]
-      return unknown unless method
+      command = COMMANDS[argv.shift]
+      return unknown_command unless command
 
-      send(method, argv)
-    rescue StaticEmbeddings::Error => e
+      public_send(command, argv)
+    rescue StaticEmbeddings::Error, ArgumentError, OptionParser::ParseError => e
       warn "#{e.class.name.split('::').last}: #{e.message}"
       1
     end
 
-    private
-
     def usage(*)
       puts HELP
       0
     end
 
-    def unknown
+    def unknown_command
       warn "unknown command"
       usage
       1
     end
 
-    def cache_path(*)
+    def cache_path_command(*)
       puts StaticEmbeddings.cache_dir
       0
     end
 
-    def convert(argv)
-      options = parse_convert_options(argv)
+    def convert_command(argv)
+      options = convert_options(argv)
       source = required_arg(argv, "usage: static_embeddings convert SOURCE_DIR [--out PATH]")
       model_id = options[:id] || File.basename(File.expand_path(source))
-      out = options[:out] || StaticEmbeddings.model_path(model_id)
-      FileUtils.mkdir_p(File.dirname(out))
+      output = options[:out] || StaticEmbeddings.model_path(model_id)
+      FileUtils.mkdir_p(File.dirname(output))
 
-      report = StaticEmbeddings.convert(source, output_path: out, model_id: model_id,
-                                                max_tokens: options[:max_tokens])
-      puts conversion_report(out, report, options[:max_tokens])
+      report = StaticEmbeddings.convert(source, output_path: output, model_id: model_id, **conversion_options(options))
+      puts conversion_report(output, report)
       0
     end
 
-    def parse_convert_options(argv)
-      options = { max_tokens: Converter::REFERENCE_MAX_TOKENS }
-      OptionParser.new do |parser|
-        parser.on("--out PATH") { |value| options[:out] = value }
-        parser.on("--id ID") { |value| options[:id] = value }
-        parser.on("--max-tokens N", Integer) { |value| options[:max_tokens] = value }
-      end.parse!(argv)
-      options
+    def convert_options(argv)
+      {}.tap do |options|
+        OptionParser.new do |parser|
+          parser.on("--out PATH") { |value| options[:out] = value }
+          parser.on("--id ID") { |value| options[:id] = value }
+          parser.on("--max-tokens N") { |value| options[:max_tokens] = parse_max_tokens(value) }
+          parser.on("--dimensions N", Integer) { |value| options[:dimensions] = value }
+          parser.on("--revision SHA") { |value| options[:source_revision] = value }
+          parser.on("--trained-mrl-dims LIST") { |value| options[:trained_mrl_dims] = value }
+        end.parse!(argv)
+      end
     end
 
-    def conversion_report(path, report, max_tokens)
+    def conversion_options(options)
+      options.select { |key, _| %i[max_tokens dimensions source_revision trained_mrl_dims].include?(key) }
+    end
+
+    def parse_max_tokens(value)
+      return :unlimited if value == "unlimited" || value == "0"
+
+      integer = Integer(value)
+      raise OptionParser::InvalidArgument, "--max-tokens must be positive or unlimited" unless integer.positive?
+      integer
+    end
+
+    def conversion_report(path, report)
+      max_tokens = report[:max_tokens].to_i.zero? ? "unlimited" : report[:max_tokens]
+      dim = report[:native_dim] == report[:dim] ? report[:dim].to_s : "#{report[:dim]} (from native #{report[:native_dim]})"
       [
         "wrote #{path}",
         "  vocab      #{report[:vocab_size]}",
-        "  dim        #{report[:dim]}",
+        "  dim        #{dim}",
         "  bytes      #{report[:bytes]}",
         "  sha256     #{report[:sha256]}",
         "  max_tokens #{max_tokens}",
@@ -100,20 +116,25 @@ module StaticEmbeddings
       ].join("\n")
     end
 
-    def verify(argv)
+    def verify_command(argv)
       result = StaticEmbeddings.verify(required_arg(argv, "usage: static_embeddings verify PATH"))
-      return puts("ok #{result[:expected]}") || 0 if result[:ok]
-
-      warn "CHECKSUM MISMATCH"
-      warn "  stored   #{result[:stored]}"
-      warn "  computed #{result[:expected]}"
-      1
-    end
-
-    def inspect_model(argv)
+      if result[:ok]
+        puts "ok #{result[:expected]}"
+        0
+      else
+        warn "CHECKSUM MISMATCH"
+        warn "  stored   #{result[:stored]}"
+        warn "  computed #{result[:expected]}"
+        1
+      end
+    end
+
+    def inspect_command(argv)
       model = StaticEmbeddings.load(required_arg(argv, "usage: static_embeddings inspect PATH"))
       puts JSON.pretty_generate(model_summary(model))
       0
+    ensure
+      model&.close
     end
 
     def model_summary(model)
@@ -130,18 +151,20 @@ module StaticEmbeddings
       }
     end
 
-    def tokenize(argv)
-      model, text = model_and_text(argv, "usage: static_embeddings tokenize PATH TEXT")
-      ids = model.tokenize(text)
-      puts JSON.generate("ids" => ids, "count" => ids.length, "unk" => ids.count(model.unk_id))
+    def tokenize_command(argv)
+      with_model_and_text(argv, "usage: static_embeddings tokenize PATH TEXT") do |model, text|
+        ids = model.tokenize(text)
+        puts JSON.generate("ids" => ids, "count" => ids.length, "unk" => ids.count(model.unk_id))
+      end
       0
     end
 
-    def embed(argv)
-      model, text = model_and_text(argv, "usage: static_embeddings embed PATH TEXT")
-      stats = model.embed_with_stats(text)
-      warn_high_unk(stats) if high_unk?(stats)
-      puts JSON.generate(stats_payload(model, stats))
+    def embed_command(argv)
+      with_model_and_text(argv, "usage: static_embeddings embed PATH TEXT") do |model, text|
+        stats = model.embed_with_stats(text)
+        warn_high_unk(stats) if high_unk?(stats)
+        puts JSON.generate(stats_payload(model, stats))
+      end
       0
     end
 
@@ -150,7 +173,7 @@ module StaticEmbeddings
         "token_count" => stats[:token_count],
         "unk_count" => stats[:unk_count],
         "truncated" => stats[:truncated],
-        "vector" => StaticEmbeddings.unpack(stats[:vector], model.dim).first.map { |v| v.round(6) }
+        "vector" => StaticEmbeddings.unpack(stats[:vector], model.dim).first.map { |value| value.round(6) }
       }
     end
 
@@ -163,16 +186,19 @@ module StaticEmbeddings
       warn "warning: #{(ratio * 100).round}% of tokens are [UNK] — wrong model for this language?"
     end
 
-    def model_and_text(argv, usage)
+    def with_model_and_text(argv, usage)
       path = argv.shift
       text = argv.join(" ")
-      abort usage if path.nil? || text.empty?
+      raise InvalidOptionError, usage if path.nil? || text.empty?
 
-      [StaticEmbeddings.load(path), text]
+      model = StaticEmbeddings.load(path)
+      yield model, text
+    ensure
+      model&.close
     end
 
     def required_arg(argv, usage)
-      argv.shift || abort(usage)
+      argv.shift || raise(InvalidOptionError, usage)
     end
   end
 end

data/lib/static_embeddings/codec.rb

@@ -0,0 +1,45 @@
+module StaticEmbeddings
+  module Codec
+    module_function
+
+    def unpack(blob, dim, format: :f32)
+      dim = Integer(dim)
+      raise InvalidOptionError, "dim must be positive" unless dim.positive?
+
+      values = decode(blob, normalize_format(format))
+      raise InvalidOptionError, "blob is not a multiple of dim" unless (values.length % dim).zero?
+
+      values.each_slice(dim).to_a
+    end
+
+    def pack(rows, format: :f32)
+      values = rows.first.is_a?(Array) ? rows.flatten(1) : rows
+      floats = values.map(&:to_f)
+
+      case normalize_format(format)
+      when :f32 then floats.pack("e*")
+      when :f16 then StaticEmbeddings.encode_f16(floats)
+      end
+    end
+
+    def decode(blob, format)
+      case format
+      when :f32
+        raise InvalidOptionError, "f32 blob byte size must be a multiple of 4" unless (blob.bytesize % 4).zero?
+        blob.unpack("e*")
+      when :f16
+        raise InvalidOptionError, "f16 blob byte size must be a multiple of 2" unless (blob.bytesize % 2).zero?
+        StaticEmbeddings.decode_f16(blob)
+      end
+    end
+
+    def normalize_format(format)
+      case format&.to_sym
+      when nil, :f32, :float32 then :f32
+      when :f16, :float16 then :f16
+      else
+        raise InvalidOptionError, "unsupported embedding format #{format.inspect} (expected :f32 or :f16)"
+      end
+    end
+  end
+end

data/lib/static_embeddings/conversion.rb

@@ -0,0 +1,58 @@
+require "json"
+require "static_embeddings/format/writer"
+require "static_embeddings/importers"
+require "static_embeddings/provenance"
+require "static_embeddings/unicode_tables"
+
+module StaticEmbeddings
+  module Conversion
+    module_function
+
+    def call(source_dir, output_path:, **options)
+      model = Importers.import(source_dir, **options)
+      provenance = Provenance.json(model, unicode_source: UnicodeTables.source_stamp)
+      result = Format::Writer.call(
+        path: output_path,
+        meta: format_meta(model),
+        tokens: model.tokens,
+        matrix: model.matrix,
+        norm_tables: UnicodeTables.packed,
+        provenance: provenance
+      )
+
+      result.merge(
+        vocab_size: model.tokens.length,
+        dim: model.dimensions.output,
+        native_dim: model.dimensions.native,
+        max_tokens: model.runtime.max_tokens,
+        provenance: JSON.parse(provenance)
+      )
+    end
+
+    def format_meta(model)
+      tokenizer = model.tokenizer
+      runtime = model.runtime
+      {
+        dim: model.dimensions.output,
+        normalization_type: runtime.normalization,
+        max_tokens_default: runtime.max_tokens,
+        add_special_tokens: runtime.add_special_tokens,
+        unk_policy: runtime.unk_policy,
+        empty_policy: runtime.empty_policy,
+        do_lower_case: tokenizer.fetch(:do_lower_case),
+        strip_accents: tokenizer.fetch(:strip_accents),
+        handle_chinese_chars: tokenizer.fetch(:handle_chinese_chars),
+        clean_text: tokenizer.fetch(:clean_text),
+        added_token_mask: tokenizer.fetch(:added_token_mask),
+        max_input_chars_per_word: tokenizer.fetch(:max_input_chars_per_word),
+        max_token_chars: tokenizer.fetch(:max_token_chars),
+        subword_prefix: tokenizer.fetch(:subword_prefix),
+        pad_id: tokenizer.fetch(:pad_id),
+        unk_id: tokenizer.fetch(:unk_id),
+        cls_id: tokenizer.fetch(:cls_id),
+        sep_id: tokenizer.fetch(:sep_id),
+        mask_id: tokenizer.fetch(:mask_id)
+      }
+    end
+  end
+end

data/lib/static_embeddings/errors.rb

@@ -1,6 +1,6 @@
 module StaticEmbeddings
-
   class ConversionError < Error; end
-
+  class InvalidSourceError < ConversionError; end
+  class InvalidOptionError < ArgumentError; end
   class ModelNotFound < Error; end
 end

data/lib/static_embeddings/format/constants.rb

@@ -0,0 +1,109 @@
+module StaticEmbeddings
+  module Format
+    MAGIC = "SEMBv1\0\0"
+    VERSION = 3
+    VERSION_WORDPIECE = VERSION
+    HEADER_SIZE = 320
+    ALIGNMENT = 64
+
+    TOKENIZER_BERT_WORDPIECE_V1 = 1
+    DTYPE_F32 = 1
+    POOLING_MEAN = 1
+    NORMALIZATION_NONE = 0
+    NORMALIZATION_L2 = 1
+    TRUNCATE_USABLE_IDS_BEFORE_POOLING = 2
+    UNK_INCLUDE = 0
+    UNK_DROP = 1
+    EMPTY_ZERO_VECTOR = 0
+    EMPTY_RAISE = 1
+
+    SLOT_EMPTY = 0xFFFFFFFF
+    HASH_SEED = 2_166_136_261
+    LOAD_FACTOR = 0.70
+    UINT32_MAX = 0xFFFFFFFF
+    UINT64_MAX = 0xFFFFFFFFFFFFFFFF
+
+    MAX_TOKEN_CHARS_OFFSET = 124
+    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,
+      vocab_hash: 144,
+      embeddings: 160,
+      norm_tables: 176,
+      provenance: 192,
+      root_trie: 208,
+      continuation_trie: 224
+    }.freeze
+
+    HEADER_U32 = {
+      8 => VERSION,
+      12 => HEADER_SIZE,
+      16 => 1,
+      28 => TOKENIZER_BERT_WORDPIECE_V1,
+      32 => DTYPE_F32,
+      36 => POOLING_MEAN,
+      48 => TRUNCATE_USABLE_IDS_BEFORE_POOLING,
+      108 => HASH_SEED
+    }.freeze
+
+    META_U32 = {
+      20 => :dim,
+      40 => :normalization_type,
+      44 => :max_tokens_default,
+      56 => :unk_policy,
+      60 => :empty_policy,
+      80 => :max_input_chars_per_word,
+      84 => :pad_id,
+      88 => :unk_id,
+      92 => :cls_id,
+      96 => :sep_id,
+      100 => :mask_id
+    }.freeze
+
+    META_BOOL = {
+      52 => :add_special_tokens,
+      64 => :do_lower_case,
+      68 => :strip_accents,
+      72 => :handle_chinese_chars,
+      76 => :clean_text
+    }.freeze
+
+    module_function
+
+    def binary_string(capacity = nil)
+      string = capacity ? String.new(capacity: capacity) : +""
+      string.force_encoding(Encoding::BINARY)
+    end
+
+    def verify(path)
+      require "static_embeddings/format/verifier"
+      Verifier.call(path)
+    end
+
+    def write(**kwargs)
+      require "static_embeddings/format/writer"
+      Writer.call(**kwargs)
+    end
+
+    def hash_bytes(string, seed = HASH_SEED)
+      require "static_embeddings/format/hash_table"
+      HashTable.hash_bytes(string, seed)
+    end
+
+    def next_power_of_two(value)
+      require "static_embeddings/format/hash_table"
+      HashTable.next_power_of_two(value)
+    end
+  end
+end

data/lib/static_embeddings/format/hash_table.rb

@@ -0,0 +1,69 @@
+require "static_embeddings/format/constants"
+
+module StaticEmbeddings
+  module Format
+    module HashTable
+      module_function
+
+      def build(tokens)
+        size = next_power_of_two((tokens.length / LOAD_FACTOR).ceil + 1)
+        slots = Array.new(size)
+        strings = Format.binary_string
+        max_probe = 0
+
+        tokens.each_with_index do |token, id|
+          bytes = token.b
+          offset = strings.bytesize
+          strings << bytes
+          probe = insert(slots, tokens, size, bytes, hash_bytes(bytes), offset, id)
+          max_probe = probe if probe > max_probe
+        end
+
+        [size, strings, pack(slots), max_probe]
+      end
+
+      def hash_bytes(string, seed = HASH_SEED)
+        string.each_byte.reduce(seed) { |hash, byte| ((hash ^ byte) * 16_777_619) & UINT32_MAX }
+      end
+
+      def next_power_of_two(value)
+        1 << (value - 1).bit_length
+      end
+
+      def insert(slots, tokens, size, bytes, hash, offset, id)
+        position = hash & (size - 1)
+        probe = 1
+
+        loop do
+          slot = slots[position]
+          unless slot
+            slots[position] = [hash, offset, bytes.bytesize, id]
+            return probe
+          end
+
+          if slot[0] == hash && slot[2] == bytes.bytesize && tokens.fetch(slot[3]).b == bytes
+            raise ArgumentError, "duplicate token in vocabulary: #{tokens.fetch(slot[3]).inspect}"
+          end
+
+          position = (position + 1) & (size - 1)
+          probe += 1
+        end
+      end
+
+      def pack(slots)
+        empty = [0, 0, 0, SLOT_EMPTY].pack("V4")
+        slots.each_with_object(Format.binary_string(slots.length * 16)) do |slot, packed|
+          packed << (slot ? slot.pack("V4") : empty)
+        end
+      end
+    end
+
+    def self.hash_bytes(string, seed = HASH_SEED)
+      HashTable.hash_bytes(string, seed)
+    end
+
+    def self.next_power_of_two(value)
+      HashTable.next_power_of_two(value)
+    end
+  end
+end

data/lib/static_embeddings/format/trie.rb

@@ -0,0 +1,78 @@
+require "static_embeddings/format/constants"
+
+module StaticEmbeddings
+  module Format
+    class Trie
+      Node = Struct.new(:terminal, :children, keyword_init: true)
+
+      def initialize
+        @nodes = [Node.new(terminal: SLOT_EMPTY, children: {})]
+      end
+
+      def insert(bytes, id)
+        return if bytes.empty?
+
+        node_index = 0
+        bytes.each_byte do |byte|
+          node = @nodes[node_index]
+          node_index = node.children[byte] ||= append_node
+        end
+
+        node = @nodes[node_index]
+        raise ArgumentError, "duplicate trie key for token id #{id}" unless node.terminal == SLOT_EMPTY
+
+        node.terminal = id
+      end
+
+      def pack
+        edges = []
+        nodes = @nodes.map do |node|
+          start = edges.length
+          node.children.sort.each { |byte, child| edges << [byte, child] }
+          [start, node.children.length, node.terminal, 0]
+        end
+
+        nodes.each_with_object(header(nodes.length, edges.length)) { |record, out| out << record.pack("V4") }
+             .tap { |out| edges.each { |edge| out << edge.pack("V2") } }
+      end
+
+      private
+
+      def append_node
+        @nodes << Node.new(terminal: SLOT_EMPTY, children: {})
+        @nodes.length - 1
+      end
+
+      def header(node_count, edge_count)
+        capacity = 16 + node_count * 16 + edge_count * 8
+        Format.binary_string(capacity).tap { |out| out << [node_count, edge_count, 0, 0].pack("V4") }
+      end
+    end
+
+    module WordPieceTrie
+      module_function
+
+      def build(tokens, prefix)
+        root = Trie.new
+        continuation = Trie.new
+        prefix_bytes = prefix.b
+
+        tokens.each_with_index do |token, id|
+          bytes = token.b
+          if continuation?(bytes, prefix_bytes)
+            body = bytes.byteslice(prefix_bytes.bytesize, bytes.bytesize - prefix_bytes.bytesize)
+            continuation.insert(body, id)
+          else
+            root.insert(bytes, id)
+          end
+        end
+
+        [root.pack, continuation.pack]
+      end
+
+      def continuation?(bytes, prefix)
+        !prefix.empty? && bytes.start_with?(prefix) && bytes.bytesize > prefix.bytesize
+      end
+    end
+  end
+end

data/lib/static_embeddings/format/verifier.rb

@@ -0,0 +1,41 @@
+require "digest"
+require "static_embeddings/errors"
+require "static_embeddings/format/constants"
+
+module StaticEmbeddings
+  module Format
+    module Verifier
+      CHUNK_BYTES = 1024 * 1024
+
+      module_function
+
+      def call(path)
+        raise InvalidModelError, "file too small" if File.size(path) < HEADER_SIZE
+
+        File.open(path, "rb") do |io|
+          header = io.read(HEADER_SIZE)
+          raise InvalidModelError, "bad magic" unless header.byteslice(0, 8) == MAGIC.b
+
+          stored = header.byteslice(CHECKSUM_OFFSET, CHECKSUM_SIZE)
+          actual = checksum(io, header)
+          { ok: stored == actual, expected: actual.unpack1("H*"), stored: stored.unpack1("H*") }
+        end
+      end
+
+      def checksum(io, header)
+        digest = Digest::SHA256.new
+        zeroed = header.dup
+        zeroed[CHECKSUM_OFFSET, CHECKSUM_SIZE] = "\0".b * CHECKSUM_SIZE
+        digest << zeroed
+
+        buffer = String.new(capacity: CHUNK_BYTES)
+        digest << buffer while io.read(CHUNK_BYTES, buffer)
+        digest.digest
+      end
+    end
+
+    def self.verify(path)
+      Verifier.call(path)
+    end
+  end
+end