Choosing the Best Embedding Model for RAG in Python

Introduction to Embedding Models for RAG

Embedding models are crucial for Retrieval-Augmented Generation (RAG) systems, as they convert text into dense vector representations. These vectors capture semantic meaning, enabling efficient similarity searches. Choosing the right embedding model can significantly impact your RAG system's performance. This presentation will guide you through the process of selecting and implementing the best embedding model for your RAG application using Python.

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

# Example of how embeddings work
text1 = "The quick brown fox jumps over the lazy dog"
text2 = "A fast auburn canine leaps above an idle hound"

# Simulate embeddings (in practice, these would be generated by a model)
embedding1 = np.random.rand(1, 300)
embedding2 = np.random.rand(1, 300)

# Calculate similarity
similarity = cosine_similarity(embedding1, embedding2)[0][0]
print(f"Similarity between texts: {similarity:.4f}")

Understanding Embedding Dimensions

Embedding dimensions play a crucial role in the performance and efficiency of your RAG system. Higher dimensions can capture more nuanced relationships but require more computational resources. Lower dimensions are faster but may lose some information. The optimal dimension often depends on your specific use case and dataset size.

import numpy as np
import matplotlib.pyplot as plt

# Generate embeddings with different dimensions
dimensions = [50, 100, 300, 1000]
num_samples = 1000

for dim in dimensions:
    embeddings = np.random.randn(num_samples, dim)
    
    # Calculate pairwise distances
    distances = np.linalg.norm(embeddings[:, np.newaxis] - embeddings, axis=2)
    
    plt.hist(distances.flatten(), bins=50, alpha=0.7, label=f'{dim}D')

plt.title('Distribution of Pairwise Distances')
plt.xlabel('Distance')
plt.ylabel('Frequency')
plt.legend()
plt.show()

Evaluating Embedding Quality

To choose the best embedding model, it's essential to evaluate its quality. One common method is to use benchmark datasets that test for semantic similarity. We'll demonstrate how to evaluate embeddings using the STS (Semantic Textual Similarity) benchmark.

from sentence_transformers import SentenceTransformer, util
import numpy as np

# Load a pre-trained model
model = SentenceTransformer('all-MiniLM-L6-v2')

# Example STS pairs
sts_pairs = [
    ("A man is playing a violin.", "A woman is playing a violin."),
    ("A dog is running in the park.", "A cat is sleeping on the couch."),
]

# Generate embeddings
embeddings = model.encode([pair[0] for pair in sts_pairs] + [pair[1] for pair in sts_pairs])

# Calculate cosine similarities
similarities = util.pytorch_cos_sim(embeddings[:len(sts_pairs)], embeddings[len(sts_pairs):])

for i, (sent1, sent2) in enumerate(sts_pairs):
    print(f"Sentence 1: {sent1}")
    print(f"Sentence 2: {sent2}")
    print(f"Similarity: {similarities[i][i]:.4f}\n")

Comparing Different Embedding Models

When selecting an embedding model for RAG, it's crucial to compare multiple options. We'll demonstrate how to compare three popular models: BERT, RoBERTa, and DistilBERT. Each model has its strengths and trade-offs in terms of performance and computational requirements.

from transformers import AutoTokenizer, AutoModel
import torch
import numpy as np

models = ['bert-base-uncased', 'roberta-base', 'distilbert-base-uncased']
sentences = [
    "The quick brown fox jumps over the lazy dog.",
    "I love machine learning and natural language processing."
]

