Skip to content

Instantly share code, notes, and snippets.

@ankane
Created February 5, 2024 04:54
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save ankane/d8ce946afe30b739a3ad3449d4a955c5 to your computer and use it in GitHub Desktop.
Save ankane/d8ce946afe30b739a3ad3449d4a955c5 to your computer and use it in GitHub Desktop.
require "bundler/inline"
gemfile do
source "https://rubygems.org"
gem "activerecord", require: "active_record"
gem "neighbor", github: "ankane/neighbor"
gem "pg"
end
ActiveRecord::Base.establish_connection adapter: "postgresql", database: "neighbor_example"
ActiveRecord::Base.logger = ActiveSupport::Logger.new(STDOUT)
ActiveRecord::Schema.define do
enable_extension "vector"
create_table :todos, force: :cascade do |t|
t.column :embedding, :vector, limit: 3
end
end
class Todo < ActiveRecord::Base
has_neighbors :embedding
default_scope { select(column_names - ["embedding"]) }
end
def generate_embedding
3.times.map { rand }
end
3.times do
Todo.create!(embedding: generate_embedding)
end
# loads embedding for nearest neighbor searches only
todo = Todo.except(:select).find(1)
pp todo.nearest_neighbors(:embedding, distance: "cosine").first(5)
# does not load embedding (Rails 7.1+)
todo = Todo.find(1)
pp Todo.with(neighbor_cte: Todo.reselect(:embedding).where(id: todo.id))
.where.not(id: todo.id)
.where.not(embedding: nil)
.order(Arel.sql("embedding <=> (SELECT embedding FROM neighbor_cte)"))
.first(5)
require "bundler/inline"
gemfile do
source "https://rubygems.org"
gem "activerecord", require: "active_record"
gem "neighbor"
gem "pg"
end
ActiveRecord::Base.establish_connection adapter: "postgresql", database: "neighbor_example"
ActiveRecord::Base.logger = ActiveSupport::Logger.new(STDOUT)
ActiveRecord::Schema.define do
enable_extension "vector"
create_table :todos, force: :cascade do |t|
end
create_table :todo_embeddings, force: :cascade do |t|
t.belongs_to :todo
t.column :embedding, :vector, limit: 3
end
end
class Todo < ActiveRecord::Base
has_one :todo_embedding
end
class TodoEmbedding < ActiveRecord::Base
belongs_to :todo
has_neighbors :embedding
end
def generate_embedding
3.times.map { rand }
end
3.times do
t = Todo.create!
t.create_todo_embedding!(embedding: generate_embedding)
end
todo = Todo.find(1)
# loads embedding for nearest neighbor searches only
pp todo
.todo_embedding
.nearest_neighbors(:embedding, distance: "cosine")
.reselect(:todo_id)
.includes(:todo)
.first(5)
.map(&:todo)
# does not load embedding (Rails 7.1+)
pp Todo.with(neighbor_cte: TodoEmbedding.select(:embedding).where(todo_id: todo.id))
.joins(:todo_embedding)
.where.not(id: todo.id)
.where.not(todo_embedding: {embedding: nil})
.order(Arel.sql("embedding <=> (SELECT embedding FROM neighbor_cte)"))
.first(5)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment