static_embeddings 0.1.1

data/lib/static_embeddings/unicode_tables.rb

@@ -0,0 +1,127 @@
+module StaticEmbeddings
+
+  module UnicodeTables
+
+    MAP_SCAN_LIMIT = 0x2FFFF
+    CODESPACE_LIMIT = 0x10FFFF
+    SURROGATES = (0xD800..0xDFFF)
+
+    ENTRY_PACK = "V6"
+
+    RE_MN = /\A\p{Mn}\z/
+    RE_PUNCT = /\A\p{P}\z/
+    RE_CONTROL = /\A(?:\p{Cc}|\p{Cf}|\p{Co})\z/
+    RE_WHITESPACE = /\A(?:\p{Zs}|[\u0085\u2028\u2029])\z/
+
+    class RangeAccumulator
+      def initialize
+        @ranges = []
+      end
+
+      def add(codepoint)
+        last = @ranges.last
+        if last && last[1] + 1 == codepoint
+          last[1] = codepoint
+        else
+          @ranges << [codepoint, codepoint]
+        end
+      end
+
+      def to_a
+        @ranges
+      end
+    end
+
+    module_function
+
+    def build
+      lower = []
+      nfd = []
+      mn = RangeAccumulator.new
+      punct = RangeAccumulator.new
+      control = RangeAccumulator.new
+      whitespace = RangeAccumulator.new
+
+      (0..CODESPACE_LIMIT).each do |cp|
+        if SURROGATES.cover?(cp)
+          control.add(cp)
+          next
+        end
+
+        ch = begin
+          cp.chr(Encoding::UTF_8)
+        rescue RangeError
+          next
+        end
+
+        mn.add(cp) if RE_MN.match?(ch)
+        punct.add(cp) if RE_PUNCT.match?(ch)
+        control.add(cp) if RE_CONTROL.match?(ch)
+        whitespace.add(cp) if RE_WHITESPACE.match?(ch)
+
+        next if cp > MAP_SCAN_LIMIT
+
+        down = ch.downcase
+        lower << [cp, down.codepoints] if down != ch
+
+        decomposed = ch.unicode_normalize(:nfd)
+        nfd << [cp, decomposed.codepoints] if decomposed != ch
+      end
+
+      { "lowercase" => lower, "NFD" => nfd }.each do |name, table|
+        long = table.find { |(_, out)| out.length > 4 }
+        next unless long
+
+        raise "#{name} mapping for U+#{format('%04X', long[0])} needs #{long[1].length} slots"
+      end
+
+      {
+        lower: lower,
+        nfd: nfd,
+        mn: mn.to_a,
+        punct: punct.to_a,
+        control: control.to_a,
+        whitespace: whitespace.to_a
+      }
+    end
+
+    def pack(tables)
+      blob = +""
+      blob.force_encoding(Encoding::BINARY)
+      blob << [
+        tables[:lower].length,
+        tables[:nfd].length,
+        tables[:mn].length,
+        tables[:punct].length,
+        tables[:control].length,
+        tables[:whitespace].length,
+        0, 0
+      ].pack("V8")
+
+      %i[lower nfd].each do |key|
+        tables[key].sort_by(&:first).each do |(cp, out)|
+          padded = out + [0] * (4 - out.length)
+          blob << ([cp, out.length] + padded).pack(ENTRY_PACK)
+        end
+      end
+
+      %i[mn punct control whitespace].each do |key|
+        tables[key].sort_by(&:first).each { |(lo, hi)| blob << [lo, hi].pack("V2") }
+      end
+
+      blob
+    end
+
+    def packed
+      @packed ||= pack(build)
+    end
+
+    def reset!
+      @packed = nil
+    end
+
+    def source_stamp
+      "ruby-#{RUBY_VERSION}p#{RUBY_PATCHLEVEL}"
+    end
+  end
+end

data/lib/static_embeddings/version.rb

@@ -0,0 +1,3 @@
+module StaticEmbeddings
+  VERSION = "0.1.1"
+end

data/lib/static_embeddings.rb