def get_embeddings(model_name, sentences):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name)
    
    inputs = tokenizer(sentences, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
    
    return outputs.last_hidden_state.mean(dim=1).numpy()

for model_name in models:
    embeddings = get_embeddings(model_name, sentences)
    print(f"{model_name} embeddings shape: {embeddings.shape}")
    print(f"Sample embedding: {embeddings[0][:5]}\n")

Fine-tuning Embedding Models

Fine-tuning can significantly improve the performance of embedding models for specific domains or tasks. We'll demonstrate how to fine-tune a BERT model for a custom similarity task using PyTorch and the Hugging Face Transformers library.

from transformers import BertTokenizer, BertModel
import torch
import torch.nn as nn
import torch.optim as optim

# Load pre-trained model and tokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

# Define a simple similarity model
class SimilarityModel(nn.Module):
    def __init__(self, bert_model):
        super(SimilarityModel, self).__init__()
        self.bert = bert_model
        self.cosine_sim = nn.CosineSimilarity(dim=1)
    
    def forward(self, input_ids1, attention_mask1, input_ids2, attention_mask2):
        output1 = self.bert(input_ids1, attention_mask=attention_mask1).last_hidden_state.mean(dim=1)
        output2 = self.bert(input_ids2, attention_mask=attention_mask2).last_hidden_state.mean(dim=1)
        return self.cosine_sim(output1, output2)

# Initialize the model and optimizer
similarity_model = SimilarityModel(model)
optimizer = optim.Adam(similarity_model.parameters(), lr=2e-5)

# Example training loop (you would need a proper dataset and more epochs)
for epoch in range(3):
    # ... training code here ...
    pass

print("Fine-tuning complete")

Optimizing Embedding Storage and Retrieval

Efficient storage and retrieval of embeddings are crucial for RAG systems. We'll explore using FAISS, a library for efficient similarity search and clustering of dense vectors, to optimize embedding storage and retrieval.

import numpy as np
import faiss

# Generate sample embeddings
num_vectors = 10000
dimension = 128
embeddings = np.random.random((num_vectors, dimension)).astype('float32')

# Create a FAISS index
index = faiss.IndexFlatL2(dimension)
index.add(embeddings)

# Perform a similarity search
k = 5  # Number of nearest neighbors to retrieve
query = np.random.random((1, dimension)).astype('float32')
distances, indices = index.search(query, k)

print(f"Indices of {k} nearest neighbors: {indices[0]}")
print(f"Distances to {k} nearest neighbors: {distances[0]}")

Handling Out-of-Vocabulary Words

Embedding models often struggle with out-of-vocabulary (OOV) words. We'll explore techniques to handle OOV words, including subword tokenization and using fastText for character n-gram embeddings.

from gensim.models import FastText
from gensim.models.fasttext import load_facebook_vectors

# Train a FastText model (for demonstration, we'll use a small corpus)
sentences = [['hello', 'world'], ['machine', 'learning'], ['natural', 'language', 'processing']]
model = FastText(sentences, min_count=1, epochs=10, sg=1)

# Get embeddings for words, including OOV words
words = ['hello', 'world', 'machine', 'learning', 'unknownword']

for word in words:
    embedding = model.wv[word]
    print(f"Embedding for '{word}': {embedding[:5]}...")  # Showing first 5 dimensions

# Load pre-trained FastText model (comment out if you don't have the model file)
# fasttext_model = load_facebook_vectors('path_to_fasttext_model.bin')

Multilingual Embedding Models

For RAG systems that need to handle multiple languages, multilingual embedding models are essential. We'll explore how to use a multilingual model like XLM-RoBERTa to generate embeddings for text in different languages.

from transformers import XLMRobertaTokenizer, XLMRobertaModel
import torch

# Load pre-trained model and tokenizer
tokenizer = XLMRobertaTokenizer.from_pretrained('xlm-roberta-base')
model = XLMRobertaModel.from_pretrained('xlm-roberta-base')

# Example sentences in different languages
sentences = [
    "Hello, how are you?",  # English
    "Bonjour, comment allez-vous?",  # French
    "こんにちは、お元気ですか？",  # Japanese
]

# Generate embeddings
inputs = tokenizer(sentences, return_tensors="pt", padding=True, truncation=True)
with torch.no_grad():
    outputs = model(**inputs)

embeddings = outputs.last_hidden_state.mean(dim=1)

for i, sentence in enumerate(sentences):
    print(f"Embedding for '{sentence}': {embeddings[i][:5]}...")  # Showing first 5 dimensions

Real-Life Example: Document Similarity

Let's explore a real-life example of using embeddings for document similarity in a RAG system. We'll use a pre-trained model to embed scientific paper abstracts and find similar documents.

from sentence_transformers import SentenceTransformer, util

# Load a pre-trained model
model = SentenceTransformer('allenai-specter')

# Example scientific paper abstracts
abstracts = [
    "We present a new method for natural language processing using deep learning techniques.",
    "This paper introduces an innovative approach to computer vision using convolutional neural networks.",
    "Our research focuses on the applications of machine learning in healthcare.",
]

# Generate embeddings
embeddings = model.encode(abstracts)

# Calculate cosine similarities
similarities = util.pytorch_cos_sim(embeddings, embeddings)

print("Similarity matrix:")
print(similarities)

# Find the most similar pair
max_sim = 0
max_pair = (0, 0)
for i in range(len(abstracts)):
    for j in range(i+1, len(abstracts)):
        if similarities[i][j] > max_sim:
            max_sim = similarities[i][j]
            max_pair = (i, j)

print(f"\nMost similar abstracts:")
print(f"1: {abstracts[max_pair[0]]}")
print(f"2: {abstracts[max_pair[1]]}")
print(f"Similarity: {max_sim:.4f}")

Real-Life Example: Question Answering

Another practical application of embeddings in RAG systems is for question answering. We'll demonstrate how to use embeddings to match questions with potential answers from a knowledge base.

from sentence_transformers import SentenceTransformer, util

# Load a pre-trained model
model = SentenceTransformer('all-MiniLM-L6-v2')

# Knowledge base (questions and answers)
kb = [
    ("What is the capital of France?", "The capital of France is Paris."),
    ("Who wrote 'To Kill a Mockingbird'?", "Harper Lee wrote 'To Kill a Mockingbird'."),
    ("What is the boiling point of water?", "The boiling point of water is 100 degrees Celsius at sea level."),
]

# User question
user_question = "What's the main city in France?"

# Encode questions and user query
kb_questions = [item[0] for item in kb]
question_embeddings = model.encode(kb_questions)
query_embedding = model.encode(user_question)

# Find most similar question
similarities = util.pytorch_cos_sim(query_embedding, question_embeddings)[0]
best_match = similarities.argmax().item()

print(f"User question: {user_question}")
print(f"Best matching question: {kb[best_match][0]}")
print(f"Answer: {kb[best_match][1]}")
print(f"Similarity score: {similarities[best_match]:.4f}")

Evaluating Embedding Performance in RAG

To ensure your chosen embedding model performs well in your RAG system, it's important to evaluate its performance in the context of your specific task. We'll demonstrate how to set up a simple evaluation pipeline for a question-answering RAG system.

from sentence_transformers import SentenceTransformer, util
import numpy as np

# Load a pre-trained model
model = SentenceTransformer('all-MiniLM-L6-v2')

# Sample dataset (question, context, answer)
dataset = [
    ("Who is the president of the United States?", "Joe Biden is the current president of the United States.", "Joe Biden"),
    ("What is the capital of Japan?", "Tokyo is the capital and largest city of Japan.", "Tokyo"),
    ("Who wrote 'Pride and Prejudice'?", "Jane Austen wrote the novel 'Pride and Prejudice' in 1813.", "Jane Austen"),
]

def evaluate_rag(model, dataset):
    correct = 0
    for question, context, answer in dataset:
        # Encode question and context
        question_embedding = model.encode(question)
        context_embedding = model.encode(context)
        
        # Calculate similarity
        similarity = util.pytorch_cos_sim(question_embedding, context_embedding)[0][0]
        
        # Simple retrieval (in a real system, you'd have multiple contexts)
        if similarity > 0.5:  # Arbitrary threshold
            # Check if answer is in context (simple string matching)
            if answer.lower() in context.lower():
                correct += 1
    
    accuracy = correct / len(dataset)
    return accuracy

accuracy = evaluate_rag(model, dataset)
print(f"RAG System Accuracy: {accuracy:.2f}")

Optimizing Embedding Model Size

For deployment in resource-constrained environments, it's often necessary to optimize the size of your embedding model. We'll explore techniques like quantization and pruning to reduce model size while maintaining performance.

import torch
from transformers import AutoTokenizer, AutoModel

# Load a pre-trained model
model_name = 'distilbert-base-uncased'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# Original model size
orig_size = sum(p.numel() for p in model.parameters())
print(f"Original model size: {orig_size} parameters")

# Quantize the model
quantized_model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)

# Quantized model size
quant_size = sum(p.numel() for p in quantized_model.parameters())
print(f"Quantized model size: {quant_size} parameters")
print(f"Size reduction: {(1 - quant_size/orig_size)*100:.2f}%")

# Example usage of quantized model
input_text = "Hello, world!"
inputs = tokenizer(input_text, return_tensors="pt")
with torch.no_grad():
    outputs = quantized_model(**inputs)

print(f"Output shape: {outputs.last_hidden_state.shape}")

Continuous Learning for Embedding Models

To keep your RAG system up-to-date with new information, it's important to implement continuous learning for your embedding models. We'll explore a simple approach to update your model with new data over time.

from transformers import AutoTokenizer, AutoModel, AdamW
import torch

# Load pre-trained model and tokenizer
model_name = 'distilbert-base-uncased'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# Function to update model with new data
def update_model(model, new_data, learning_rate=1e-5, epochs=3):
    optimizer = AdamW(model.parameters(), lr=learning_rate)
    model.train()
    
    for epoch in range(epochs):
        for text in new_data:
            inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True)
            outputs = model(**inputs)
            loss = outputs.loss
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
    
    return model

# Example usage
new_data = [
    "Artificial intelligence is revolutionizing various industries.",
    "Machine learning models can process vast amounts of data quickly.",
    "Natural language processing enables computers to understand human language."
]

