Improved examples [skip ci]

Author: ankane

examples/cohere/example.rb

@@ -9,7 +9,7 @@
 conn.exec("CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1024))")
 
 # https://docs.cohere.com/reference/embed
-def fetch_embeddings(texts, input_type)
+def embed(texts, input_type)
   url = "https://api.cohere.com/v1/embed"
   headers = {
     "Authorization" => "Bearer #{ENV.fetch("CO_API_KEY")}",
@@ -31,13 +31,13 @@ def fetch_embeddings(texts, input_type)
   "The cat is purring",
   "The bear is growling"
 ]
-embeddings = fetch_embeddings(input, "search_document")
+embeddings = embed(input, "search_document")
 input.zip(embeddings) do |content, embedding|
   conn.exec_params("INSERT INTO documents (content, embedding) VALUES ($1, $2)", [content, embedding])
 end
 
 query = "forest"
-query_embedding = fetch_embeddings([query], "search_query")[0]
+query_embedding = embed([query], "search_query")[0]
 result = conn.exec_params("SELECT content FROM documents ORDER BY embedding <~> $1 LIMIT 5", [query_embedding])
 result.each do |row|
   puts row["content"]

examples/openai/example.rb

@@ -11,7 +11,7 @@
 
 # https://platform.openai.com/docs/guides/embeddings/how-to-get-embeddings
 # input can be an array with 2048 elements
-def fetch_embeddings(input)
+def embed(input)
   url = "https://api.openai.com/v1/embeddings"
   headers = {
     "Authorization" => "Bearer #{ENV.fetch("OPENAI_API_KEY")}",
@@ -31,14 +31,14 @@ def fetch_embeddings(input)
   "The cat is purring",
   "The bear is growling"
 ]
-embeddings = fetch_embeddings(input)
-
+embeddings = embed(input)
 input.zip(embeddings) do |content, embedding|
   conn.exec_params("INSERT INTO documents (content, embedding) VALUES ($1, $2)", [content, embedding])
 end
 
-document_id = 1
-result = conn.exec_params("SELECT content FROM documents WHERE id != $1 ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = $1) LIMIT 5", [document_id])
+query = "forest"
+query_embedding = embed([query])[0]
+result = conn.exec_params("SELECT content FROM documents ORDER BY embedding <=> $1 LIMIT 5", [query_embedding])
 result.each do |row|
   puts row["content"]
 end

examples/openai/example_sequel.rb

@@ -20,7 +20,7 @@ class Document < Sequel::Model
 
 # https://platform.openai.com/docs/guides/embeddings/how-to-get-embeddings
 # input can be an array with 2048 elements
-def fetch_embeddings(input)
+def embed(input)
   url = "https://api.openai.com/v1/embeddings"
   headers = {
     "Authorization" => "Bearer #{ENV.fetch("OPENAI_API_KEY")}",
@@ -40,13 +40,14 @@ def fetch_embeddings(input)
   "The cat is purring",
   "The bear is growling"
 ]
-embeddings = fetch_embeddings(input)
+embeddings = embed(input)
 
 documents = []
 input.zip(embeddings) do |content, embedding|
   documents << {content: content, embedding: Pgvector.encode(embedding)}
 end
 Document.multi_insert(documents)
 
-document = Document.first
-pp document.nearest_neighbors(:embedding, distance: "cosine").limit(5).map(&:content)
+query = "forest"
+query_embedding = embed([query])[0]
+pp Document.nearest_neighbors(:embedding, query_embedding, distance: "cosine").limit(5).map(&:content)