@@ -0,0 +1,118 @@
+require "json"
+require "static_embeddings/version"
+
+begin
+  require "static_embeddings/static_embeddings"
+rescue LoadError
+  ext_dir = File.expand_path("static_embeddings", __dir__)
+  so_path = %w[.so .bundle].lazy.map { |ext| File.join(ext_dir, "static_embeddings#{ext}") }
+                                .find { |path| File.file?(path) }
+  unless so_path
+    raise LoadError, "Could not find the compiled StaticEmbeddings extension. " \
+                     "Run: bundle exec rake compile"
+  end
+
+  require so_path
+end
+
+require "static_embeddings/errors"
+require "static_embeddings/paths"
+require "static_embeddings/format"
+require "static_embeddings/unicode_tables"
+require "static_embeddings/safetensors"
+require "static_embeddings/converter"
+require "static_embeddings/reference"
+require "static_embeddings/model"
+
+module StaticEmbeddings
+  class << self
+    def load(path, verify: false)
+      expanded = File.expand_path(path.to_s)
+      raise ModelNotFound, "no model at #{expanded}" unless File.file?(expanded)
+
+      if verify
+        result = Format.verify(expanded)
+        raise InvalidModelError, "checksum mismatch for #{expanded}" unless result[:ok]
+      end
+
+      Model.new(expanded)
+    end
+
+    def load_builtin(name = :demo, verify: false)
+      load(Paths.builtin_path(name), verify: verify)
+    end
+
+    def builtin_available?(name = :demo)
+      Paths.builtin_available?(name)
+    end
+
+    def cache_dir
+      Paths.cache_dir
+    end
+
+    def model_path(model_id)
+      Paths.model_path(model_id)
+    end
+
+    def load_model(model_id, verify: false)
+      path = model_path(model_id)
+      unless File.file?(path)
+        raise ModelNotFound,
+              "model #{model_id.inspect} is not installed. " \
+              "Convert it first: static_embeddings convert <hf-dir> --id #{model_id}"
+      end
+      load(path, verify: verify)
+    end
+
+    def convert(source_dir, output_path:, model_id: nil, max_tokens: Converter::REFERENCE_MAX_TOKENS)
+      converter = Converter.new(source_dir)
+      converter.convert(output_path: output_path, model_id: model_id, max_tokens: max_tokens)
+      converter.report
+    end
+
+    def verify(path)
+      Format.verify(File.expand_path(path.to_s))
+    end
+
+    def unpack(blob, dim, format: :f32)
+      raise ArgumentError, "dim must be positive" unless dim.to_i.positive?
+
+      floats =
+        case normalize_format(format)
+        when :f32
+          raise ArgumentError, "f32 blob byte size must be a multiple of 4" unless (blob.bytesize % 4).zero?
+
+          blob.unpack("e*")
+        when :f16
+          raise ArgumentError, "f16 blob byte size must be a multiple of 2" unless (blob.bytesize % 2).zero?
+
+          decode_f16(blob)
+        end
+
+      raise ArgumentError, "blob is not a multiple of dim" unless (floats.length % dim).zero?
+
+      floats.each_slice(dim).to_a
+    end
+
+    def pack(rows, format: :f32)
+      flat = rows.first.is_a?(Array) ? rows.flatten(1) : rows
+
+      case normalize_format(format)
+      when :f32 then flat.map(&:to_f).pack("e*")
+      when :f16 then encode_f16(flat.map(&:to_f))
+      end
+    end
+
+    def normalize_format(format)
+      case format&.to_sym
+      when nil, :f32, :float32
+        :f32
+      when :f16, :float16
+        :f16
+      else
+        raise ArgumentError, "unsupported embedding format #{format.inspect} (expected :f32 or :f16)"
+      end
+    end
+
+  end
+end

data/static_embeddings.gemspec

