Quick cohere embeddings demo.

compute_embeddings.py

import json
import os

import cohere
import numpy as np

names = ["Vet", "Police", "Engineer"]
descriptions = [
    "Vets are like doctors but for animals.",
    "Police protect and serve the community.",
    "Engineers build things.",
]


co = cohere.Client(os.environ["COHERE_API_KEY"])
response = co.embed(texts=descriptions)
embeddings = np.array(response.embeddings)

with open("embeddings.npy", "wb") as f:
    np.save(f, embeddings)

with open("descriptions.json", "w") as f:
    json.dump(descriptions, f)

with open("names.json", "w") as f:
    json.dump(names, f)

find_similar.py

import json
import os

import cohere
import numpy as np

question = "can my dog eat chocolate?"


def cosine_distance(A, B):
    return np.dot(A, B) / (np.linalg.norm(A) * np.linalg.norm(B))


with open("embeddings.npy", "rb") as f:
    embeddings = np.load(f)

with open("descriptions.json", "r") as f:
    descriptions = json.load(f)

with open("names.json", "r") as f:
    names = json.load(f)


co = cohere.Client(os.environ["COHERE_API_KEY"])
response = co.embed(texts=[question])
question_embedding = np.array(response.embeddings[0])

best_name = None
best_score = -1000
for name, embedding in zip(names, embeddings):
    score = cosine_distance(embedding, question_embedding)
    if score > best_score:
        best_name = name
        best_score = score

print(best_name)