如何为你的 LLM 应用选择最合适的 Embedding 模型

如果你正在构建 2024 年的生成式人工智能（GenAI）应用，你现在可能已经听过几次 "嵌入（embedding） "这个词了，而且每周都能看到新的嵌入模型上架。 那么，为什么会有这么多人突然关心起嵌入这个自 20 世纪 50 年代就存在的概念呢？ 如果嵌入如此重要，而且您必须使用它们，那么您该如何在众多嵌入模型中做出选择呢？ 本教程将涵盖以下内容：

• 什么是嵌入？
• 嵌入在 RAG 应用程序中的重要性
• 如何为您的 RAG 应用程序选择最佳嵌入模型
• 评估嵌入模型

什么是嵌入和嵌入模型？

嵌入是代表文本、图像、音频、视频等信息的数字（矢量）数组。 这些数字共同捕捉数据的语义和其他重要特征。 这样做的直接后果是，在向量空间中，语义相似的实体彼此靠近，而不相似的实体则相距较远。 为清晰起见，请看下图对高维向量空间的描述：

在自然语言处理（NLP）中，嵌入模型是一种算法，旨在学习和生成给定信息的嵌入。 在当今的人工智能应用中，通常使用大型语言模型（LLM）创建嵌入模型，这些模型在海量语料库中进行训练，并使用尖端算法学习数据中的复杂语义关系。

什么是 RAG（简述）？

顾名思义，RAG（Retrieval-augmented generation）旨在利用从知识库中检索的数据提高预训练 LLM 生成的质量。 RAG 的成功在于从知识库中检索出最相关的结果。 这就是嵌入的作用所在。 RAG 管道看起来像这样：

在上述管道中，我们可以看到基因人工智能应用中常用的检索方法–即语义搜索。 在这种技术中，嵌入模型用于创建用户查询和知识库信息的向量表示。 这样，在给定用户查询及其嵌入的情况下，我们就可以根据文档嵌入与查询嵌入的相似程度，从知识库中检索出最相关的源文档。 然后将检索到的文档、用户查询和任何用户提示作为上下文传递给 LLM，以生成对用户问题的回答。

为您的 RAG 应用程序选择最佳嵌入模型

如上所述，嵌入是 RAG 的核心。 但是，面对如此众多的嵌入模型，我们该如何选择最适合我们使用情况的模型呢？

Hugging Face 上的 MTEB Leaderboard（很不幸，网络问题无法直接看到） 是寻找最佳嵌入模型的一个好的开始。 它是专有和开源文本嵌入模型的最新列表，并附有关于每个嵌入模型在各种嵌入任务（如检索、摘要等）中表现的统计数据。

不过，我们也有国产化替代，加油，fighting！！！

基准是一个很好的起点，但请记住，这些结果都是自我报告的，而且是在数据集上进行的基准测试，可能无法准确代表您正在处理的数据。 此外，一些嵌入模型可能会将 MTEB 数据集包含在其训练数据中，因为这些数据集是公开的。 因此，即使您根据基准结果选择了嵌入模型，我们也建议您在自己的数据集上对其进行评估。 我们将在本教程的稍后部分了解如何进行评估，但首先，让我们仔细看看排行榜。

让我们来看看 "总体 "选项卡，因为它提供了每种嵌入模型的全面总结。 不过，请注意，我们按照检索平均值列对排行榜进行了排序。 这是因为 RAG 是一项检索任务，我们希望看到最好的检索嵌入模型排在最前面。 我们将忽略与其他任务相对应的栏目，重点关注以下栏目：

• 检索平均值（Retrieval Average）： 代表多个数据集的平均归一化折现累积增益（NDCG）@ 10。 NDCG 是衡量检索系统性能的常用指标。 NDCG 越高，表明嵌入模型在检索结果列表中相关项目的排名越靠前。
• 模型大小（Model Size）： 嵌入模型大小（GB）。 通过它可以了解运行模型所需的计算资源。 虽然检索性能与模型大小成正比，但需要注意的是，模型大小对延迟也有直接影响。 在生产设置中，延迟与性能的权衡变得尤为重要。
• 最大标记数（Max Tokens）： 可压缩到单个嵌入式内容中的标记数。 通常情况下，您不希望将超过一段文字（约 100 个标记符）放入一个嵌入模型中。 因此，即使嵌入模型的最大标记数为 512，也应该绰绰有余了。
• 嵌入尺寸（Embedding Dimensions）： 嵌入向量的长度。 较小的嵌入向量推理速度更快，存储效率更高，而较多的维度则可以捕捉数据中细微的细节和关系。 最终，我们希望在捕捉数据的复杂性和运行效率之间取得良好的平衡。