@@ -0,0 +1,45 @@
+require_relative "lib/static_embeddings/version"
+
+Gem::Specification.new do |spec|
+  spec.name    = "static_embeddings"
+  spec.version = StaticEmbeddings::VERSION
+  spec.authors = ["Roman Haydarov"]
+  spec.email   = ["romnhajdarov@gmail.com"]
+
+  spec.summary     = "Fast local text embeddings for Ruby — no ONNX, no Rust, no network"
+  spec.description = "A small C-extension runtime for Model2Vec-style static embedding " \
+                     "models. Models are converted offline into a flat mmap-able .semb " \
+                     "file; at runtime the gem tokenizes (BERT WordPiece), looks up rows " \
+                     "and mean-pools them. Releases the GVL on large native work, rejects internal " \
+                     "thread fan-out, and links nothing but libc."
+  spec.homepage = "https://github.com/roman-haidarov/static_embeddings"
+  spec.license  = "MIT"
+  spec.required_ruby_version = ">= 3.1.0"
+
+  spec.metadata["homepage_uri"]    = spec.homepage
+  spec.metadata["source_code_uri"] = spec.homepage
+  spec.metadata["changelog_uri"]   = "#{spec.homepage}/blob/main/CHANGELOG.md"
+
+  spec.files = Dir[
+    "lib/**/*.rb",
+    "lib/models/*.semb",
+    "exe/*",
+    "ext/**/*.{c,h,rb}",
+    "docs/**/*.md",
+    "tools/*.rb",
+    "static_embeddings.gemspec",
+    "README.md",
+    "CHANGELOG.md",
+    "LICENSE.txt"
+  ].reject { |path| File.basename(path) == ".DS_Store" }
+
+  spec.bindir        = "exe"
+  spec.executables   = ["static_embeddings"]
+  spec.require_paths = ["lib"]
+  spec.extensions    = ["ext/static_embeddings/extconf.rb"]
+
+  spec.add_development_dependency "bundler", "~> 2.0"
+  spec.add_development_dependency "minitest", "~> 5.0"
+  spec.add_development_dependency "rake", "~> 13.0"
+  spec.add_development_dependency "rake-compiler", "~> 1.2"
+end

data/tools/benchmark.rb

@@ -0,0 +1,38 @@
+#!/usr/bin/env ruby
+
+$LOAD_PATH.unshift(File.expand_path("../lib", __dir__))
+require "static_embeddings"
+require "benchmark"
+
+path = ARGV[0] || File.expand_path("../lib/models/demo.semb", __dir__)
+model = StaticEmbeddings.load(path)
+model.warmup!
+
+CORPUS = (1..5_000).map do |i|
+  "локальный поиск по тексту номер #{i} hello world static embedding vector search"
+end
+
+tokens = CORPUS.sum { |t| model.tokenize(t).length }
+bytes = CORPUS.sum(&:bytesize)
+
+def timed(label, iterations)
+  elapsed = Benchmark.realtime { iterations.times { yield } }
+  [label, elapsed]
+end
+
+single = Benchmark.realtime { 2_000.times { model.embed(CORPUS.first) } }
+batch1 = Benchmark.realtime { model.embed_batch(CORPUS) }
+tok_only = Benchmark.realtime { CORPUS.each { |t| model.tokenize(t) } }
+
+per_token_ns = (batch1 / tokens) * 1e9
+per_token_per_100dim = per_token_ns / (model.dim / 100.0)
+
+puts "model            #{model.model_id} dim=#{model.dim} vocab=#{model.vocab_size}"
+puts "corpus           #{CORPUS.length} texts, #{tokens} tokens, #{bytes} bytes"
+puts
+puts "single embed     #{((single / 2_000) * 1e6).round(1)} us/call"
+puts "batch           #{(CORPUS.length / batch1).round} texts/s, #{(tokens / batch1).round} tokens/s"
+puts "tokenize only    #{((tok_only / bytes) * 1e9).round(1)} ns/input byte"
+puts
+puts "normalised       #{per_token_ns.round(1)} ns/token"
+puts "                 #{per_token_per_100dim.round(1)} ns/token/100dim"

data/tools/build_demo_model.rb

@@ -0,0 +1,16 @@
+#!/usr/bin/env ruby
+
+$LOAD_PATH.unshift(File.expand_path("../lib", __dir__))
+require "static_embeddings"
+require_relative "make_fixture_model"
+require "fileutils"
+require "tmpdir"
+
+target = File.expand_path("../lib/models/demo.semb", __dir__)
+FileUtils.mkdir_p(File.dirname(target))
+
+Dir.mktmpdir do |dir|
+  FixtureModel.write(dir)
+  report = StaticEmbeddings.convert(dir, output_path: target, model_id: "demo/tiny-wordpiece")
+  puts "wrote #{target} (#{report[:bytes]} bytes, vocab #{report[:vocab_size]}, dim #{report[:dim]})"
+end

data/tools/check_model2vec_parity.rb

