static_embeddings 0.1.5 → 1.5.6

data/lib/static_embeddings/format/writer.rb

@@ -0,0 +1,131 @@
+require "digest"
+require "static_embeddings/format/constants"
+require "static_embeddings/format/hash_table"
+require "static_embeddings/format/trie"
+
+module StaticEmbeddings
+  module Format
+    module Writer
+      module_function
+
+      def call(path:, meta:, tokens:, matrix:, norm_tables:, provenance:)
+        hash_size, strings, hash_blob, max_probe = HashTable.build(tokens)
+        root_trie, continuation_trie = WordPieceTrie.build(tokens, meta.fetch(:subword_prefix))
+        payloads = payloads(strings, hash_blob, matrix, norm_tables, provenance, root_trie, continuation_trie)
+        sections, file_size = layout(payloads)
+        header = build_header(meta, tokens.length, hash_size, max_probe, sections)
+        checksum = write_file(path, header, payloads, sections)
+
+        { bytes: file_size, sha256: checksum.unpack1("H*"), hash_table_size: hash_size }
+      end
+
+      def payloads(strings, hash_blob, matrix, norm_tables, provenance, root_trie, continuation_trie)
+        {
+          vocab_strings: strings,
+          vocab_hash: hash_blob,
+          embeddings: matrix,
+          norm_tables: norm_tables,
+          provenance: provenance,
+          root_trie: root_trie,
+          continuation_trie: continuation_trie
+        }
+      end
+
+      def layout(payloads)
+        offset = HEADER_SIZE
+        sections = payloads.each_with_object({}) do |(name, payload), out|
+          offset += (ALIGNMENT - (offset % ALIGNMENT)) % ALIGNMENT
+          out[name] = [offset, payload.bytesize]
+          offset += payload.bytesize
+        end
+        [sections, offset]
+      end
+
+      def write_file(path, header, payloads, sections)
+        digest = Digest::SHA256.new
+        File.open(path, "wb") do |io|
+          write_chunk(io, digest, header)
+          offset = HEADER_SIZE
+
+          payloads.each do |name, payload|
+            target = sections.fetch(name).first
+            padding = target - offset
+            write_chunk(io, digest, "\0".b * padding) if padding.positive?
+            write_payload(io, digest, payload)
+            offset = target + payload.bytesize
+          end
+        end
+
+        digest.digest.tap do |checksum|
+          File.open(path, "r+b") do |io|
+            io.seek(CHECKSUM_OFFSET, IO::SEEK_SET)
+            io.write(checksum)
+          end
+        end
+      end
+
+      def write_payload(io, digest, payload)
+        if payload.respond_to?(:each_chunk)
+          payload.each_chunk { |chunk| write_chunk(io, digest, chunk) }
+        else
+          write_chunk(io, digest, payload)
+        end
+      end
+
+      def write_chunk(io, digest, chunk)
+        io.write(chunk)
+        digest << chunk
+      end
+
+      def build_header(meta, vocab_size, hash_size, max_probe, sections)
+        header = "\0".b * HEADER_SIZE
+        header[0, 8] = MAGIC.b
+
+        HEADER_U32.each { |offset, value| put_u32(header, offset, value) }
+        META_U32.each { |offset, key| put_u32(header, offset, meta.fetch(key)) }
+        META_BOOL.each { |offset, key| put_u32(header, offset, meta.fetch(key) ? 1 : 0) }
+        put_u32(header, 24, vocab_size)
+        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)) }
+        header
+      end
+
+      def put_prefix(header, prefix)
+        bytes = prefix.b
+        raise ArgumentError, "subword prefix too long" if bytes.bytesize > 8
+
+        put_u32(header, 112, bytes.bytesize)
+        header[116, 8] = bytes.ljust(8, "\0")
+      end
+
+      def put_section(header, offset, section)
+        put_u64(header, offset, section.fetch(0))
+        put_u64(header, offset + 8, section.fetch(1))
+      end
+
+      def put_u32(buffer, offset, value)
+        integer = Integer(value)
+        unless integer.between?(0, UINT32_MAX)
+          raise ArgumentError, "u32 value out of range: #{integer}"
+        end
+        buffer[offset, 4] = [integer].pack("V")
+      end
+
+      def put_u64(buffer, offset, value)
+        integer = Integer(value)
+        unless integer.between?(0, UINT64_MAX)
+          raise ArgumentError, "u64 value out of range: #{integer}"
+        end
+        buffer[offset, 8] = [integer & UINT32_MAX, integer >> 32].pack("V2")
+      end
+    end
+
+    def self.write(**kwargs)
+      Writer.call(**kwargs)
+    end
+  end
+end

