Added Cohere example [skip ci]

ankane · ankane · commit 182fe8491326 · 2024-07-15T17:04:13.000-07:00
diff --git a/README.md b/README.md
@@ -27,6 +27,7 @@ And follow the instructions for your database library:
 Or check out some examples:
 
 - [Embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/openai_embeddings.py) with OpenAI
+- [Binary embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/cohere_embeddings.py) with Cohere
 - [Sentence embeddings](https://github.com/pgvector/pgvector-python/blob/master/examples/sentence_embeddings.py) with SentenceTransformers
 - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search_rrf.py) with SentenceTransformers (Reciprocal Rank Fusion)
 - [Hybrid search](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search.py) with SentenceTransformers (cross-encoder)
diff --git a/examples/cohere_embeddings.py b/examples/cohere_embeddings.py
@@ -0,0 +1,34 @@
+import cohere
+import numpy as np
+from pgvector.psycopg import register_vector, Bit
+import psycopg
+
+conn = psycopg.connect(dbname='pgvector_example', autocommit=True)
+
+conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
+register_vector(conn)
+
+conn.execute('DROP TABLE IF EXISTS documents')
+conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1024))')
+
+
+def fetch_embeddings(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]
+
+
+input = [
+    'The dog is barking',
+    'The cat is purring',
+    'The bear is growling'
+]
+embeddings = fetch_embeddings(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]
+result = conn.execute('SELECT content FROM documents ORDER BY embedding <~> %s LIMIT 5', (Bit(query_embedding),)).fetchall()
+for row in result:
+    print(row[0])
diff --git a/examples/requirements.txt b/examples/requirements.txt
@@ -1,3 +1,4 @@
+cohere
 datasets
 gensim
 imagehash

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+cohere`
`1`	`2`	`datasets`
`2`	`3`	`gensim`
`3`	`4`	`imagehash`