neighbor

> 编程语言

Open source

Nearest neighbor search for Rails

840 stars0 likes2 views

Website GitHub

About

Nearest neighbor search for Rails

Neighbor

Nearest neighbor search for Rails

Supports:

Postgres (pgvector and cube)
MariaDB 11.8+
MySQL 9.7+ (searching requires HeatWave)
SQLite

Also available for Redis and S3 Vectors

Installation

Add this line to your application’s Gemfile:

ruby

gem "neighbor"

For Postgres

Neighbor supports two extensions for Postgres: pgvector and cube. cube ships with Postgres, while pgvector supports more dimensions and approximate nearest neighbor search.

For pgvector, install the extension and run:

bash

rails generate neighbor:vector
rails db:migrate

For cube, run:

bash

rails generate neighbor:cube
rails db:migrate

Getting Started

Create a migration

ruby

class AddEmbeddingToItems  0.9, 1 => 1.3, 2 => 1.1}, 3)
Item.nearest_neighbors(:embedding, embedding, distance: "euclidean").first(5)

cube

Distance

Supported values are:

euclidean
cosine
taxicab
chebyshev

For cosine distance with cube, vectors must be normalized before being stored.

ruby

class Item  "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

ruby

input = [
  "The dog is barking",
  "The cat is purring",
  "The bear is growling"
]
embeddings = embed(input)

Store the embeddings

ruby

documents = []
input.zip(embeddings) do |content, embedding|
  documents  "Bearer #{ENV.fetch("CO_API_KEY")}",
    "Content-Type" => "application/json"
  }
  data = {
    texts: input,
    model: "embed-v4.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

ruby

input = [
  "The dog is barking",
  "The cat is purring",
  "The bear is growling"
]
embeddings = embed(input, "search_document")

Store the embeddings

ruby

documents = []
input.zip(embeddings) do |content, embedding|
  documents (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

Document.create!(content: "The dog is barking")
Document.create!(content: "The cat is purring")
Document.create!(content: "The bear is growling")

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

Document.find_each do |document|
  embedding = embed.(document.content, **embed_options)
  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)

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).first(5)

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)).first(5).map { |v| results[v[:doc_id]] }

See the complete code

Sparse Search

You can generate sparse embeddings locally with Transformers.rb.

Generate a model

bash

rails generate model Document content:text embedding:sparsevec{30522}
rails db:migrate

And add has_neighbors

ruby

class Document < ApplicationRecord
  has_neighbors :embedding
end

Load a model to generate embeddings

…

Pass your input

ruby

input = [
  "The dog is barking",
  "The cat is purring",
  "The bear is growling"
]
embeddings = model.embed(input)

Store the embeddings

ruby

documents = []
input.zip(embeddings) do |content, embedding|
  documents << {content: content, embedding: Neighbor::SparseVector.new(embedding)}
end
Document.insert_all!(documents)

Embed the search query

ruby

query = "forest"
query_embedding = model.embed([query])[0]

And search the documents

ruby

Document.nearest_neighbors(:embedding, Neighbor::SparseVector.new(query_embedding), distance: "inner_product").first(5).map(&:content)

See the complete code

Disco Recommendations

You can use Neighbor for online item-based recommendations with Disco. We’ll use MovieLens data for this example.

Generate a model

bash

rails generate model Movie name:string factors:cube
rails db:migrate

And add has_neighbors

ruby

class Movie < ApplicationRecord
  has_neighbors :factors, dimensions: 20, normalize: true
end

Fit the recommender

ruby

data = Disco.load_movielens
recommender = Disco::Recommender.new(factors: 20)
recommender.fit(data)

Store the item factors

ruby

movies = []
recommender.item_ids.each do |item_id|
  movies << {name: item_id, factors: recommender.item_factors(item_id)}
end
Movie.create!(movies)

And get similar movies

ruby

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 pgvector

History

View the changelog

Contributing

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:

bash

git clone https://github.com/ankane/neighbor.git
cd neighbor
bundle install

# Postgres
createdb neighbor_test
bundle exec rake test:postgresql

# SQLite
bundle exec rake test:sqlite
bundle exec rake test:sqlitevec
bundle exec rake test:vec1

# MariaDB
docker run -e MARIADB_ALLOW_EMPTY_ROOT_PASSWORD=1 -e MARIADB_DATABASE=neighbor_test -p 3307:3306 mariadb:12.3
bundle exec rake test:mariadb

# MySQL
docker run -e MYSQL_ALLOW_EMPTY_PASSWORD=1 -e MYSQL_DATABASE=neighbor_test -p 3306:3306 mysql:9.7
bundle exec rake test:mysql