data/lib/static_embeddings/format.rb

@@ -1,338 +1,3 @@
-require "digest"
-
-module StaticEmbeddings
-  module Format
-    MAGIC = "SEMBv1\0\0"
-    VERSION = 3
-    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
-
-    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
-
-    VERIFY_CHUNK_BYTES = 1024 * 1024
-
-    module_function
-
-    def hash_bytes(str, seed = HASH_SEED)
-      str.each_byte.reduce(seed) { |h, b| ((h ^ b) * 16_777_619) & 0xFFFFFFFF }
-    end
-
-    def next_power_of_two(n)
-      1 << (n - 1).bit_length
-    end
-
-    def build_hash_table(tokens)
-      size = next_power_of_two((tokens.length / LOAD_FACTOR).ceil + 1)
-      slots = Array.new(size)
-      strings = binary_string
-      max_probe = 0
-
-      tokens.each_with_index do |token, id|
-        bytes = token.b
-        offset = strings.bytesize
-        strings << bytes
-        probe = insert_slot!(slots, tokens, size, bytes, hash_bytes(bytes), offset, id)
-        max_probe = probe if probe > max_probe
-      end
-
-      [size, strings, pack_slots(slots), max_probe]
-    end
-
-    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))
-      payloads = {
-        vocab_strings: strings,
-        vocab_hash: hash_blob,
-        embeddings: matrix,
-        norm_tables: norm_tables,
-        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)
-
-      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)
-      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 = streaming_checksum(io, header)
-
-        { ok: stored == actual, expected: actual.unpack1("H*"), stored: stored.unpack1("H*") }
-      end
-    end
-
-    def streaming_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: VERIFY_CHUNK_BYTES)
-      digest << buffer while io.read(VERIFY_CHUNK_BYTES, buffer)
-      digest.digest
-    end
-
-    def binary_string(capacity = nil)
-      str = capacity ? String.new(capacity: capacity) : +""
-      str.force_encoding(Encoding::BINARY)
-    end
-
-    def insert_slot!(slots, tokens, size, bytes, hash, offset, id)
-      pos = hash & (size - 1)
-      probe = 1
-      loop do
-        slot = slots[pos]
-        unless slot
-          slots[pos] = [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
-
-        pos = (pos + 1) & (size - 1)
-        probe += 1
-      end
-    end
-
-    def pack_slots(slots)
-      empty = [0, 0, 0, SLOT_EMPTY].pack("V4")
-      packed = binary_string(slots.length * 16)
-      slots.each { |slot| packed << (slot ? slot.pack("V4") : empty) }
-      packed
-    end
-
-    def build_wordpiece_tries(tokens, prefix)
-      root = TrieBuilder.new
-      continuation = TrieBuilder.new
-      prefix_bytes = prefix.b
-
-      tokens.each_with_index do |token, id|
-        bytes = token.b
-        if !prefix_bytes.empty? && bytes.start_with?(prefix_bytes) && bytes.bytesize > prefix_bytes.bytesize
-          continuation.insert(bytes.byteslice(prefix_bytes.bytesize, bytes.bytesize - prefix_bytes.bytesize), id)
-        else
-          root.insert(bytes, id)
-        end
-      end
-
-      [root.pack, continuation.pack]
-    end
-
-    class TrieBuilder
-      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]
-          child = node.children[byte]
-          unless child
-            child = @nodes.length
-            node.children[byte] = child
-            @nodes << Node.new(terminal: SLOT_EMPTY, children: {})
-          end
-          node_index = child
-        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 = []
-        node_records = @nodes.map do |node|
-          start = edges.length
-          node.children.sort_by { |byte, _| byte }.each do |byte, child|
-            edges << [byte, child]
-          end
-          [start, node.children.length, node.terminal, 0]
-        end
-
-        packed = Format.binary_string(16 + node_records.length * 16 + edges.length * 8)
-        packed << [node_records.length, edges.length, 0, 0].pack("V4")
-        node_records.each { |record| packed << record.pack("V4") }
-        edges.each { |edge| packed << edge.pack("V2") }
-        packed
-      end
-    end
-
-    def layout_sections(payloads)
-      offset = HEADER_SIZE
-      sections = {}
-      payloads.each do |name, payload|
-        padding = (ALIGNMENT - (offset % ALIGNMENT)) % ALIGNMENT
-        offset += padding
-        sections[name] = [offset, payload.bytesize]
-        offset += payload.bytesize
-      end
-      [sections, offset]
-    end
-
-    def build_header(meta, vocab_size, hash_size, max_probe, sections)
-      header = "\0".b * HEADER_SIZE
-      header[0, 8] = MAGIC.b
-
-      HEADER_U32.each { |offset, value| put_u32(header, offset, value) }
-      META_U32.each { |offset, key| put_u32(header, offset, meta.fetch(key)) }
-      META_BOOL.each { |offset, key| put_u32(header, offset, meta.fetch(key) ? 1 : 0) }
-
-      put_u32(header, 24, vocab_size)
-      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)) }
-
-      header
-    end
-
-    def put_prefix(header, prefix)
-      bytes = prefix.b
-      raise ArgumentError, "subword prefix too long" if bytes.bytesize > 8
-
-      put_u32(header, 112, bytes.bytesize)
-      header[116, 8] = bytes.ljust(8, "\0")
-    end
-
-    def put_section(header, offset, section)
-      off, size = section
-      put_u64(header, offset, off)
-      put_u64(header, offset + 8, size)
-    end
-
-    def put_u32(buffer, offset, value)
-      buffer[offset, 4] = [value].pack("V")
-    end
-
-    def put_u64(buffer, offset, value)
-      buffer[offset, 8] = [value & 0xFFFFFFFF, value >> 32].pack("V2")
-    end
-
-  end
-end
+require "static_embeddings/format/constants"
+require "static_embeddings/format/verifier"
+require "static_embeddings/format/writer"

