-
Notifications
You must be signed in to change notification settings - Fork 14
/
example.rb
37 lines (30 loc) · 917 Bytes
/
example.rb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
require "bundler/setup"
require "active_record"
require "informers"
require "neighbor"
ActiveRecord::Base.establish_connection adapter: "postgresql", database: "neighbor_test"
ActiveRecord::Schema.verbose = false
ActiveRecord::Schema.define do
enable_extension "vector"
create_table :documents, force: true do |t|
t.text :content
t.vector :embedding, limit: 384
end
end
class Document < ActiveRecord::Base
has_neighbors :embedding
end
model = Informers.pipeline("embedding", "sentence-transformers/all-MiniLM-L6-v2")
input = [
"The dog is barking",
"The cat is purring",
"The bear is growling"
]
embeddings = model.(input)
documents = []
input.zip(embeddings) do |content, embedding|
documents << {content: content, embedding: embedding}
end
Document.insert_all!(documents)
document = Document.first
pp document.nearest_neighbors(:embedding, distance: "cosine").first(5).map(&:content)