Nearest neighbor search for Rails and Postgres
Add this line to your application’s Gemfile:
gem "neighbor"
Neighbor supports two extensions: cube and vector. cube ships with Postgres, while vector supports more dimensions and approximate nearest neighbor search.
For cube, run:
rails generate neighbor:cube
rails db:migrate
For vector, install pgvector and run:
rails generate neighbor:vector
rails db:migrate
Create a migration
class AddEmbeddingToItems < ActiveRecord::Migration[7.2]
def change
add_column :items, :embedding, :cube
# or
add_column :items, :embedding, :vector, limit: 3 # dimensions
end
end
Add to your model
class Item < ApplicationRecord
has_neighbors :embedding
end
Update the vectors
item.update(embedding: [1.0, 1.2, 0.5])
Get the nearest neighbors to a record
item.nearest_neighbors(:embedding, distance: "euclidean").first(5)
Get the nearest neighbors to a vector
Item.nearest_neighbors(:embedding, [0.9, 1.3, 1.1], distance: "euclidean").first(5)
Supported values are:
euclidean
cosine
taxicab
chebyshev
(cube only)inner_product
(vector only)hamming
(vector only)jaccard
(vector only)
For cosine distance with cube, vectors must be normalized before being stored.
class Item < ApplicationRecord
has_neighbors :embedding, normalize: true
end
For inner product with cube, see this example.
Records returned from nearest_neighbors
will have a neighbor_distance
attribute
nearest_item = item.nearest_neighbors(:embedding, distance: "euclidean").first
nearest_item.neighbor_distance
The cube data type can have up to 100 dimensions by default. See the Postgres docs 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.
For cube, it’s a good idea to specify the number of dimensions to ensure all records have the same number.
class Item < ApplicationRecord
has_neighbors :embedding, dimensions: 3
end
For vector, add an approximate index to speed up queries. Create a migration with:
class AddIndexToItemsEmbedding < ActiveRecord::Migration[7.2]
def change
add_index :items, :embedding, using: :hnsw, opclass: :vector_l2_ops
# or
add_index :items, :embedding, using: :ivfflat, opclass: :vector_l2_ops
end
end
Use :vector_cosine_ops
for cosine distance and :vector_ip_ops
for inner product.
Set the size of the dynamic candidate list with HNSW
Item.connection.execute("SET hnsw.ef_search = 100")
Or the number of probes with IVFFlat
Item.connection.execute("SET ivfflat.probes = 3")
Use the halfvec
type to store half-precision vectors
class AddEmbeddingToItems < ActiveRecord::Migration[7.2]
def change
add_column :items, :embedding, :halfvec, limit: 3 # dimensions
end
end
Index vectors at half precision for smaller indexes
class AddIndexToItemsEmbedding < ActiveRecord::Migration[7.2]
def change
add_index :items, "(embedding::halfvec(3)) vector_l2_ops", using: :hnsw
end
end
Get the nearest neighbors [unreleased]
Item.nearest_neighbors(:embedding, [0.9, 1.3, 1.1], distance: "euclidean", precision: "half").first(5)
Use the bit
type to store binary vectors
class AddEmbeddingToItems < ActiveRecord::Migration[7.2]
def change
add_column :items, :embedding, :bit, limit: 3 # dimensions
end
end
Get the nearest neighbors by Hamming distance
Item.nearest_neighbors(:embedding, "101", distance: "hamming").first(5)
Use expression indexing for binary quantization
class AddIndexToItemsEmbedding < ActiveRecord::Migration[7.2]
def change
add_index :items, "(binary_quantize(embedding)::bit(3)) bit_hamming_ops", using: :hnsw
end
end
Use the sparsevec
type to store sparse vectors
class AddEmbeddingToItems < ActiveRecord::Migration[7.2]
def change
add_column :items, :embedding, :sparsevec, limit: 3 # dimensions
end
end
Get the nearest neighbors
embedding = Neighbor::SparseVector.new({0 => 0.9, 1 => 1.3, 2 => 1.1}, 3)
Item.nearest_neighbors(:embedding, embedding, distance: "euclidean").first(5)
Generate a model
rails generate model Document content:text embedding:vector{1536}
rails db:migrate
And add has_neighbors
class Document < ApplicationRecord
has_neighbors :embedding
end
Create a method to call the embeddings API
def fetch_embeddings(input)
url = "https://api.openai.com/v1/embeddings"
headers = {
"Authorization" => "Bearer #{ENV.fetch("OPENAI_API_KEY")}",
"Content-Type" => "application/json"
}
data = {
input: input,
model: "text-embedding-3-small"
}
response = Net::HTTP.post(URI(url), data.to_json, headers).tap(&:value)
JSON.parse(response.body)["data"].map { |v| v["embedding"] }
end
Pass your input
input = [
"The dog is barking",
"The cat is purring",
"The bear is growling"
]
embeddings = fetch_embeddings(input)
Store the embeddings
documents = []
input.zip(embeddings) do |content, embedding|
documents << {content: content, embedding: embedding}
end
Document.insert_all!(documents)
And get similar documents
document = Document.first
document.nearest_neighbors(:embedding, distance: "cosine").first(5).map(&:content)
See the complete code
Generate a model
rails generate model Document content:text embedding:bit{1024}
rails db:migrate
And add has_neighbors
class Document < ApplicationRecord
has_neighbors :embedding
end
Create a method to call the embed API
def fetch_embeddings(input, input_type)
url = "https://api.cohere.com/v1/embed"
headers = {
"Authorization" => "Bearer #{ENV.fetch("CO_API_KEY")}",
"Content-Type" => "application/json"
}
data = {
texts: input,
model: "embed-english-v3.0",
input_type: input_type,
embedding_types: ["ubinary"]
}
response = Net::HTTP.post(URI(url), data.to_json, headers).tap(&:value)
JSON.parse(response.body)["embeddings"]["ubinary"].map { |e| e.map { |v| v.chr.unpack1("B*") }.join }
end
Pass your input
input = [
"The dog is barking",
"The cat is purring",
"The bear is growling"
]
embeddings = fetch_embeddings(input, "search_document")
Store the embeddings
documents = []
input.zip(embeddings) do |content, embedding|
documents << {content: content, embedding: embedding}
end
Document.insert_all!(documents)
Embed the search query
query = "forest"
query_embedding = fetch_embeddings([query], "search_query")[0]
And search the documents
Document.nearest_neighbors(:embedding, query_embedding, distance: "hamming").first(5).map(&:content)
See the complete code
You can generate embeddings locally with Transformers.rb.
Generate a model
rails generate model Document content:text embedding:vector{384}
rails db:migrate
And add has_neighbors
class Document < ApplicationRecord
has_neighbors :embedding
end
Load a model
model = Transformers::SentenceTransformer.new("sentence-transformers/all-MiniLM-L6-v2")
Pass your input
input = [
"The dog is barking",
"The cat is purring",
"The bear is growling"
]
embeddings = model.encode(input)
Store the embeddings
documents = []
input.zip(embeddings) do |content, embedding|
documents << {content: content, embedding: embedding}
end
Document.insert_all!(documents)
And get similar documents
document = Document.first
document.nearest_neighbors(:embedding, distance: "cosine").first(5).map(&:content)
See the complete code
You can generate sparse embeddings locally with Transformers.rb.
Generate a model
rails generate model Document content:text embedding:sparsevec{30522}
rails db:migrate
And add has_neighbors
class Document < ApplicationRecord
has_neighbors :embedding
end
Load a model to generate embeddings
class EmbeddingModel
def initialize(model_id)
@model = Transformers::AutoModelForMaskedLM.from_pretrained(model_id)
@tokenizer = Transformers::AutoTokenizer.from_pretrained(model_id)
@special_token_ids = @tokenizer.special_tokens_map.map { |_, token| @tokenizer.vocab[token] }
end
def embed(input)
feature = @tokenizer.(input, padding: true, truncation: true, return_tensors: "pt", return_token_type_ids: false)
output = @model.(**feature)[0]
values = Torch.max(output * feature[:attention_mask].unsqueeze(-1), dim: 1)[0]
values = Torch.log(1 + Torch.relu(values))
values[0.., @special_token_ids] = 0
values.to_a
end
end
model = EmbeddingModel.new("opensearch-project/opensearch-neural-sparse-encoding-v1")
Pass your input
input = [
"The dog is barking",
"The cat is purring",
"The bear is growling"
]
embeddings = model.embed(input)
Store the embeddings
documents = []
input.zip(embeddings) do |content, embedding|
documents << {content: content, embedding: Neighbor::SparseVector.new(embedding)}
end
Document.insert_all!(documents)
Embed the search query
query = "forest"
query_embedding = model.embed([query])[0]
And search the documents
Document.nearest_neighbors(:embedding, Neighbor::SparseVector.new(query_embedding), distance: "inner_product").first(5).map(&:content)
See the complete code
You can use Neighbor for online item-based recommendations with Disco. We’ll use MovieLens data for this example.
Generate a model
rails generate model Movie name:string factors:cube
rails db:migrate
And add has_neighbors
class Movie < ApplicationRecord
has_neighbors :factors, dimensions: 20, normalize: true
end
Fit the recommender
data = Disco.load_movielens
recommender = Disco::Recommender.new(factors: 20)
recommender.fit(data)
Store the item factors
movies = []
recommender.item_ids.each do |item_id|
movies << {name: item_id, factors: recommender.item_factors(item_id)}
end
Movie.insert_all!(movies)
And get similar movies
movie = Movie.find_by(name: "Star Wars (1977)")
movie.nearest_neighbors(:factors, distance: "cosine").first(5).map(&:name)
See the complete code for cube and vector
View the changelog
Everyone is encouraged to help improve this project. Here are a few ways you can help:
- Report bugs
- Fix bugs and submit pull requests
- Write, clarify, or fix documentation
- Suggest or add new features
To get started with development:
git clone https://github.com/ankane/neighbor.git
cd neighbor
bundle install
createdb neighbor_test
bundle exec rake test