updated_model = update_model(model, new_data)
print("Model updated with new data")

Embedding Visualization for Model Inspection

Visualizing embeddings can provide insights into how your model represents different concepts. We'll use t-SNE to reduce the dimensionality of our embeddings and plot them in 2D space.

from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
from transformers import AutoTokenizer, AutoModel
import torch

model_name = 'distilbert-base-uncased'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

words = ["cat", "dog", "fish", "bird", "lion", "tiger", "elephant", "mouse"]

# Generate embeddings
inputs = tokenizer(words, padding=True, truncation=True, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)
embeddings = outputs.last_hidden_state[:, 0, :].numpy()

# Reduce dimensionality with t-SNE
tsne = TSNE(n_components=2, random_state=42)
reduced_embeddings = tsne.fit_transform(embeddings)

# Plot the embeddings
plt.figure(figsize=(10, 8))
for i, word in enumerate(words):
    x, y = reduced_embeddings[i]
    plt.scatter(x, y)
    plt.annotate(word, (x, y))

plt.title("t-SNE visualization of word embeddings")
plt.xlabel("Dimension 1")
plt.ylabel("Dimension 2")
plt.show()

Additional Resources

For further exploration of embedding models and their applications in RAG systems, consider the following resources:

1. "Efficient Estimation of Word Representations in Vector Space" by Mikolov et al. (2013) ArXiv: https://arxiv.org/abs/1301.3781
2. "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding" by Devlin et al. (2018) ArXiv: https://arxiv.org/abs/1810.04805
3. "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks" by Reimers and Gurevych (2019) ArXiv: https://arxiv.org/abs/1908.10084
4. "RoBERTa: A Robustly Optimized BERT Pretraining Approach" by Liu et al. (2019) ArXiv: https://arxiv.org/abs/1907.11692

These papers provide in-depth information on various embedding techniques and their applications in natural language processing tasks.