@@ -0,0 +1,107 @@
+require "json"
+require "optparse"
+require "static_embeddings"
+
+options = {
+  min_cosine: 1.0 - 1e-6,
+  max_abs: 1e-5,
+  ids: true
+}
+
+parser = OptionParser.new do |opts|
+  opts.on("--model PATH") { |value| options[:model] = value }
+  opts.on("--oracle PATH") { |value| options[:oracle] = value }
+  opts.on("--min-cosine N", Float) { |value| options[:min_cosine] = value }
+  opts.on("--max-abs N", Float) { |value| options[:max_abs] = value }
+  opts.on("--[no-]ids") { |value| options[:ids] = value }
+end
+parser.parse!(ARGV)
+unless options[:model] && options[:oracle]
+  abort "usage: ruby -Ilib tools/check_model2vec_parity.rb --model MODEL.semb --oracle oracle.json"
+end
+
+model = StaticEmbeddings.load(options[:model], verify: true)
+payload = JSON.parse(File.read(options[:oracle], encoding: "UTF-8"))
+rows = payload.is_a?(Hash) ? payload.fetch("rows") : payload
+oracle_max_length = payload.is_a?(Hash) ? payload["max_length"] : nil
+
+if oracle_max_length && oracle_max_length != model.max_tokens
+  warn "WARNING oracle max_length=#{oracle_max_length} but model.max_tokens=#{model.max_tokens}"
+end
+
+def dot(a, b)
+  a.zip(b).sum { |x, y| x * y }
+end
+
+def norm(a)
+  Math.sqrt(a.sum { |x| x * x })
+end
+
+min_cosine = 1.0
+max_abs_all = 0.0
+vector_failures = []
+id_failures = []
+id_rows_checked = 0
+
+rows.each_with_index do |row, i|
+  text = row.fetch("text")
+  ref = row.fetch("vector")
+  ref_ids = row["token_ids"]
+
+  problems = []
+
+  if options[:ids] && ref_ids
+    id_rows_checked += 1
+    got_ids = model.tokenize(text)
+    if got_ids != ref_ids
+      id_failures << i
+      first = got_ids.zip(ref_ids).index { |a, b| a != b } || [got_ids.length, ref_ids.length].min
+      problems << "ids differ at #{first} (got #{got_ids.length}, ref #{ref_ids.length})"
+    end
+  end
+
+  got = model.embed(text).unpack("e*")
+  max_abs = ref.zip(got).map { |a, b| (a - b).abs }.max || 0.0
+  ref_zero = ref.all?(&:zero?)
+  got_zero = got.all?(&:zero?)
+  cosine =
+    if ref_zero && got_zero
+      1.0
+    elsif ref_zero || got_zero
+      0.0
+    else
+      dot(ref, got) / (norm(ref) * norm(got))
+    end
+
+  min_cosine = [min_cosine, cosine].min
+  max_abs_all = [max_abs_all, max_abs].max
+
+  unless cosine >= options[:min_cosine] && max_abs <= options[:max_abs]
+    vector_failures << i
+    problems << "vector out of tolerance"
+  end
+
+  status = problems.empty? ? "ok" : "FAIL"
+  label = text.bytesize > 64 ? "#{text[0, 32].inspect}...(#{text.bytesize}B)" : text.inspect
+  line = format("%s idx=%02d cos=%.10f max_abs=%.8g bytes=%d text=%s",
+                status, i, cosine, max_abs, text.bytesize, label)
+  line += "  [#{problems.join('; ')}]" unless problems.empty?
+  puts line
+end
+
+puts "rows=#{rows.length}"
+puts "id_rows_checked=#{id_rows_checked}"
+puts "min_cosine=#{min_cosine}"
+puts "max_abs_all=#{max_abs_all}"
+puts "token_id_failures=#{id_failures.inspect}"
+puts "vector_failures=#{vector_failures.inspect}"
+
+if id_rows_checked.zero? && options[:ids]
+  warn "WARNING oracle has no token_ids; regenerate it with the current tools/model2vec_oracle.py"
+end
+
+unless id_failures.empty? && vector_failures.empty?
+  abort "parity failed: token ids #{id_failures.inspect}, vectors #{vector_failures.inspect}"
+end
+
+puts "parity OK"

data/tools/make_fixture_model.rb

