static_embeddings 0.1.1

data/lib/static_embeddings/format.rb

@@ -0,0 +1,289 @@
+require "digest"
+
+module StaticEmbeddings
+  module Format
+    MAGIC = "SEMBv1\0\0"
+    VERSION = 2
+    HEADER_SIZE = 320
+    ALIGNMENT = 64
+
+    TOKENIZER_BERT_WORDPIECE_V1 = 1
+    DTYPE_F32 = 1
+    POOLING_MEAN = 1
+    NORMALIZATION_NONE = 0
+    NORMALIZATION_L2 = 1
+    TRUNCATE_IDS_BEFORE_POOLING = 1
+    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
+
+    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_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 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))
+      sections, body = build_body(
+        vocab_strings: strings,
+        vocab_hash: hash_blob,
+        embeddings: matrix,
+        norm_tables: norm_tables,
+        provenance: provenance,
+        root_trie: root_trie,
+        continuation_trie: continuation_trie
+      )
+
+      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 }
+    end
+
+    def verify(path)
+      data = File.binread(path)
+      raise InvalidModelError, "file too small" if data.bytesize < HEADER_SIZE
+      raise InvalidModelError, "bad magic" unless data.byteslice(0, 8) == MAGIC.b
+
+      stored = data.byteslice(CHECKSUM_OFFSET, CHECKSUM_SIZE)
+      actual = checksum_for_verify(data)
+
+      { ok: stored == actual, expected: actual.unpack1("H*"), stored: stored.unpack1("H*") }
+    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 build_body(payloads)
+      body = binary_string
+      sections = {}
+      payloads.each do |name, payload|
+        align_body!(body)
+        sections[name] = [HEADER_SIZE + body.bytesize, payload.bytesize]
+        body << payload
+      end
+      [sections, body]
+    end
+
+    def align_body!(body)
+      padding = (ALIGNMENT - ((HEADER_SIZE + body.bytesize) % ALIGNMENT)) % ALIGNMENT
+      body << "\0".b * padding if padding.positive?
+    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_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
+
+    def checksum_for_verify(data)
+      zeroed = data.dup
+      zeroed[CHECKSUM_OFFSET, CHECKSUM_SIZE] = "\0".b * CHECKSUM_SIZE
+      Digest::SHA256.digest(zeroed)
+    end
+  end
+end

data/lib/static_embeddings/model.rb

@@ -0,0 +1,48 @@
+require "json"
+
+module StaticEmbeddings
+  class Model
+    attr_reader :path
+
+    def provenance
+      raw = provenance_json
+      return {} if raw.nil?
+
+      JSON.parse(raw)
+    rescue JSON::ParserError, EncodingError => e
+      raise InvalidModelError, "invalid provenance JSON: #{e.message}"
+    end
+
+    def model_id
+      provenance["source_model_id"]
+    end
+
+    def embed_array(text, **opts)
+      format = opts.key?(:format) ? opts[:format] : :f32
+      StaticEmbeddings.unpack(embed(text, **opts), dim, format: format).first
+    end
+
+    def embed_batch_arrays(texts, **opts)
+      format = opts.key?(:format) ? opts[:format] : :f32
+      StaticEmbeddings.unpack(embed_batch(texts, **opts), dim, format: format)
+    end
+
+    def cosine_top_k(query_blob, matrix_blob, k, **opts)
+      raise ArgumentError, "dim: is set by the model" if opts.key?(:dim)
+
+      StaticEmbeddings.cosine_top_k(query_blob, matrix_blob, k, **opts.merge(dim: dim))
+    end
+
+    def dot_top_k(query_blob, matrix_blob, k, **opts)
+      raise ArgumentError, "dim: is set by the model" if opts.key?(:dim)
+
+      StaticEmbeddings.dot_top_k(query_blob, matrix_blob, k, **opts.merge(dim: dim))
+    end
+
+    def to_s
+      "#<StaticEmbeddings::Model #{model_id || path} dim=#{dim} vocab=#{vocab_size}>"
+    end
+
+    alias inspect to_s
+  end
+end

data/lib/static_embeddings/paths.rb

@@ -0,0 +1,29 @@
+module StaticEmbeddings
+  module Paths
+    BUILTIN_DIR = File.expand_path("../models", __dir__)
+    BUILTIN_MODELS = { demo: "demo.semb" }.freeze
+
+    module_function
+
+    def cache_dir(env = ENV)
+      env["STATIC_EMBEDDINGS_CACHE"] ||
+        File.join(env["XDG_CACHE_HOME"] || File.join(Dir.home, ".cache"), "static_embeddings")
+    end
+
+    def model_path(model_id, env = ENV)
+      File.join(cache_dir(env), "models", "#{model_id}.semb")
+    end
+
+    def builtin_path(name)
+      file = BUILTIN_MODELS.fetch(name) do
+        raise ModelNotFound, "unknown builtin model #{name.inspect}"
+      end
+      File.join(BUILTIN_DIR, file)
+    end
+
+    def builtin_available?(name = :demo)
+      file = BUILTIN_MODELS[name]
+      !file.nil? && File.file?(File.join(BUILTIN_DIR, file))
+    end
+  end
+end

data/lib/static_embeddings/reference.rb

@@ -0,0 +1,191 @@
+module StaticEmbeddings
+  class Reference
+    CJK_RANGES = [
+      0x4E00..0x9FFF, 0x3400..0x4DBF, 0x20000..0x2A6DF, 0x2A700..0x2B73F,
+      0x2B740..0x2B81F, 0x2B820..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_MN = /\A\p{Mn}\z/
+    RE_WHITESPACE = /\A(?:\p{Zs}|[\u0085\u2028\u2029])\z/
+
+    attr_reader :meta
+
+    def initialize(tokens:, matrix:, meta:)
+      @tokens = tokens
+      @vocab = tokens.each_with_index.to_h
+      @matrix = matrix
+      @meta = meta
+      @dim = meta.fetch(:dim)
+    end
+
+    def self.from_source_dir(dir, max_tokens: Converter::REFERENCE_MAX_TOKENS)
+      tokenizer = JSON.parse(File.binread(File.join(dir, "tokenizer.json")))
+      config_path = File.join(dir, "config.json")
+      config = File.file?(config_path) ? JSON.parse(File.binread(config_path)) : {}
+      tokens = tokens_from(tokenizer)
+      tensor = Safetensors.read(File.join(dir, "model.safetensors"))[:tensors].values.first
+      rows, dim = tensor[:shape]
+      floats = tensor[:bytes].unpack("e*")
+      normalizer = tokenizer["normalizer"] || {}
+
+      new(tokens: tokens,
+          matrix: Array.new(rows) { |i| floats[i * dim, dim] },
+          meta: meta_from(tokenizer, config, normalizer, dim, max_tokens))
+    end
+
+    def self.tokens_from(tokenizer)
+      vocab = tokenizer.dig("model", "vocab")
+      vocab.each_with_object(Array.new(vocab.length)) { |(token, id), tokens| tokens[id] = token }
+    end
+
+    def self.meta_from(tokenizer, config, normalizer, dim, max_tokens)
+      lowercase = normalizer.fetch("lowercase", true)
+      {
+        dim: dim,
+        lowercase: lowercase,
+        strip_accents: normalizer["strip_accents"].nil? ? lowercase : normalizer["strip_accents"],
+        clean_text: normalizer.fetch("clean_text", true),
+        handle_chinese_chars: normalizer.fetch("handle_chinese_chars", true),
+        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
+      }
+    end
+
+    def normalize_text(text)
+      chars = text.chars
+      chars = clean_chars(chars) if @meta[:clean_text]
+      chars = split_chinese(chars) if @meta[:handle_chinese_chars]
+      chars = strip_accents(chars) if @meta[:strip_accents]
+      chars = lowercase(chars) if @meta[:lowercase]
+      chars.join
+    end
+
+    def pre_tokenize(normalized)
+      normalized.split(" ").flat_map { |word| split_punctuation(word) }
+    end
+
+    def tokenize(text, max_tokens: @meta[:max_tokens])
+      ids = []
+      pre_tokenize(normalize_text(text)).each { |word| ids.concat(wordpiece(word)) }
+      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 }
+      return Array.new(@dim, 0.0) if used.empty?
+
+      vector = pooled(used)
+      @meta[:normalize] ? l2_normalize(vector) : vector
+    end
+
+    private
+
+    def clean_chars(chars)
+      chars.filter_map do |ch|
+        cp = ch.ord
+        next nil if cp.zero? || cp == 0xFFFD || control?(ch)
+
+        whitespace?(ch) ? " " : ch
+      end
+    end
+
+    def split_chinese(chars)
+      chars.flat_map { |ch| cjk?(ch.ord) ? [" ", ch, " "] : [ch] }
+    end
+
+    def strip_accents(chars)
+      chars.flat_map { |ch| ch.unicode_normalize(:nfd).chars }.reject { |ch| RE_MN.match?(ch) }
+    end
+
+    def lowercase(chars)
+      chars.flat_map { |ch| ch.downcase.chars }
+    end
+
+    def split_punctuation(word)
+      out = []
+      current = +""
+      word.each_char do |ch|
+        if punctuation?(ch)
+          out << current unless current.empty?
+          out << ch
+          current = +""
+        else
+          current << ch
+        end
+      end
+      out << current unless current.empty?
+      out
+    end
+
+    def wordpiece(word)
+      chars = word.chars
+      return [unk_id] if chars.length > @meta[:max_input_chars_per_word]
+
+      ids = []
+      start = 0
+      while start < chars.length
+        found, finish = longest_piece(chars, start)
+        return [unk_id] if found.nil?
+
+        ids << found
+        start = finish
+      end
+      ids
+    end
+
+    def longest_piece(chars, start)
+      finish = chars.length
+      while start < finish
+        piece = chars[start...finish].join
+        piece = "###{piece}" if start.positive?
+        id = @vocab[piece]
+        return [id, finish] unless id.nil?
+
+        finish -= 1
+      end
+      [nil, nil]
+    end
+
+    def pooled(ids)
+      acc = Array.new(@dim, 0.0)
+      ids.each do |id|
+        row = @matrix[id]
+        @dim.times { |i| acc[i] += row[i] }
+      end
+      scale = 1.0 / ids.length
+      acc.map { |v| v * scale }
+    end
+
+    def l2_normalize(vec)
+      norm = Math.sqrt(vec.sum { |v| v * v })
+      norm.positive? ? vec.map { |v| v / norm } : vec
+    end
+
+    def unk_id
+      @vocab.fetch(@meta[:unk_token])
+    end
+
+    def cjk?(codepoint)
+      CJK_RANGES.any? { |range| range.cover?(codepoint) }
+    end
+
+    def punctuation?(char)
+      cp = char.ord
+      ASCII_PUNCT.any? { |range| range.cover?(cp) } || RE_PUNCT.match?(char)
+    end
+
+    def control?(char)
+      !ASCII_SPACES.include?(char) && RE_CONTROL.match?(char)
+    end
+
+    def whitespace?(char)
+      ASCII_SPACES.include?(char) || RE_WHITESPACE.match?(char)
+    end
+  end
+end

data/lib/static_embeddings/safetensors.rb

@@ -0,0 +1,87 @@
+require "json"
+
+module StaticEmbeddings
+  module Safetensors
+    MAX_HEADER_BYTES = 100 * 1024 * 1024
+    SUPPORTED_DTYPES = { "F32" => 4 }.freeze
+
+    module_function
+
+    def read(path)
+      data = File.binread(path)
+      header, body_offset = parse_header(data)
+      body_size = data.bytesize - body_offset
+
+      tensors = header.each_with_object({}) do |(name, spec), acc|
+        next if name == "__metadata__"
+
+        acc[name] = tensor_from(data, body_offset, body_size, name, spec)
+      end
+
+      { metadata: header["__metadata__"] || {}, tensors: tensors }
+    end
+
+    def write(path, name, shape, floats)
+      body = floats.pack("e*")
+      header = JSON.generate(name => { "dtype" => "F32", "shape" => shape, "data_offsets" => [0, body.bytesize] })
+      padded = header << (" " * ((8 - (header.bytesize % 8)) % 8))
+      File.binwrite(path, [padded.bytesize].pack("Q<") << padded << body)
+    end
+
+    def parse_header(data)
+      raise ConversionError, "safetensors file is shorter than its length prefix" if data.bytesize < 8
+
+      header_len = data.byteslice(0, 8).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
+
+      raw_header = data.byteslice(8, header_len)
+      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) }
+    end
+
+    def checked_offsets(name, offsets, body_size)
+      unless offsets.is_a?(Array) && offsets.length == 2
+        raise ConversionError, "tensor #{name}: malformed data_offsets"
+      end
+
+      begin_off, end_off = offsets
+      unless begin_off.is_a?(Integer) && end_off.is_a?(Integer) && begin_off >= 0 &&
+             end_off >= begin_off && end_off <= body_size
+        raise ConversionError, "tensor #{name}: data_offsets out of bounds"
+      end
+
+      [begin_off, end_off]
+    end
+
+    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 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
+    end
+  end
+end