neighbor 0.4.1 → 0.4.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 8aa6de2790d94de9411b0142836b2ad181a411e299fce4b98357b96ac4161183
-  data.tar.gz: 2924d7f15f5b36bc89ee72372c1bfeb373d99481269696a9a9dcc41f90201f38
+  metadata.gz: e8d611fd277cd48d309b2a087fdeb22f39f43d8ef81fcab57763bd5b4b2e48b3
+  data.tar.gz: fe0a5f7e4aa1ebd81f8c5849be67dc0f3c948d53d313b80259def04b9b7e9e84
 SHA512:
-  metadata.gz: 2bc1b3ee6d5b1ee0ab175b017e753cf958bd8ceb1ef2a23ba769770dfebf54eec251ac59c8f5f3b6ca56efcbad1763c34622b94924a017622c2f78fc8740f762
-  data.tar.gz: d946dda99833964582f63863b2d898fea6bf065312cf60aec873631df96195e1a54375606ad9c9cc0f767937cdb7ea38b0d9990efcbbeab15ccbb11f8a2020ef
+  metadata.gz: 0d9d9d0be9f2929f1eab7e5df52a0606ca8c220ea74e6ef349c2ab77e5da44bd4b98ed4d2271345b2d679882f1790ac54b8373c13a69f75607935322bdb68754
+  data.tar.gz: 834c7d6e26be9b6d8fc262280048e5104cd022587fed30d030f6c0edd849966a49a08bc793f1b49764b1fc5afc014c1463dd8acacfb5739bc507f4a77d3281a1

data/CHANGELOG.md

@@ -1,3 +1,11 @@
+## 0.4.3 (2024-09-02)
+
+- Added `rrf` method
+
+## 0.4.2 (2024-08-27)
+
+- Fixed error with `nil` values
+
 ## 0.4.1 (2024-08-26)
 
 - Added `precision` option

data/README.md

@@ -14,7 +14,7 @@ gem "neighbor"
 
 ## Choose An Extension
 