data/lib/static_embeddings/importers/model2vec.rb

@@ -0,0 +1,52 @@
+module StaticEmbeddings
+  module Importers
+    module Model2Vec
+      SOURCE_FILES = %w[tokenizer.json config.json tokenizer_config.json model.safetensors].freeze
+      DEFAULT_MAX_TOKENS = 512
+      FAMILY = "model2vec"
+      ORACLE = "model2vec.StaticModel"
+
+      module_function
+
+      def call(root, model_id: nil, max_tokens: nil, dimensions: nil,
+               source_revision: nil, trained_mrl_dims: nil)
+        tokenizer = Support.load_object(File.join(root, "tokenizer.json"))
+        config = Support.load_object(File.join(root, "config.json"), optional: true) || {}
+        tokenizer_config = Support.load_object(File.join(root, "tokenizer_config.json"), optional: true) || {}
+        profile, tokens = BertWordPiece.compile(tokenizer, tokenizer_config)
+        payload, native_dim = Support.extract_matrix(File.join(root, "model.safetensors"), tokens.length)
+        output_dim = Support.resolve_dimensions(dimensions, native_dim)
+
+        Canonical.model(
+          tokens: tokens,
+          matrix: Support.slice_matrix(payload, native_dim, output_dim),
+          dimensions: Canonical.dimensions(
+            native: native_dim,
+            output: output_dim,
+            trained: Support.normalize_mrl_dims(trained_mrl_dims, native_dim: native_dim)
+          ),
+          runtime: Canonical.runtime(
+            normalization: normalization(config),
+            unk_policy: Format::UNK_DROP,
+            empty_policy: Format::EMPTY_ZERO_VECTOR,
+            max_tokens: Support.resolve_max_tokens(max_tokens, DEFAULT_MAX_TOKENS)
+          ),
+          tokenizer: BertWordPiece.runtime_meta(profile, tokens),
+          source: Canonical.source(
+            family: FAMILY,
+            model: model_id || File.basename(root),
+            revision: source_revision,
+            oracle: ORACLE,
+            files_sha256: Support.digest_files(SOURCE_FILES.map { |name| File.join(root, name) }),
+            tokenizer_class: profile[:tokenizer_class],
+            config_seq_length: config["seq_length"]
+          )
+        )
+      end
+
+      def normalization(config)
+        config.fetch("normalize", false) ? Format::NORMALIZATION_L2 : Format::NORMALIZATION_NONE
+      end
+    end
+  end
+end

data/lib/static_embeddings/importers/sentence_transformers_static.rb