样例

每个嵌入模型所需的库略有不同，但常见的库如下：

• datasets： 用于访问 Hugging Face Hub 上可用数据集的 Python 库
• sentence-transformers： 处理文本和图像嵌入的框架
• numpy： 提供数组数学运算工具的 Python 库
• pandas： 用于数据分析、探索和操作的 Python 库
• tdqm： 显示循环进度表的 Python 模块

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

Voyage AI 的附加功能：voyageai： 与 OpenAI API 交互的 Python 库

!pip install -qU voyageai

OpenAI 的附加功能：openai： 与 OpenAI API 交互的 Python 库

!pip install -qU openai

此外，UAE: transformers： Python 库，提供与 Hugging Face 上提供的预训练模型交互的 API

!pip install -qU transformers

第 2 步：设置先决条件 OpenAI 和 Voyage AI 模型通过 API 提供。 因此，您需要获取 API 密钥，并将其提供给相应的客户端。

import os
import getpass

初始化 Voyage AI 客户端：

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

初始化 OpenAI 客户端：

from openai import OpenAI
os.environ[“OPENAI_API_KEY”] = getpass.getpass(“OpenAI API Key:”)
openai_client = OpenAI()

下载评估数据集

如前所述，我们将使用 MongoDB 的 cosmopedia-wikihow 分块数据集。 该数据集相当大（超过 100 万个文档）。 因此，我们将以流的方式抓取前 25k 条记录，而不是将整个数据集下载到磁盘。

from datasets import load_dataset
import pandas as pd

# Use streaming=True to load the dataset without downloading it fully
data = load_dataset("MongoDB/cosmopedia-wikihow-chunked", split="train", streaming=True)
# Get first 25k records from the dataset
data_head = data.take(25000)
df = pd.DataFrame(data_head)

# Use this if you want the full dataset
# data = load_dataset("MongoDB/cosmopedia-wikihow-chunked", split="train")
# df = pd.DataFrame(data)

数据分析

现在我们有了数据集，让我们进行一些简单的数据分析，并对数据进行一些正确性检查，以确保我们没有发现任何明显的错误：

# Ensuring length of dataset is what we expect i.e. 25k
len(df)

# Previewing the contents of the data
df.head()

# Only keep records where the text field is not null
df = df[df["text"].notna()]

# Number of unique documents in the dataset
df.doc_id.nunique()

创建嵌入函数

现在，让我们为每个嵌入模型创建嵌入函数。

对于 voyage-lite-02-instruct，我们需要

def get_embeddings(docs: List[str], input_type: str, model:str="voyage-lite-02-instruct") -> List[List[float]]:
    """
    Get embeddings using the Voyage AI API.

    Args:
        docs (List[str]): List of texts to embed
        input_type (str): Type of input to embed. Can be "document" or "query".
        model (str, optional): Model name. Defaults to "voyage-lite-02-instruct".

    Returns:
        List[List[float]]: Array of embedddings
    """
    response = voyage_client.embed(docs, model=model, input_type=input_type)
    return response.embeddings

上述嵌入函数将文本列表（文档）和输入类型作为参数，并返回一个嵌入列表。 输入类型可以是文档或查询，这取决于我们是嵌入文档列表还是用户查询。 Voyage 利用该值在输入前添加特殊提示，以提高检索质量。

对于 text-embedding-3-large，我们需要

def get_embeddings(docs: List[str], model: str="text-embedding-3-large") -> List[List[float]]:
    """
    Generate embeddings using the OpenAI API.

    Args:
        docs (List[str]): List of texts to embed
        model (str, optional): Model name. Defaults to "text-embedding-3-large".

    Returns:
        List[float]: Array of embeddings
    """
    # replace newlines, which can negatively affect performance.
    docs = [doc.replace("\n", " ") for doc in docs]
    response = openai_client.embeddings.create(input=docs, model=model)
    response = [r.embedding for r in response.data]
    return response

OpenAI 模型的嵌入函数与之前的函数类似，但有一些主要区别–没有输入_类型参数，API 返回一个嵌入对象列表，需要对该列表进行解析才能得到最终的嵌入列表。 应用程序接口的响应示例如下：

{
  "data": [
    {
      "embedding": [
        0.018429679796099663,
        -0.009457024745643139
    .
    .
    .
      ],
      "index": 0,
      "object": "embedding"
    }
  ],
  "model": "text-embedding-3-large",
  "object": "list",
  "usage": {
    "prompt_tokens": 183,
    "total_tokens": 183
  }
}

对于UAE-large-V1：

from typing import List
from transformers import AutoModel, AutoTokenizer
import torch

# Instruction to append to user queries, to improve retrieval
RETRIEVAL_INSTRUCT = "Represent this sentence for searching relevant passages:"

# Check if CUDA (GPU support) is available, and set the device accordingly
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
# Load the UAE-Large-V1 model from the Hugging Face 
model = AutoModel.from_pretrained('WhereIsAI/UAE-Large-V1').to(device)
# Load the tokenizer associated with the UAE-Large-V1 model
tokenizer = AutoTokenizer.from_pretrained('WhereIsAI/UAE-Large-V1')

# Decorator to disable gradient calculations
@torch.no_grad()
def get_embeddings(docs: List[str], input_type: str) -> List[List[float]]:
    """
    Get embeddings using the UAE-Large-V1 model.

    Args:
        docs (List[str]): List of texts to embed
        input_type (str): Type of input to embed. Can be "document" or "query".

    Returns:
        List[List[float]]: Array of embedddings
    """
    # Prepend retrieval instruction to queries
    if input_type == "query":
        docs = ["{}{}".format(RETRIEVAL_INSTRUCT, q) for q in docs]
    # Tokenize input texts
    inputs = tokenizer(docs, padding=True, truncation=True, return_tensors='pt', max_length=512).to(device)
    # Pass tokenized inputs to the model, and obtain the last hidden state
    last_hidden_state = model(**inputs, return_dict=True).last_hidden_state
    # Extract embeddings from the last hidden state
    embeddings = last_hidden_state[:, 0]
    return embeddings.cpu().numpy()

UAE-Large-V1 模型是 Hugging Face Model Hub 上的一个开源模型。 首先，我们需要从 Hugging Face 下载模型及其标记符。 我们使用 Auto 类（即 Transformers 库中的 AutoModel 和 AutoTokenizer）进行下载，它会自动推断出底层模型架构，在本例中就是 BERT。 接下来，我们使用 .to(device) 将模型加载到 GPU 上，因为我们有一个可用的 GPU。 UAE 模型的嵌入函数与 Voyage 模型非常相似，它将文本（文档）列表和输入类型作为参数，并返回一个嵌入列表。 首先对输入文本进行标记化处理，包括填充（针对短序列）和截断（针对长序列），以确保输入模型的长度一致–在本例中为 512，由 max_length 参数定义。 return_tensors 的 pt 值表示标记化的输出应该是 PyTorch 张量。

然后将标记化文本传递给模型进行推理，并提取最后一个隐藏层（last_hidden_state）。 这一层是模型对整个输入序列的最终学习表示。 然而，最终嵌入只从第一个标记中提取，在基于转换器的模型中，第一个标记通常是一个特殊标记（BERT 中的[CLS]）。 由于转换器中的自我关注机制，序列中每个标记的表示都会受到所有其他标记的影响，因此这个标记可以作为整个序列的集合表示。 最后，我们使用 .cpu() 将嵌入移回 CPU，并使用 .numpy() 将 PyTorch 张量转换为 numpy 数组。

评估

如前所述，我们将根据嵌入延迟和检索质量对模型进行评估。 测量嵌入延迟 为了测量嵌入延迟，我们将创建一个本地向量存储，这基本上是整个数据集的嵌入列表。 这里的延迟定义为为整个数据集创建嵌入所需的时间。

from tqdm.auto import tqdm

# Get all the texts in the dataset
texts = df["text"].tolist()

# Number of samples in a single batch
batch_size = 128

embeddings = []
# Generate embeddings in batches
for i in tqdm(range(0, len(texts), batch_size)):
    end = min(len(texts), i+batch_size)
    batch = texts[i:end]
    # Generate embeddings for current batch
    batch_embeddings = get_embeddings(batch)
    # Add to the list of embeddings
    embeddings.extend(batch_embeddings)

我们首先创建一个要嵌入的所有文本的列表，并设置批量大小。 voyage-lite-02-instruct 模型的批量大小限制为 128，因此为了保持一致，我们对所有模型都使用相同的大小。 我们迭代文本列表，在每次迭代中抓取 batch_size 数量的样本，获取该批次的嵌入结果，并将其添加到我们的 "向量存储 "中。 在我们的硬件上生成嵌入结果所需的时间如下：

OpenAI 模型的延迟最低。 但要注意的是，它的嵌入维数也是其他两种模型的三倍。 OpenAI 还按使用的代币收费，因此该模型的存储和推理成本会随着时间的推移而增加。 虽然 UAE 模型是所有模型中速度最慢的（尽管推理是在 GPU 上运行的），但由于它是开源的，因此还有量化、蒸馏等优化空间。

当然，评估是需要定量指标及其相关技术的，建议结合数据多多探索学习。