How to Choose the Best Embedding Model for Your LLM Application

1
! pip install -qU datasets sentence-transformers numpy pandas tqdm

1
! pip install -qU voyageai

1
! pip install -qU openai

1
! pip install -qU transformers

1
import os
2
import getpass

1
import voyageai
2
VOYAGE_API_KEY = getpass.getpass("Voyage API Key:")
3
voyage_client = voyageai.Client(api_key=VOYAGE_API_KEY)

1
from openai import OpenAI
2
os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
3
openai_client = OpenAI()

1
from datasets import load_dataset
2
import pandas as pd
3
4
# Use streaming=True to load the dataset without downloading it fully
5
data = load_dataset("MongoDB/cosmopedia-wikihow-chunked", split="train", streaming=True)
6
# Get first 25k records from the dataset
7
data_head = data.take(25000)
8
df = pd.DataFrame(data_head)
9
10
# Use this if you want the full dataset
11
# data = load_dataset("MongoDB/cosmopedia-wikihow-chunked", split="train")
12
# df = pd.DataFrame(data)

1
# Ensuring length of dataset is what we expect i.e. 25k
2
len(df)
3
4
# Previewing the contents of the data
5
df.head()
6
7
# Only keep records where the text field is not null
8
df = df[df["text"].notna()]
9
10
# Number of unique documents in the dataset
11
df.doc_id.nunique()

1
def get_embeddings(docs: List[str], input_type: str, model:str="voyage-lite-02-instruct") -> List[List[float]]:
2
    """
3
    Get embeddings using the Voyage AI API.
4
5
    Args:
6
        docs (List[str]): List of texts to embed
7
        input_type (str): Type of input to embed. Can be "document" or "query".
8
        model (str, optional): Model name. Defaults to "voyage-lite-02-instruct".
9
10
    Returns:
11
        List[List[float]]: Array of embedddings
12
    """
13
    response = voyage_client.embed(docs, model=model, input_type=input_type)
14
    return response.embeddings

1
def get_embeddings(docs: List[str], model: str="text-embedding-3-large") -> List[List[float]]:
2
    """
3
    Generate embeddings using the OpenAI API.
4
5
    Args:
6
        docs (List[str]): List of texts to embed
7
        model (str, optional): Model name. Defaults to "text-embedding-3-large".
8
9
    Returns:
10
        List[float]: Array of embeddings
11
    """
12
    # replace newlines, which can negatively affect performance.
13
    docs = [doc.replace("\n", " ") for doc in docs]
14
    response = openai_client.embeddings.create(input=docs, model=model)
15
    response = [r.embedding for r in response.data]
16
    return response

1
{
2
  "data": [
3
    {
4
      "embedding": [
5
        0.018429679796099663,
6
        -0.009457024745643139
7
    .
8
    .
9
    .
10
      ],
11
      "index": 0,
12
      "object": "embedding"
13
    }
14
  ],
15
  "model": "text-embedding-3-large",
16
  "object": "list",
17
  "usage": {
18
    "prompt_tokens": 183,
19
    "total_tokens": 183
20
  }
21
}

1
from typing import List
2
from transformers import AutoModel, AutoTokenizer
3
import torch
4
5
# Instruction to append to user queries, to improve retrieval
6
RETRIEVAL_INSTRUCT = "Represent this sentence for searching relevant passages:"
7
8
# Check if CUDA (GPU support) is available, and set the device accordingly
9
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
10
# Load the UAE-Large-V1 model from the Hugging Face 
11
model = AutoModel.from_pretrained('WhereIsAI/UAE-Large-V1').to(device)
12
# Load the tokenizer associated with the UAE-Large-V1 model
13
tokenizer = AutoTokenizer.from_pretrained('WhereIsAI/UAE-Large-V1')
14
15
# Decorator to disable gradient calculations
16
@torch.no_grad()
17
def get_embeddings(docs: List[str], input_type: str) -> List[List[float]]:
18
    """
19
    Get embeddings using the UAE-Large-V1 model.
20
21
    Args:
22
        docs (List[str]): List of texts to embed
23
        input_type (str): Type of input to embed. Can be "document" or "query".
24
25
    Returns:
26
        List[List[float]]: Array of embedddings
27
    """
28
    # Prepend retrieval instruction to queries
29
    if input_type == "query":
30
        docs = ["{}{}".format(RETRIEVAL_INSTRUCT, q) for q in docs]
31
    # Tokenize input texts
32
    inputs = tokenizer(docs, padding=True, truncation=True, return_tensors='pt', max_length=512).to(device)
33
    # Pass tokenized inputs to the model, and obtain the last hidden state
34
    last_hidden_state = model(**inputs, return_dict=True).last_hidden_state
35
    # Extract embeddings from the last hidden state
36
    embeddings = last_hidden_state[:, 0]
37
    return embeddings.cpu().numpy()

1
from tqdm.auto import tqdm
2
3
# Get all the texts in the dataset
4
texts = df["text"].tolist()
5
6
# Number of samples in a single batch
7
batch_size = 128
8
9
embeddings = []
10
# Generate embeddings in batches
11
for i in tqdm(range(0, len(texts), batch_size)):
12
    end = min(len(texts), i+batch_size)
13
    batch = texts[i:end]
14
    # Generate embeddings for current batch
15
    batch_embeddings = get_embeddings(batch)
16
    # Add to the list of embeddings
17
    embeddings.extend(batch_embeddings)