@@ -0,0 +1,103 @@
+require "pathname"
+
+module StaticEmbeddings
+  module Importers
+    module SentenceTransformersStatic
+      TYPE = "sentence_transformers.models.StaticEmbedding"
+      FAMILY = "sentence_transformers_static"
+      ORACLE = "sentence_transformers.SentenceTransformer.encode"
+      DEFAULT_MAX_TOKENS = 0
+
+      module_function
+
+      def call(root, model_id: nil, max_tokens: nil, dimensions: nil,
+               source_revision: nil, trained_mrl_dims: nil)
+        module_dir = module_directory(root)
+        tokenizer = Support.load_object(File.join(module_dir, "tokenizer.json"))
+        tokenizer_config = Support.load_object(File.join(module_dir, "tokenizer_config.json"), optional: true) || {}
+        config = Support.load_object(File.join(module_dir, "config.json"), optional: true) || {}
+        profile, tokens = BertWordPiece.compile(tokenizer, tokenizer_config)
+        payload, native_dim = Support.extract_matrix(File.join(module_dir, "model.safetensors"), tokens.length)
+        output_dim = Support.resolve_dimensions(dimensions, native_dim)
+
+        Canonical.model(
+          tokens: tokens,
+          matrix: Support.slice_matrix(payload, native_dim, output_dim),
+          dimensions: Canonical.dimensions(
+            native: native_dim,
+            output: output_dim,
+            trained: Support.normalize_mrl_dims(trained_mrl_dims, native_dim: native_dim)
+          ),
+          runtime: Canonical.runtime(
+            normalization: Format::NORMALIZATION_NONE,
+            unk_policy: Format::UNK_INCLUDE,
+            empty_policy: Format::EMPTY_ZERO_VECTOR,
+            max_tokens: Support.resolve_max_tokens(max_tokens, DEFAULT_MAX_TOKENS)
+          ),
+          tokenizer: BertWordPiece.runtime_meta(profile, tokens),
+          source: Canonical.source(
+            family: FAMILY,
+            model: model_id || File.basename(root),
+            revision: source_revision,
+            oracle: ORACLE,
+            files_sha256: source_digests(root, module_dir),
+            tokenizer_class: profile[:tokenizer_class],
+            config_seq_length: config["seq_length"]
+          )
+        )
+      end
+
+      def module_directory(root)
+        spec = module_spec(root)
+        contained_directory(root, spec.fetch("path"))
+      end
+
+      def module_spec(root)
+        modules = Support.load_json(File.join(root, "modules.json"))
+        reject_source("modules.json is not an array") unless modules.is_a?(Array)
+        reject_source("modules.json declares #{modules.length} modules, expected exactly one StaticEmbedding") unless modules.length == 1
+
+        spec = modules.first
+        reject_source("modules.json[0] is not an object") unless spec.is_a?(Hash)
+        reject_source("modules.json[0].type is #{spec['type'].inspect}, expected #{TYPE}") unless spec["type"] == TYPE
+        reject_source("modules.json[0] has no path") unless spec.key?("path")
+        reject_source("modules.json[0].path is not a string") unless spec["path"].is_a?(String)
+        spec
+      end
+
+      def contained_directory(root, relative)
+        relative = relative.to_s
+        reject_source("module path contains a NUL byte") if relative.include?("\0")
+        reject_source("module path #{relative.inspect} is absolute") if Pathname.new(relative).absolute?
+
+        root_real = File.realpath(root)
+        candidate = relative.empty? ? root : File.expand_path(relative, root)
+        reject_source("module path #{relative.inspect} is not a directory") unless File.directory?(candidate)
+
+        real = File.realpath(candidate)
+        prefix = root_real.end_with?(File::SEPARATOR) ? root_real : "#{root_real}#{File::SEPARATOR}"
+        reject_source("module path #{relative.inspect} escapes the source directory") unless real == root_real || real.start_with?(prefix)
+        real
+      rescue Errno::ENOENT
+        reject_source("module path #{relative.inspect} does not exist")
+      end
+
+      def source_digests(root, module_dir)
+        paths = [
+          File.join(root, "modules.json"),
+          File.join(module_dir, "tokenizer.json"),
+          File.join(module_dir, "tokenizer_config.json"),
+          File.join(module_dir, "config.json"),
+          File.join(module_dir, "model.safetensors")
+        ]
+        Support.digest_files(paths, root: root)
+      end
+
+      def reject_source(message)
+        raise UnsupportedModelError,
+              "#{message}. Sentence Transformers conversion accepts exactly one #{TYPE} module " \
+              "using #{BertWordPiece::TOKENIZER_PROFILE}."
+      end
+    end
+  end
+end