static_embeddings 0.1.5 → 1.5.7

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

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