-Neighbor supports two extensions: [cube](https://www.postgresql.org/docs/current/cube.html) and [vector](https://github.com/pgvector/pgvector). cube ships with Postgres, while vector supports more dimensions and approximate nearest neighbor search.
+Neighbor supports two extensions: [cube](https://www.postgresql.org/docs/current/cube.html) and [pgvector](https://github.com/pgvector/pgvector). cube ships with Postgres, while pgvector supports more dimensions and approximate nearest neighbor search.
 
 For cube, run:
 
@@ -23,7 +23,7 @@ rails generate neighbor:cube
 rails db:migrate
 ```
 
-For vector, [install pgvector](https://github.com/pgvector/pgvector#installation) and run:
+For pgvector, [install the extension](https://github.com/pgvector/pgvector#installation) and run:
 
 ```sh
 rails generate neighbor:vector
@@ -70,17 +70,30 @@ Get the nearest neighbors to a vector
 Item.nearest_neighbors(:embedding, [0.9, 1.3, 1.1], distance: "euclidean").first(5)
 ```
 
-## Distance
+Records returned from `nearest_neighbors` will have a `neighbor_distance` attribute
+
+```ruby
+nearest_item = item.nearest_neighbors(:embedding, distance: "euclidean").first
+nearest_item.neighbor_distance
+```
+
+See the additional docs for:
+
+- [cube](#cube)
+- [pgvector](#pgvector)
+
+Or check out some [examples](#examples)
+
+## cube
+
+### Distance
 
 Supported values are:
 
 - `euclidean`
 - `cosine`
 - `taxicab`
-- `chebyshev` (cube only)
-- `inner_product` (vector only)
-- `hamming` (vector only)
-- `jaccard` (vector only)
+- `chebyshev`
 
 For cosine distance with cube, vectors must be normalized before being stored.
 
@@ -90,18 +103,11 @@ class Item < ApplicationRecord
 end
 ```
 
-For inner product with cube, see [this example](examples/disco_user_recs_cube.rb).
-
-Records returned from `nearest_neighbors` will have a `neighbor_distance` attribute
-
-```ruby
-nearest_item = item.nearest_neighbors(:embedding, distance: "euclidean").first
-nearest_item.neighbor_distance
-```
+For inner product with cube, see [this example](examples/disco/user_recs_cube.rb).
 
-## Dimensions
+### Dimensions
 
-The cube data type can have up to 100 dimensions by default. See the [Postgres docs](https://www.postgresql.org/docs/current/cube.html) for how to increase this. The vector data type can have up to 16,000 dimensions, and vectors with up to 2,000 dimensions can be indexed.
+The `cube` type can have up to 100 dimensions by default. See the [Postgres docs](https://www.postgresql.org/docs/current/cube.html) for how to increase this.
 
 For cube, it’s a good idea to specify the number of dimensions to ensure all records have the same number.
 
@@ -111,9 +117,26 @@ class Item < ApplicationRecord
 end
 ```
 
-## Indexing
+## pgvector
 
-For vector, add an approximate index to speed up queries. Create a migration with:
+### Distance
+
+Supported values are:
+
+- `euclidean`
+- `inner_product`
+- `cosine`
+- `taxicab`
+- `hamming`
+- `jaccard`
+
+### Dimensions
+
+The `vector` type can have up to 16,000 dimensions, and vectors with up to 2,000 dimensions can be indexed.
+
+### Indexing
+
+Add an approximate index to speed up queries. Create a migration with:
 
 ```ruby
 class AddIndexToItemsEmbedding < ActiveRecord::Migration[7.2]
@@ -139,7 +162,7 @@ Or the number of probes with IVFFlat
 Item.connection.execute("SET ivfflat.probes = 3")
 ```
 
-## Half-Precision Vectors
+### Half-Precision Vectors
 
 Use the `halfvec` type to store half-precision vectors
 
@@ -151,7 +174,7 @@ class AddEmbeddingToItems < ActiveRecord::Migration[7.2]
 end
 ```
 
-## Half-Precision Indexing
+### Half-Precision Indexing
 
 Index vectors at half precision for smaller indexes
 
@@ -169,7 +192,7 @@ Get the nearest neighbors
 Item.nearest_neighbors(:embedding, [0.9, 1.3, 1.1], distance: "euclidean", precision: "half").first(5)
 ```
 
-## Binary Vectors
+### Binary Vectors
 
 Use the `bit` type to store binary vectors
 
@@ -187,7 +210,7 @@ Get the nearest neighbors by Hamming distance
 Item.nearest_neighbors(:embedding, "101", distance: "hamming").first(5)
 ```
 
-## Binary Quantization
+### Binary Quantization
 
 Use expression indexing for binary quantization
 
@@ -199,7 +222,7 @@ class AddIndexToItemsEmbedding < ActiveRecord::Migration[7.2]
 end
 ```
 
-## Sparse Vectors
+### Sparse Vectors
 
 Use the `sparsevec` type to store sparse vectors
 
@@ -220,11 +243,12 @@ Item.nearest_neighbors(:embedding, embedding, distance: "euclidean").first(5)
 
 ## Examples
 
-- [OpenAI Embeddings](#openai-embeddings)
-- [Cohere Embeddings](#cohere-embeddings)
-- [Sentence Embeddings](#sentence-embeddings)
-- [Sparse Embeddings](#sparse-embeddings)
-- [Disco Recommendations](#disco-recommendations)
+- [Embeddings](#openai-embeddings) with OpenAI
+- [Binary embeddings](#cohere-embeddings) with Cohere
+- [Sentence embeddings](#sentence-embeddings) with Informers
+- [Hybrid search](#hybrid-search) with Informers
+- [Sparse search](#sparse-search) with Transformers.rb
+- [Recommendations](#disco-recommendations) with Disco
 
 ### OpenAI Embeddings
 
@@ -388,7 +412,7 @@ end
 Load a [model](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2)
 
 ```ruby
-model = Informers::Model.new("sentence-transformers/all-MiniLM-L6-v2")
+model = Informers.pipeline("embedding", "sentence-transformers/all-MiniLM-L6-v2")
 ```
 
 Pass your input
@@ -399,7 +423,7 @@ input = [
   "The cat is purring",
   "The bear is growling"
 ]
-embeddings = model.embed(input)
+embeddings = model.(input)
 ```
 
 Store the embeddings
@@ -421,7 +445,86 @@ document.nearest_neighbors(:embedding, distance: "cosine").first(5).map(&:conten
 
 See the [complete code](examples/informers/example.rb)
 
-### Sparse Embeddings
+### Hybrid Search
+
+You can use Neighbor for hybrid search with [Informers](https://github.com/ankane/informers).
+
+Generate a model
+
+```sh
+rails generate model Document content:text embedding:vector{768}
+rails db:migrate
+```
+
+And add `has_neighbors` and a scope for keyword search
+
+```ruby
+class Document < ApplicationRecord
+  has_neighbors :embedding
+
+  scope :search, ->(query) {
+    where("to_tsvector(content) @@ plainto_tsquery(?)", query)
+      .order(Arel.sql("ts_rank_cd(to_tsvector(content), plainto_tsquery(?)) DESC", query))
+  }
+end
+```
+
+Create some documents
+
+```ruby
+texts = [
+  "The dog is barking",
+  "The cat is purring",
+  "The bear is growling"
+]
+documents = Document.create!(texts.map { |v| {content: v} })
+```
+
+Generate an embedding for each document
+
+```ruby
+embed = Informers.pipeline("embedding", "Snowflake/snowflake-arctic-embed-m-v1.5")
+embed_options = {model_output: "sentence_embedding", pooling: "none"} # specific to embedding model
+embeddings = embed.(documents.map(&:content), **embed_options)
+
+documents.zip(embeddings) do |document, embedding|
+  document.update!(embedding: embedding)
+end
+```
+
+Perform keyword search
+
+```ruby
+query = "growling bear"
+keyword_results = Document.search(query).limit(20).load_async
+```
+
+And semantic search in parallel (the query prefix is specific to the [embedding model](https://huggingface.co/Snowflake/snowflake-arctic-embed-m-v1.5))
+
+```ruby
+query_prefix = "Represent this sentence for searching relevant passages: "
+query_embedding = embed.(query_prefix + query, **embed_options)
+semantic_results =
+  Document.nearest_neighbors(:embedding, query_embedding, distance: "cosine").limit(20).load_async
+```
+
+To combine the results, use Reciprocal Rank Fusion (RRF)
+
+```ruby
+Neighbor::Reranking.rrf(keyword_results, semantic_results)
+```
+
+Or a reranking model
+
+```ruby
+rerank = Informers.pipeline("reranking", "mixedbread-ai/mxbai-rerank-xsmall-v1")
+results = (keyword_results + semantic_results).uniq
+rerank.(query, results.map(&:content), top_k: 5).map { |v| results[v[:doc_id]] }
+```
+
+See the [complete code](examples/hybrid/example.rb)
+
+### Sparse Search
 
 You can generate sparse embeddings locally with [Transformers.rb](https://github.com/ankane/transformers-ruby).
 
@@ -533,7 +636,7 @@ movies = []
 recommender.item_ids.each do |item_id|
   movies << {name: item_id, factors: recommender.item_factors(item_id)}
 end
-Movie.insert_all!(movies)
+Movie.create!(movies)
 ```
 
 And get similar movies
@@ -543,7 +646,7 @@ movie = Movie.find_by(name: "Star Wars (1977)")
 movie.nearest_neighbors(:factors, distance: "cosine").first(5).map(&:name)
 ```
 
-See the complete code for [cube](examples/disco/item_recs_cube.rb) and [vector](examples/disco/item_recs_vector.rb)
+See the complete code for [cube](examples/disco/item_recs_cube.rb) and [pgvector](examples/disco/item_recs_vector.rb)
 
 ## History
 

data/lib/neighbor/reranking.rb

@@ -0,0 +1,27 @@
+module Neighbor
+  module Reranking
+    def self.rrf(first_ranking, *rankings, k: 60)
+      rankings.unshift(first_ranking)
+
+      ranks = []
+      results = []
+      rankings.each do |ranking|
+        ranks << ranking.map.with_index.to_h { |v, i| [v, i + 1] }
+        results.concat(ranking)
+      end
+
+      results =
+        results.uniq.map do |result|
+          score =
+            ranks.sum do |rank|
+              r = rank[result]
+              r ? 1.0 / (k + r) : 0.0
+            end
+
+          {result: result, score: score}
+        end
+
+      results.sort_by { |v| -v[:score] }
+    end
+  end
+end

data/lib/neighbor/type/cube.rb

@@ -6,7 +6,7 @@ module Neighbor
       end
 
       def serialize(value)
-        if value.respond_to?(:to_a)
+        if Utils.array?(value)
           value = value.to_a
           if value.first.is_a?(Array)
             value = value.map { |v| serialize_point(v) }.join(", ")
@@ -20,7 +20,7 @@ module Neighbor
       private
 
       def cast_value(value)
-        if value.respond_to?(:to_a)
+        if Utils.array?(value)
           value.to_a
         elsif value.is_a?(Numeric)
           [value]

data/lib/neighbor/type/halfvec.rb

@@ -6,7 +6,7 @@ module Neighbor
       end
 
       def serialize(value)
-        if value.respond_to?(:to_a)
+        if Utils.array?(value)
           value = "[#{value.to_a.map(&:to_f).join(",")}]"
         end
         super(value)
@@ -17,7 +17,7 @@ module Neighbor
       def cast_value(value)
         if value.is_a?(String)
           value[1..-1].split(",").map(&:to_f)
-        elsif value.respond_to?(:to_a)
+        elsif Utils.array?(value)
           value.to_a
         else
           raise "can't cast #{value.class.name} to halfvec"

data/lib/neighbor/type/sparsevec.rb

@@ -19,7 +19,7 @@ module Neighbor
           value
         elsif value.is_a?(String)
           SparseVector.from_text(value)
-        elsif value.respond_to?(:to_a)
+        elsif Utils.array?(value)
           value = SparseVector.new(value.to_a)
         else
           raise "can't cast #{value.class.name} to sparsevec"

data/lib/neighbor/type/vector.rb

@@ -6,7 +6,7 @@ module Neighbor
       end
 
       def serialize(value)
-        if value.respond_to?(:to_a)
+        if Utils.array?(value)
           value = "[#{value.to_a.map(&:to_f).join(",")}]"
         end
         super(value)
@@ -17,7 +17,7 @@ module Neighbor
       def cast_value(value)
         if value.is_a?(String)
           value[1..-1].split(",").map(&:to_f)
-        elsif value.respond_to?(:to_a)
+        elsif Utils.array?(value)
           value.to_a
         else
           raise "can't cast #{value.class.name} to vector"

data/lib/neighbor/utils.rb

@@ -38,5 +38,9 @@ module Neighbor
       # could also throw error
       norm > 0 ? value.map { |v| v / norm } : value
     end
+
+    def self.array?(value)
+      !value.nil? && value.respond_to?(:to_a)
+    end
   end
 end

data/lib/neighbor/version.rb

@@ -1,3 +1,3 @@
 module Neighbor
-  VERSION = "0.4.1"
+  VERSION = "0.4.3"
 end

data/lib/neighbor.rb

@@ -2,6 +2,7 @@
 require "active_support"
 
 # modules
+require_relative "neighbor/reranking"
 require_relative "neighbor/sparse_vector"
 require_relative "neighbor/utils"
 require_relative "neighbor/version"

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: neighbor
 version: !ruby/object:Gem::Version
-  version: 0.4.1
+  version: 0.4.3
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2024-08-27 00:00:00.000000000 Z
+date: 2024-09-02 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -40,6 +40,7 @@ files:
 - lib/neighbor.rb
 - lib/neighbor/model.rb
 - lib/neighbor/railtie.rb
+- lib/neighbor/reranking.rb
 - lib/neighbor/sparse_vector.rb
 - lib/neighbor/type/cube.rb
 - lib/neighbor/type/halfvec.rb