@@ -0,0 +1,165 @@
+#!/usr/bin/env ruby
+
+require "json"
+require "fileutils"
+
+$LOAD_PATH.unshift(File.expand_path("../lib", __dir__))
+require "static_embeddings/version"
+require "static_embeddings/safetensors"
+
+module FixtureModel
+  DIM = 8
+  SPECIALS = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"].freeze
+
+  WORDS = %w[
+    hello world ruby gem static embedding vector search text token
+    the a of and for local fast model cafe naive
+    привет мир поиск текст вектор модель локальный быстрый
+  ].freeze
+
+  PUNCTUATION = %w[. , ! ? - : ; ' " ( ) /].freeze
+  SUBWORDS = %w[##ing ##ed ##s ##er ##ик ##ов ##ый].freeze
+
+  module_function
+
+  def vocab_tokens
+    [SPECIALS, ascii_letters, digits, cyrillic_letters, ascii_subwords,
+     cyrillic_subwords, PUNCTUATION, WORDS, SUBWORDS].each_with_object([]) do |group, tokens|
+      group.each { |token| tokens << token }
+    end.uniq
+  end
+
+  def matrix_floats(vocab_size)
+    state = 0x12345678
+    Array.new(vocab_size * DIM) do
+      state = (state * 1_103_515_245 + 12_345) & 0x7FFFFFFF
+      ((state % 20_000) - 10_000) / 10_000.0
+    end
+  end
+
+  def write(dir)
+    FileUtils.mkdir_p(dir)
+    tokens = vocab_tokens
+    write_tokenizer(dir, tokens)
+    write_config(dir)
+    write_tokenizer_config(dir)
+    write_tensor(dir, tokens.length)
+    { dir: dir, vocab_size: tokens.length, dim: DIM }
+  end
+
+  def ascii_letters
+    ("a".."z").to_a
+  end
+
+  def digits
+    ("0".."9").to_a
+  end
+
+  def cyrillic_letters
+    ("а".."я").to_a
+  end
+
+  def ascii_subwords
+    ascii_letters.map { |char| "###{char}" }
+  end
+
+  def cyrillic_subwords
+    cyrillic_letters.map { |char| "###{char}" }
+  end
+
+  def write_tokenizer(dir, tokens)
+    File.write(File.join(dir, "tokenizer.json"), JSON.pretty_generate(tokenizer_payload(tokens)))
+  end
+
+  def tokenizer_payload(tokens)
+    {
+      "version" => "1.0",
+      "truncation" => nil,
+      "padding" => nil,
+      "added_tokens" => added_tokens,
+      "normalizer" => normalizer,
+      "pre_tokenizer" => { "type" => "BertPreTokenizer" },
+      "post_processor" => nil,
+      "decoder" => nil,
+      "model" => wordpiece_model(tokens)
+    }
+  end
+
+  def added_tokens
+    SPECIALS.each_with_index.map do |content, index|
+      {
+        "id" => index,
+        "content" => content,
+        "single_word" => false,
+        "lstrip" => false,
+        "rstrip" => false,
+        "normalized" => false,
+        "special" => true
+      }
+    end
+  end
+
+  def normalizer
+    {
+      "type" => "BertNormalizer",
+      "clean_text" => true,
+      "handle_chinese_chars" => true,
+      "strip_accents" => nil,
+      "lowercase" => true
+    }
+  end
+
+  def wordpiece_model(tokens)
+    {
+      "type" => "WordPiece",
+      "unk_token" => "[UNK]",
+      "continuing_subword_prefix" => "##",
+      "max_input_chars_per_word" => 100,
+      "vocab" => tokens.each_with_index.to_h
+    }
+  end
+
+  def write_config(dir)
+    File.write(File.join(dir, "config.json"), JSON.pretty_generate(config_payload))
+  end
+
+  def config_payload
+    {
+      "model_type" => "model2vec",
+      "hidden_dim" => DIM,
+      "seq_length" => 1_000_000,
+      "normalize" => true,
+      "tokenizer_name" => "fixture/tiny-wordpiece"
+    }
+  end
+
+  def write_tokenizer_config(dir)
+    File.write(File.join(dir, "tokenizer_config.json"), JSON.pretty_generate(tokenizer_config_payload))
+  end
+
+  def tokenizer_config_payload
+    {
+      "tokenizer_class" => "BertTokenizer",
+      "do_lower_case" => true,
+      "strip_accents" => nil,
+      "tokenize_chinese_chars" => true,
+      "model_max_length" => 1_000_000,
+      "unk_token" => "[UNK]"
+    }
+  end
+
+  def write_tensor(dir, vocab_size)
+    StaticEmbeddings::Safetensors.write(
+      File.join(dir, "model.safetensors"),
+      "embeddings",
+      [vocab_size, DIM],
+      matrix_floats(vocab_size)
+    )
+  end
+end
+
+if $PROGRAM_NAME == __FILE__
+  target = ARGV[0] || File.expand_path("../test/fixtures/tiny-wordpiece", __dir__)
+  info = FixtureModel.write(target)
+  puts "wrote #{info[:vocab_size]} tokens, dim #{info[:dim]} -> #{info[:dir]}"
+end