Improved examples [skip ci]

Author: ankane

examples/cohere/example.py

@@ -12,7 +12,7 @@
 conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1024))')
 
 
-def fetch_embeddings(input, input_type):
+def embed(input, input_type):
     co = cohere.Client()
     response = co.embed(texts=input, model='embed-english-v3.0', input_type=input_type, embedding_types=['ubinary'])
     return [np.unpackbits(np.array(embedding, dtype=np.uint8)) for embedding in response.embeddings.ubinary]
@@ -23,12 +23,12 @@ def fetch_embeddings(input, input_type):
     'The cat is purring',
     'The bear is growling'
 ]
-embeddings = fetch_embeddings(input, 'search_document')
+embeddings = embed(input, 'search_document')
 for content, embedding in zip(input, embeddings):
     conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, Bit(embedding)))
 
 query = 'forest'
-query_embedding = fetch_embeddings([query], 'search_query')[0]
+query_embedding = embed([query], 'search_query')[0]
 result = conn.execute('SELECT content FROM documents ORDER BY embedding <~> %s LIMIT 5', (Bit(query_embedding),)).fetchall()
 for row in result:
     print(row[0])

examples/openai/example.py

@@ -1,3 +1,4 @@
+import numpy as np
 from openai import OpenAI
 from pgvector.psycopg import register_vector
 import psycopg
@@ -10,20 +11,24 @@
 conn.execute('DROP TABLE IF EXISTS documents')
 conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(1536))')
 
+
+def embed(input):
+    client = OpenAI()
+    response = client.embeddings.create(input=input, model='text-embedding-3-small')
+    return [v.embedding for v in response.data]
+
+
 input = [
     'The dog is barking',
     'The cat is purring',
     'The bear is growling'
 ]
-
-client = OpenAI()
-response = client.embeddings.create(input=input, model='text-embedding-3-small')
-embeddings = [v.embedding for v in response.data]
-
+embeddings = embed(input)
 for content, embedding in zip(input, embeddings):
-    conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding))
+    conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, np.array(embedding)))
 
-document_id = 1
-neighbors = conn.execute('SELECT content FROM documents WHERE id != %(id)s ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = %(id)s) LIMIT 5', {'id': document_id}).fetchall()
-for neighbor in neighbors:
-    print(neighbor[0])
+query = 'forest'
+query_embedding = embed([query])[0]
+result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (np.array(query_embedding),)).fetchall()
+for row in result:
+    print(row[0])

examples/sentence_transformers/example.py

@@ -10,19 +10,19 @@
 conn.execute('DROP TABLE IF EXISTS documents')
 conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(384))')
 
+model = SentenceTransformer('all-MiniLM-L6-v2')
+
 input = [
     'The dog is barking',
     'The cat is purring',
     'The bear is growling'
 ]
-
-model = SentenceTransformer('all-MiniLM-L6-v2')
 embeddings = model.encode(input)
-
 for content, embedding in zip(input, embeddings):
     conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding))
 
-document_id = 1
-neighbors = conn.execute('SELECT content FROM documents WHERE id != %(id)s ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = %(id)s) LIMIT 5', {'id': document_id}).fetchall()
-for neighbor in neighbors:
-    print(neighbor[0])
+query = 'forest'
+query_embedding = model.encode(query)
+result = conn.execute('SELECT content FROM documents ORDER BY embedding <=> %s LIMIT 5', (query_embedding,)).fetchall()
+for row in result:
+    print(row[0])

examples/sparse_search/example.py

@@ -20,7 +20,7 @@
 special_token_ids = [tokenizer.vocab[token] for token in tokenizer.special_tokens_map.values()]
 
 
-def fetch_embeddings(input):
+def embed(input):
     feature = tokenizer(
         input,
         padding=True,
@@ -42,12 +42,12 @@ def fetch_embeddings(input):
     'The cat is purring',
     'The bear is growling'
 ]
-embeddings = fetch_embeddings(input)
+embeddings = embed(input)
 for content, embedding in zip(input, embeddings):
     conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, SparseVector(embedding)))
 
 query = 'forest'
-query_embedding = fetch_embeddings([query])[0]
+query_embedding = embed([query])[0]
 result = conn.execute('SELECT content FROM documents ORDER BY embedding <#> %s LIMIT 5', (SparseVector(query_embedding),)).fetchall()
 for row in result:
     print(row[0])