Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save juanmirod/6429d962b6ffb71f2bc9a3cdb27cbd71 to your computer and use it in GitHub Desktop.
Save juanmirod/6429d962b6ffb71f2bc9a3cdb27cbd71 to your computer and use it in GitHub Desktop.
I talk with JARVIS about how to do semantic search with OpenAI embeddings

You: do you know the openAI embeddings API?

JARVIS:* Yes, I am familiar with the OpenAI embeddings API. It allows you to generate high-quality vector representations of text. These embeddings can be used for a variety of natural language processing tasks, such as sentiment analysis, language translation, and text classification.

You: I would like to calculate the embeddings of some texts and store them in a file to do semantic searches on them, how could I do that?

JARVIS:* To calculate the embeddings of your texts and store them in a file, you can use the OpenAI API along with a programming language such as Python. Here's an example script that you can modify to suit your needs:

import openai
import json

# First, you need to set up your OpenAI API credentials
openai.api_key = "YOUR_API_KEY"

# Define your texts
texts = [
    "The quick brown fox jumps over the lazy dog",
    "The cat in the hat",
    "To be or not to be, that is the question"
]

# Calculate the embeddings for each text
embeddings = []
for text in texts:
    response = openai.Embedding.create(
        engine="text-davinci-002",
        input=text
    )
    embeddings.append(response["embedding"])

# Write the embeddings to a file
with open("embeddings.json", "w") as f:
    json.dump(embeddings, f)

In this example, you first set up your OpenAI API credentials. Then, you define your texts as a list of strings. Next, you loop through each text and use the OpenAI Embedding API to calculate the embedding for that text. Finally, you write the embeddings to a file in JSON format.

You can modify this script to use your own texts and file names. Once you have the embeddings stored in a file, you can use them to perform semantic searches using a search algorithm of your choice.

You: but you are not storing the texts alongside the embeddings, how I would recover the texts later?

JARVIS:* You're right, in the previous example I only stored the embeddings, not the original texts. If you want to store both the texts and the embeddings together, you can modify the script like this:

import openai
import json

# First, you need to set up your OpenAI API credentials
openai.api_key = "YOUR_API_KEY"

# Define your texts
texts = [
    "The quick brown fox jumps over the lazy dog",
    "The cat in the hat",
    "To be or not to be, that is the question"
]

# Calculate the embeddings for each text
embeddings = []
for text in texts:
    response = openai.Embedding.create(
        engine="text-davinci-002",
        input=text
    )
    embeddings.append({"text": text, "embedding": response["embedding"]})

# Write the embeddings to a file
with open("embeddings.json", "w") as f:
    json.dump(embeddings, f)

In this modified script, I use a list of dictionaries to store both the text and its corresponding embedding. The embeddings list now contains a dictionary for each text, with keys "text" and "embedding". The "text" key contains the original text, and the "embedding" key contains the embedding calculated by the OpenAI API.

By using this modified script, you can recover the original texts along with their embeddings when you load the file later.

You: how could I do the semantic search on that file using for example cosine distance?

JARVIS:* To perform a semantic search on the file containing the embeddings using cosine distance, you can use a programming language such as Python along with a library such as NumPy. Here's an example script that you can modify to suit your needs:

import json
import numpy as np

# Load the embeddings from the file
with open("embeddings.json", "r") as f:
    embeddings = json.load(f)

# Define the query text
query = "The quick brown fox"

# Calculate the embedding for the query text
response = openai.Embedding.create(
    engine="text-davinci-002",
    input=query
)
query_embedding = response["embedding"]

# Calculate the cosine similarity between the query embedding and each text embedding
similarities = []
for item in embeddings:
    text_embedding = np.array(item["embedding"])
    similarity = np.dot(text_embedding, query_embedding) / (np.linalg.norm(text_embedding) * np.linalg.norm(query_embedding))
    similarities.append({"text": item["text"], "similarity": similarity})

# Sort the list of similarities in descending order
similarities = sorted(similarities, key=lambda x: x["similarity"], reverse=True)

# Print the top 3 most similar texts
for i in range(3):
    print(similarities[i]["text"], similarities[i]["similarity"])

In this example, I first load the embeddings from the file into a list of dictionaries. Each dictionary contains both the original text and its corresponding embedding. Then, I define the query text and calculate its embedding using the OpenAI API.

Next, I loop through each text embedding in the list and calculate its cosine similarity with the query embedding using NumPy. I store the similarity score along with the original text in a list of dictionaries.

Finally, I sort the list of similarities in descending order and print the top 3 most similar texts along with their similarity scores.

You can modify this script to use your own file names and query text. You can also experiment with different similarity metrics and ranking algorithms to optimize the search results.

You: this has a O(N) complexity because you always are calculating the cosine similarity againts all the stored embeddings, is there a way to reduce that complexity on search? maybe sorting the registers before saving them or doing some kind of pruning of the registers to calculate the consine distance only on some of them?

JARVIS:* Yes, you're correct that the previous example has a time complexity of O(N) because it compares the query embedding with every stored embedding. To reduce the search time, you can use an approximate nearest neighbor (ANN) algorithm to search for the most similar embeddings.

One popular ANN algorithm is called locality-sensitive hashing (LSH). LSH is a technique that hashes high-dimensional vectors to a lower-dimensional space in a way that similar vectors are more likely to hash to the same bucket. This allows you to quickly search for similar vectors by only comparing vectors within the same bucket.

To use LSH for approximate nearest neighbor search, you can use a library such as Annoy, which implements LSH in Python. Here's an example script that uses Annoy to perform semantic search on the file containing the embeddings:

import json
from annoy import AnnoyIndex

# Load the embeddings from the file
with open("embeddings.json", "r") as f:
    embeddings = json.load(f)

# Define the query text
query = "The quick brown fox"

# Calculate the embedding for the query text
response = openai.Embedding.create(
    engine="text-davinci-002",
    input=query
)
query_embedding = response["embedding"]

# Build the Annoy index
index = AnnoyIndex(len(embeddings[0]["embedding"]), metric="angular")
for i, item in enumerate(embeddings):
    index.add_item(i, item["embedding"])
index.build(50)  # 50 trees

# Search for the most similar embeddings
similarities = index.get_nns_by_vector(query_embedding, 3, include_distances=True)

# Print the top 3 most similar texts
for i in range(3):
    print(embeddings[similarities[0][i]]["text"], similarities[1][i])

In this example, I first load the embeddings from the file into a list of dictionaries. Then, I define the query text and calculate its embedding using the OpenAI API.

Next, I build an Annoy index from the embeddings. The AnnoyIndex class allows you to add items to the index and build it with a specified number of trees. The metric parameter specifies the distance metric to use, which in this case is the cosine similarity.

Once the index is built, I use the get_nns_by_vector() method to search for the most similar embeddings to the query embedding. The get_nns_by_vector() method returns the indices of the most similar embeddings, along with their similarity scores.

Finally, I print the top 3 most similar texts along with their similarity scores.

You can modify this script to use your own file names and query text. You can also experiment with different parameters for the Annoy index to optimize the search results.

You: q

You: Make a summary of the conversation in 5 words or less.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment