Comparing Vector Embedding Models in Python

Introduction

In this lesson, we will delve into the world of vector embeddings and how they can be used to compare different models. Vector embeddings are numerical representations of text that capture semantic meaning, allowing us to perform tasks like similarity comparison. We will focus on understanding cosine similarity, a crucial metric for comparing the similarity between vector embeddings. By the end of this lesson, you will understand how to calculate cosine similarity and its importance in evaluating the effectiveness of embedding models.

Understanding Cosine Similarity

\mathbf{A}

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Generating OpenAI Embeddings and Calculating Cosine Similarity

To demonstrate the use of cosine similarity, we will generate embeddings for three sentences using the OpenAI pre-trained model and calculate their cosine similarity. This will help us understand how similar or different the sentences are based on their embeddings.

Python

from openai import OpenAI
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# Initialize OpenAI client
client = OpenAI()

def get_openai_embedding(text: str) -> np.ndarray:
    """Fetch a 1536-dimensional embedding for 'text' using OpenAI."""
    response = client.embeddings.create(
        model="text-embedding-ada-002",
        input=text
    )
    # Convert to a NumPy array for convenience
    return np.array(response.data[0].embedding)

# Three test sentences
anchor_text = "I love pizza."
similar_text = "I enjoy pizza a lot."
different_text = "Penguins are cute animals."

# OpenAI embeddings
anchor_oa = get_openai_embedding(anchor_text)
similar_oa = get_openai_embedding(similar_text)
different_oa = get_openai_embedding(different_text)

sim_oa_similar = cosine_similarity([anchor_oa], [similar_oa])[0][0]
sim_oa_different = cosine_similarity([anchor_oa], [different_oa])[0][0]

print("OpenAI embeddings (text-embedding-ada-002):")
print(f"  Anchor vs. Similar:   {sim_oa_similar:.4f}")
print(f"  Anchor vs. Different: {sim_oa_different:.4f}")

from openai import OpenAI
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# Initialize OpenAI client
client = OpenAI()

def get_openai_embedding(text: str) -> np.ndarray:
    """Fetch a 1536-dimensional embedding for 'text' using OpenAI."""
    response = client.embeddings.create(
        model="text-embedding-ada-002",
        input=text
    )
    # Convert to a NumPy array for convenience
    return np.array(response.data[0].embedding)

# Three test sentences
anchor_text = "I love pizza."
similar_text = "I enjoy pizza a lot."
different_text = "Penguins are cute animals."

# OpenAI embeddings
anchor_oa = get_openai_embedding(anchor_text)
similar_oa = get_openai_embedding(similar_text)
different_oa = get_openai_embedding(different_text)

sim_oa_similar = cosine_similarity([anchor_oa], [similar_oa])[0][0]
sim_oa_different = cosine_similarity([anchor_oa], [different_oa])[0][0]

print("OpenAI embeddings (text-embedding-ada-002):")
print(f"  Anchor vs. Similar:   {sim_oa_similar:.4f}")
print(f"  Anchor vs. Different: {sim_oa_different:.4f}")

Output:

text

OpenAI embeddings (text-embedding-ada-002):
  Anchor vs. Similar:   0.9603
  Anchor vs. Different: 0.7973

OpenAI embeddings (text-embedding-ada-002):
  Anchor vs. Similar:   0.9603
  Anchor vs. Different: 0.7973

We use cosine_similarity from sklearn.metrics.pairwise to calculate the similarity scores between the embeddings. The [0][0] indexing is used because cosine_similarity returns a 2D array (a matrix) even when comparing two single vectors. The [0][0] accesses the first element of this matrix, which contains the similarity score between the two vectors.

The output shows that for the OpenAI model, the "Anchor vs. Similar" pair has a higher cosine similarity score compared to the "Anchor vs. Different" pair, indicating that the sentences "I love pizza." and "I enjoy pizza a lot." are more semantically similar than "I love pizza." and "Penguins are cute animals."

Generating Hugging Face Embeddings and Calculating Cosine Similarity

Next, we will use the Hugging Face pre-trained model to generate embeddings for the same three sentences and calculate their cosine similarity. This will allow us to compare the performance of the two models in capturing semantic similarities and differences.

Python

from sentence_transformers import SentenceTransformer

# Hugging Face embeddings
model = SentenceTransformer('all-MiniLM-L6-v2')
# Encode all three sentences in one call
hf_embeds = model.encode([anchor_text, similar_text, different_text])
anchor_hf, similar_hf, different_hf = hf_embeds

sim_hf_similar = cosine_similarity([anchor_hf], [similar_hf])[0][0]
sim_hf_different = cosine_similarity([anchor_hf], [different_hf])[0][0]

print("Hugging Face embeddings (all-MiniLM-L6-v2):")
print(f"  Anchor vs. Similar:   {sim_hf_similar:.4f}")
print(f"  Anchor vs. Different: {sim_hf_different:.4f}")

from sentence_transformers import SentenceTransformer

# Hugging Face embeddings
model = SentenceTransformer('all-MiniLM-L6-v2')
# Encode all three sentences in one call
hf_embeds = model.encode([anchor_text, similar_text, different_text])
anchor_hf, similar_hf, different_hf = hf_embeds

sim_hf_similar = cosine_similarity([anchor_hf], [similar_hf])[0][0]
sim_hf_different = cosine_similarity([anchor_hf], [different_hf])[0][0]

print("Hugging Face embeddings (all-MiniLM-L6-v2):")
print(f"  Anchor vs. Similar:   {sim_hf_similar:.4f}")
print(f"  Anchor vs. Different: {sim_hf_different:.4f}")

Output:

text

Hugging Face embeddings (all-MiniLM-L6-v2):
  Anchor vs. Similar:   0.8996
  Anchor vs. Different: 0.1977

Hugging Face embeddings (all-MiniLM-L6-v2):
  Anchor vs. Similar:   0.8996
  Anchor vs. Different: 0.1977

For the Hugging Face model, the output also shows that the "Anchor vs. Similar" pair has a higher cosine similarity score compared to the "Anchor vs. Different" pair. It's important to note that the dimensions of the embeddings need to be the same to perform cosine similarity calculations. This is why we cannot directly compare the embeddings from the two models for a given text, as the OpenAI model produces a 1536-dimensional output, while the Hugging Face model produces a 384-dimensional output.

Comparing Vector Embedding Models in Python

Introduction

Understanding Cosine Similarity

Join the 1M+ learners on CodeSignal

Generating OpenAI Embeddings and Calculating Cosine Similarity

Generating Hugging Face Embeddings and Calculating Cosine Similarity

Conclusion