static_embeddings 0.1.5 → 1.5.7

data/lib/static_embeddings/bert_wordpiece.rb

@@ -0,0 +1,191 @@
+require "static_embeddings/errors"
+require "static_embeddings/format/constants"
+
+module StaticEmbeddings
+  module BertWordPiece
+    TOKENIZER_PROFILE = "BERT_WORDPIECE_V1"
+    ALLOWED_NORMALIZER_KEYS = %w[type clean_text handle_chinese_chars strip_accents lowercase].freeze
+    STANDARD_SPECIAL_TOKENS = {
+      "[PAD]" => Format::ADDED_PAD,
+      "[UNK]" => Format::ADDED_UNK,
+      "[CLS]" => Format::ADDED_CLS,
+      "[SEP]" => Format::ADDED_SEP,
+      "[MASK]" => Format::ADDED_MASK
+    }.freeze
+
+    module_function
+
+    def compile(tokenizer, tokenizer_config = {})
+      profile = audit(tokenizer, tokenizer_config)
+      tokens = vocabulary(tokenizer)
+      validate_added_token_ids(profile.fetch(:added_tokens), tokens)
+      [profile.freeze, tokens]
+    end
+
+    def runtime_meta(profile, tokens)
+      prefix = profile.fetch(:continuing_subword_prefix)
+      {
+        tokenizer_type: Format::TOKENIZER_BERT_WORDPIECE_V1,
+        tokenizer_profile: TOKENIZER_PROFILE,
+        do_lower_case: profile.fetch(:lowercase),
+        strip_accents: profile.fetch(:strip_accents),
+        handle_chinese_chars: profile.fetch(:handle_chinese_chars),
+        clean_text: profile.fetch(:clean_text),
+        added_token_mask: profile.fetch(:added_token_mask),
+        max_input_chars_per_word: profile.fetch(:max_input_chars_per_word),
+        max_token_chars: max_token_chars(tokens, prefix),
+        subword_prefix: prefix,
+        pad_id: token_id(tokens, "[PAD]", 0),
+        unk_id: token_id(tokens, profile.fetch(:unk_token)),
+        cls_id: token_id(tokens, "[CLS]", 0),
+        sep_id: token_id(tokens, "[SEP]", 0),
+        mask_id: token_id(tokens, "[MASK]", 0)
+      }.freeze
+    end
+
+    def audit(tokenizer, tokenizer_config)
+      unsupported!("tokenizer.json is not an object") unless tokenizer.is_a?(Hash)
+      unsupported!("tokenizer_config.json is not an object") unless tokenizer_config.is_a?(Hash)
+      model = tokenizer["model"] || unsupported!("tokenizer.json has no model section")
+      normalizer = tokenizer["normalizer"] || unsupported!("tokenizer has no normalizer")
+      pre_tokenizer = tokenizer["pre_tokenizer"]
+      unsupported!("tokenizer model section is not an object") unless model.is_a?(Hash)
+      unsupported!("tokenizer normalizer is not an object") unless normalizer.is_a?(Hash)
+
+      audit_model(model)
+      audit_normalizer(normalizer)
+      audit_pre_tokenizer(pre_tokenizer)
+      added_tokens = audit_added_tokens(tokenizer["added_tokens"] || [])
+      lowercase = flag(normalizer, "lowercase", tokenizer_config["do_lower_case"], true)
+      clean_text = flag(normalizer, "clean_text", nil, true)
+      max_input_chars = Integer(model.fetch("max_input_chars_per_word", 100))
+      unsupported!("max_input_chars_per_word must be positive") unless max_input_chars.positive?
+      unsupported!("clean_text=false is not supported by the runtime") unless clean_text
+
+      {
+        lowercase: lowercase,
+        strip_accents: strip_accents(normalizer, lowercase),
+        clean_text: clean_text,
+        handle_chinese_chars: flag(normalizer, "handle_chinese_chars",
+                                   tokenizer_config["tokenize_chinese_chars"], true),
+        continuing_subword_prefix: model.fetch("continuing_subword_prefix", "##"),
+        unk_token: model.fetch("unk_token", "[UNK]"),
+        max_input_chars_per_word: max_input_chars,
+        tokenizer_class: tokenizer_config["tokenizer_class"],
+        added_tokens: added_tokens,
+        added_token_mask: added_tokens.reduce(0) do |mask, token|
+          mask | STANDARD_SPECIAL_TOKENS.fetch(token.fetch("content"))
+        end
+      }
+    end
+
+    def audit_model(model)
+      unsupported!("tokenizer model.type is #{model['type'].inspect}, expected WordPiece") unless model["type"] == "WordPiece"
+
+      prefix = model.fetch("continuing_subword_prefix", "##")
+      unsupported!("continuing_subword_prefix #{prefix.inspect} is not supported") unless prefix == "##"
+
+      unk_token = model.fetch("unk_token", "[UNK]")
+      unsupported!("unk_token #{unk_token.inspect} is not supported") unless unk_token == "[UNK]"
+    end
+
+    def audit_normalizer(normalizer)
+      unsupported!("normalizer type #{normalizer['type'].inspect} is not BertNormalizer") unless normalizer["type"] == "BertNormalizer"
+
+      unknown = normalizer.keys - ALLOWED_NORMALIZER_KEYS
+      unsupported!("normalizer has unsupported keys #{unknown.inspect}") unless unknown.empty?
+    end
+
+    def audit_pre_tokenizer(pre_tokenizer)
+      return if pre_tokenizer.is_a?(Hash) && pre_tokenizer["type"] == "BertPreTokenizer"
+
+      unsupported!("pre_tokenizer type #{pre_tokenizer && pre_tokenizer['type'].inspect} is not BertPreTokenizer")
+    end
+
+    def audit_added_tokens(tokens)
+      unsupported!("added_tokens is not an array") unless tokens.is_a?(Array)
+      unsupported!("added_tokens contains a non-object entry") unless tokens.all? { |token| token.is_a?(Hash) }
+      bad = tokens.reject { |token| standard_special?(token) }
+      unsupported!("tokenizer declares non-standard added_tokens #{bad.map { |token| token['content'] }.inspect}") unless bad.empty?
+
+      tokens.each do |token|
+        content = token["content"].to_s
+        unsupported!("added token #{content.inspect} contains whitespace") if content.match?(/\s/)
+        unsupported!("added token #{content.inspect} uses lstrip/rstrip/single_word") if token["lstrip"] || token["rstrip"] || token["single_word"]
+        unsupported!("added token #{content.inspect} must use normalized=false") unless token["normalized"] == false
+      end
+
+      duplicate = tokens.group_by { |token| token["content"] }.find { |_, rows| rows.length > 1 }
+      unsupported!("duplicate added token #{duplicate.first.inspect}") if duplicate
+      tokens.freeze
+    end
+
+    def vocabulary(tokenizer)
+      vocab = tokenizer.dig("model", "vocab") || unsupported!("tokenizer.json has no model.vocab")
+      unsupported!("tokenizer model.vocab is not an object") unless vocab.is_a?(Hash)
+      tokens = Array.new(vocab.length)
+      vocab.each { |token, id| assign_vocab(tokens, token, id) }
+      missing = tokens.index(nil)
+      invalid!("vocab has a hole at id #{missing}") if missing
+      tokens
+    end
+
+    def assign_vocab(tokens, token, id)
+      unsupported!("vocab id #{id.inspect} is not an integer") unless id.is_a?(Integer)
+      invalid!("vocab id #{id} for #{token.inspect} is outside 0...#{tokens.length}") unless id.between?(0, tokens.length - 1)
+      invalid!("duplicate vocab id #{id}") unless tokens[id].nil?
+      tokens[id] = token
+    end
+
+    def validate_added_token_ids(added_tokens, tokens)
+      added_tokens.each do |token|
+        content = token.fetch("content")
+        id = token["id"]
+        unsupported!("added token #{content.inspect} has non-integer id #{id.inspect}") unless id.is_a?(Integer)
+        valid = id.between?(0, tokens.length - 1) && tokens[id] == content
+        unsupported!("added token #{content.inspect} id #{id} does not match model.vocab") unless valid
+      end
+    end
+
+    def flag(hash, key, fallback, default)
+      return !!hash[key] if hash.key?(key) && !hash[key].nil?
+      return !!fallback unless fallback.nil?
+
+      default
+    end
+
+    def strip_accents(normalizer, lowercase)
+      value = normalizer["strip_accents"]
+      normalizer.key?("strip_accents") && !value.nil? ? !!value : lowercase
+    end
+
+    def standard_special?(token)
+      token["special"] && STANDARD_SPECIAL_TOKENS.key?(token["content"])
+    end
+
+    def max_token_chars(tokens, prefix)
+      tokens.reduce(1) do |maximum, token|
+        body = token.start_with?(prefix) ? token[prefix.length..] : token
+        [maximum, body.each_char.count].max
+      end
+    end
+
+    def token_id(tokens, token, fallback = nil)
+      id = tokens.index(token)
+      return id unless id.nil?
+      return fallback unless fallback.nil?
+
+      invalid!("vocabulary has no #{token.inspect}")
+    end
+
+    def unsupported!(message)
+      raise UnsupportedModelError,
+            "#{message}. Supported tokenizer profile: #{TOKENIZER_PROFILE}; " \
+            "other tokenizer behaviour requires a separate runtime capability."
+    end
+
+    def invalid!(message)
+      raise ConversionError, message
+    end
+  end
+end

