vecsearch 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 23bcc9e270a754cd56fbc2013084719f6567c9f284fe0504e2e0bec666035c67
+  data.tar.gz: 4ddb00bc0c604bf1a7f3fc47561804e97e894df4382c0611105cdea24c4a50d8
+SHA512:
+  metadata.gz: ab92719d0d0c8f3e14e002758c01b73493b4b4c578de81658e8bd952dd0a0eec1217d2061c4c5020a1246c6d74d5094ea12978d82661aec3ff97388cb8defad9
+  data.tar.gz: 23b554f0bdd20cfaea4f954dc989c3b01d52f174598d69d450de71a2bd69f5c720427d3b01389901a4ee9af4d9c6178fa82b5b2edd891ea332376deb52a0abc4

data/.rubocop.yml

@@ -0,0 +1,13 @@
+AllCops:
+  TargetRubyVersion: 2.6
+
+Style/StringLiterals:
+  Enabled: true
+  EnforcedStyle: double_quotes
+
+Style/StringLiteralsInInterpolation:
+  Enabled: true
+  EnforcedStyle: double_quotes
+
+Layout/LineLength:
+  Max: 120

data/Gemfile

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+source "https://rubygems.org"
+
+# Specify your gem's dependencies in vecsearch.gemspec
+gemspec
+
+gem "rake", "~> 13.0"
+
+gem "rubocop", "~> 1.21"

data/Gemfile.lock

@@ -0,0 +1,56 @@
+PATH
+  remote: .
+  specs:
+    vecsearch (0.1.0)
+      faiss
+      ffi
+
+GEM
+  remote: https://rubygems.org/
+  specs:
+    ast (2.4.2)
+    base64 (0.1.1)
+    faiss (0.3.0)
+      numo-narray
+      rice (>= 4.0.2)
+    ffi (1.16.3)
+    json (2.6.3)
+    language_server-protocol (3.17.0.3)
+    numo-narray (0.9.2.1)
+    parallel (1.23.0)
+    parser (3.2.2.4)
+      ast (~> 2.4.1)
+      racc
+    racc (1.7.1)
+    rainbow (3.1.1)
+    rake (13.0.6)
+    regexp_parser (2.8.1)
+    rexml (3.2.6)
+    rice (4.1.0)
+    rubocop (1.56.4)
+      base64 (~> 0.1.1)
+      json (~> 2.3)
+      language_server-protocol (>= 3.17.0)
+      parallel (~> 1.10)
+      parser (>= 3.2.2.3)
+      rainbow (>= 2.2.2, < 4.0)
+      regexp_parser (>= 1.8, < 3.0)
+      rexml (>= 3.2.5, < 4.0)
+      rubocop-ast (>= 1.28.1, < 2.0)
+      ruby-progressbar (~> 1.7)
+      unicode-display_width (>= 2.4.0, < 3.0)
+    rubocop-ast (1.29.0)
+      parser (>= 3.2.1.0)
+    ruby-progressbar (1.13.0)
+    unicode-display_width (2.5.0)
+
+PLATFORMS
+  arm64-darwin-22
+
+DEPENDENCIES
+  rake (~> 13.0)
+  rubocop (~> 1.21)
+  vecsearch!
+
+BUNDLED WITH
+   2.4.10

data/LICENSE.txt

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2023 Burke Libbey
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

data/README.md

