如何为你的 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 数量的样本,获取该批次的嵌入结果,并将其添加到我们的 "向量存储 "中。 在我们的硬件上生成嵌入结果所需的时间如下:

ModelBatch SizeDimensionsTime
text-embedding-3-large12830724m 17s
voyage-lite-02-instruct128102411m 14s
UAE-large-V1128102419m 50s

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

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

  • 19
    点赞
  • 7
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值