data/lib/static_embeddings/canonical.rb

@@ -0,0 +1,50 @@
+module StaticEmbeddings
+  module Canonical
+    Dimensions = Struct.new(:native, :output, :trained, keyword_init: true)
+    Runtime = Struct.new(:normalization, :unk_policy, :empty_policy, :max_tokens, :add_special_tokens,
+                         keyword_init: true)
+    Source = Struct.new(:family, :model, :revision, :oracle, :files_sha256, :tokenizer_class,
+                        :config_seq_length, keyword_init: true)
+    Model = Struct.new(:tokens, :matrix, :dimensions, :runtime, :tokenizer, :source, keyword_init: true)
+
+    module_function
+
+    def dimensions(native:, output:, trained: nil)
+      Dimensions.new(native: native, output: output, trained: trained&.freeze).freeze
+    end
+
+    def runtime(normalization:, unk_policy:, empty_policy:, max_tokens:, add_special_tokens: false)
+      Runtime.new(
+        normalization: normalization,
+        unk_policy: unk_policy,
+        empty_policy: empty_policy,
+        max_tokens: max_tokens,
+        add_special_tokens: add_special_tokens
+      ).freeze
+    end
+
+    def source(family:, model:, oracle:, files_sha256:, revision: nil, tokenizer_class: nil,
+               config_seq_length: nil)
+      Source.new(
+        family: family,
+        model: model,
+        revision: revision,
+        oracle: oracle,
+        files_sha256: files_sha256.freeze,
+        tokenizer_class: tokenizer_class,
+        config_seq_length: config_seq_length
+      ).freeze
+    end
+
+    def model(tokens:, matrix:, dimensions:, runtime:, tokenizer:, source:)
+      Model.new(
+        tokens: tokens.freeze,
+        matrix: matrix,
+        dimensions: dimensions,
+        runtime: runtime,
+        tokenizer: tokenizer.freeze,
+        source: source
+      ).freeze
+    end
+  end
+end

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,67 +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)
-      require "static_embeddings/converter"
-      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}",
@@ -101,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)
@@ -131,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
 
@@ -151,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
 
@@ -164,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