@@ -0,0 +1,38 @@
+# Vecsearch
+
+Vecsearch is an all-in-one vector search library for ruby that uses a 4-bit (Q4_1) quantization of
+[gte-tiny](https://huggingface.co/TaylorAI/gte-tiny) by using [bert.cpp](https://github.com/skeskinen/bert.cpp) (a
+[GGML](https://ggml.ai/) implementation of [BERT](https://arxiv.org/abs/1810.04805) via
+[FFI](https://github.com/ffi/ffi)), and an in-process [FAISS](https://github.com/facebookresearch/faiss) index.
+
+Vecsearch embeds pre-built dynamic libraries for `libbert` and `libggml`, as well as a quantized model checkpoint for
+gte-tiny (total size: 14MB).
+
+Currently only ARM64 macOS is supported, purely because I haven't bothered to build other dylibs yet. There is nothing
+difficult about this.
+
+## Usage
+
+```ruby
+require 'vecsearch'
+
+vs = Vecsearch.new
+vs << "sharks with freaking laser beams"
+vs << "hello"
+vs << "the sky is green"
+
+puts(vs.nearest("hey there")) # => "hello"
+```
+## Bugs
+
+Yes
+
+## Limitations / TODO
+
+* I haven't got the mean-pooling part of gte-tiny working. It seems to work well
+  enough without it but we should do that and assert that ours generates
+  approximately the same embedding as the canonical model.
+* Batching looks unimplemented in bert.cpp; it would be nice for prefilling the
+  index.
+* Add more builds for platforms other than darwin/amd64.
+* Probably add a way to fetch an unquantized model, maybe other models entirely?

data/Rakefile

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+require "bundler/gem_tasks"
+require "rubocop/rake_task"
+
+RuboCop::RakeTask.new
+
+task default: :rubocop

data/demo.rb

@@ -0,0 +1,15 @@
+require 'vecsearch'
+
+vsi = Vecsearch.new
+vsi << "sharks with freaking laser beams"
+vsi << "hello"
+vsi << "the sky is green"
+
+puts("nearest record to 'hey there': #{vsi.nearest("hey there")}")
+
+t = Time.now
+N = 1000
+N.times { vsi.nearest("hey there") }
+elapsed = Time.now - t
+puts "#{(elapsed / (N/1000.0)).round(1)} ms per query"
+

data/dev.yml

@@ -0,0 +1,3 @@
+up:
+  - ruby: 3.2.2
+  - bundler

data/lib/vecsearch/gte_tiny.rb

@@ -0,0 +1,91 @@
+require 'ffi'
+# require 'narray'
+
+class Vecsearch
+  class GTETiny
+    VENDOR = File.expand_path('../../vendor', __dir__)
+
+    module CStdio
+      extend FFI::Library
+      ffi_lib 'c'
+
+      attach_function :fflush, [:pointer], :int
+    end
+
+    module Bert
+      extend FFI::Library
+      ffi_lib File.expand_path('libbert.dylib', VENDOR)
+
+      attach_function :bert_load_from_file, [:string], :pointer
+      attach_function :bert_n_embd, [:pointer], :int
+      attach_function :bert_encode_batch, [:pointer, :int, :int, :int, :pointer, :pointer], :void
+    end
+
+    GTE_BIN = File.expand_path('gte-tiny-q4_1.ggml.bin', VENDOR)
+    MAX_TOKENS = 512
+
+    def initialize(fname=GTE_BIN)
+      suppress_streams do
+        @ctx = Bert.bert_load_from_file(fname)
+        @n_embd = Bert.bert_n_embd(@ctx)
+        sleep(0.1)
+      end
+    end
+
+    def suppress_streams
+      prev_stdout = STDOUT.dup
+      STDOUT.reopen("/dev/null", "w")
+      STDOUT.sync = true
+      yield
+    ensure
+      CStdio.fflush(nil) # Regular STDOUT.flush doesn't do it.
+      STDOUT.reopen(prev_stdout)
+    end
+
+    def encode(sentence, n_threads: 1)
+      # Encode the sentence into token embeddings
+      token_embeddings = encode_batch([sentence], n_threads: 1)
+
+      # Pool the token embeddings into a sentence embedding
+      # For simplicity, we'll use an attention mask of all ones
+      attention_mask = Array.new(token_embeddings.first.length, 1)
+      # sentence_embedding = mean_pooling(token_embeddings, attention_mask)
+      # sentence_embedding
+      token_embeddings
+    end
+
+    def encode_batch(input, n_threads: 1)
+      # Create an array of pointers to the input strings
+      input_ptrs = input.map { |str| FFI::MemoryPointer.from_string(str) }
+
+      # Create a pointer to the array of input pointers
+      input_ptrs_ptr = FFI::MemoryPointer.new(:pointer, input_ptrs.length)
+      input_ptrs_ptr.write_array_of_pointer(input_ptrs)
+
+      # Create an output buffer for each input string
+      output_ptrs = input.map { FFI::MemoryPointer.new(:float, @n_embd) }
+
+      # Create a pointer to the array of output pointers
+      output_ptrs_ptr = FFI::MemoryPointer.new(:pointer, output_ptrs.length)
+      output_ptrs_ptr.write_array_of_pointer(output_ptrs)
+
+      Bert.bert_encode_batch(@ctx, n_threads, MAX_TOKENS, input.length, input_ptrs_ptr, output_ptrs_ptr)
+
+      # Convert the output buffers to Ruby arrays
+      output = output_ptrs.map { |ptr| ptr.read_array_of_float(@n_embd) }
+
+      output
+    end
+
+    # def mean_pooling(token_embeddings, attention_mask)
+    #   token_embeddings_na = NArray.to_na(token_embeddings)
+    #   attention_mask_na = NArray.to_na(attention_mask)
+
+    #   input_mask_expanded = attention_mask_na.expand_dims(-1).repeat(token_embeddings_na.shape[-1], -1)
+
+    #   sentence_embeddings = (token_embeddings_na * input_mask_expanded).sum(1) / input_mask_expanded.sum(1).clip(1e-9)
+
+    #   sentence_embeddings.to_a
+    # end
+  end
+end

data/lib/vecsearch/version.rb

@@ -0,0 +1,3 @@
+class Vecsearch
+  VERSION = "0.1.0"
+end

data/lib/vecsearch.rb

@@ -0,0 +1,27 @@
+require 'vecsearch/version'
+require 'faiss'
+
+class Vecsearch
+  autoload(:GTETiny, 'vecsearch/gte_tiny')
+
+  def initialize(records=[])
+    @faiss = Faiss::IndexFlatL2.new(384)
+    @gte = GTETiny.new
+    @texts = []
+    records.each { |rec| self << rec }
+  end
+
+  def <<(str)
+    emb = @gte.encode(str)
+    @texts << str
+    @faiss.add(emb)
+  end
+
+  def nearest(str) = query(str, 1)[0]
+
+  def query(str, n)
+    emb = @gte.encode(str)
+    _scores, indexes = @faiss.search(emb, 2)
+    [].tap { |res| indexes.each { |idx| res << @texts[idx] } }
+  end
+end

data/vecsearch.gemspec

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+require_relative "lib/vecsearch/version"
+
+Gem::Specification.new do |spec|
+  spec.name = "vecsearch"
+  spec.version = Vecsearch::VERSION
+  spec.authors = ["Burke Libbey"]
+  spec.email = ["burke.libbey@shopify.com"]
+
+  spec.summary = "All-in-one simple vector search class for ruby."
+  spec.description = "All-in-one simple vector search class for ruby."
+  spec.homepage = "https://github.com/shopify/vecsearch"
+  spec.license = "MIT"
+  spec.required_ruby_version = ">= 2.6.0"
+
+  spec.metadata["homepage_uri"] = spec.homepage
+  spec.metadata["source_code_uri"] = spec.homepage
+
+  # Specify which files should be added to the gem when it is released.
+  # The `git ls-files -z` loads the files in the RubyGem that have been added into git.
+  spec.files = Dir.chdir(__dir__) do
+    `git ls-files -z`.split("\x0").reject do |f|
+      (File.expand_path(f) == __FILE__) || f.start_with?(*%w[bin/ test/ spec/ features/ .git .circleci appveyor])
+    end
+  end
+  spec.bindir = "exe"
+  spec.executables = spec.files.grep(%r{\Aexe/}) { |f| File.basename(f) }
+  spec.require_paths = ["lib"]
+
+  spec.add_dependency 'ffi'
+  spec.add_dependency 'faiss'
+  # spec.add_dependency 'narray' # for mean pooling
+end

metadata

@@ -0,0 +1,88 @@
+--- !ruby/object:Gem::Specification
+name: vecsearch
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Burke Libbey
+autorequire:
+bindir: exe
+cert_chain: []
+date: 2023-10-06 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: ffi
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: faiss
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+description: All-in-one simple vector search class for ruby.
+email:
+- burke.libbey@shopify.com
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- ".rubocop.yml"
+- Gemfile
+- Gemfile.lock
+- LICENSE.txt
+- README.md
+- Rakefile
+- demo.rb
+- dev.yml
+- lib/vecsearch.rb
+- lib/vecsearch/gte_tiny.rb
+- lib/vecsearch/version.rb
+- vecsearch.gemspec
+- vendor/gte-tiny-q4_1.ggml.bin
+- vendor/libbert.dylib
+- vendor/libggml.dylib
+homepage: https://github.com/shopify/vecsearch
+licenses:
+- MIT
+metadata:
+  homepage_uri: https://github.com/shopify/vecsearch
+  source_code_uri: https://github.com/shopify/vecsearch
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 2.6.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.4.10
+signing_key:
+specification_version: 4
+summary: All-in-one simple vector search class for ruby.
+test_files: []