在LLM中，如何选择嵌入模型？

Rank	Model	Model Size (Million Parameters)	Memory Usage (GB, fp32)	Embedding Dimensions	Max Tokens	Average (56 datasets)	Classification Average (12 datasets)	Clustering Average (11 datasets)	PairClassification Average (3 datasets)	Reranking Average (4 datasets)	Retrieval Average (15 datasets)	STS Average (10 datasets)	Summarization Average (1 datasets)
1	NV-Embed-v2	7851	29.25	4096	32768	72.31	90.37	58.46	88.67	60.65	62.65	84.31	30.7
2	bge-en-icl	7111	26.49	4096	32768	71.67	88.95	57.89	88.14	59.86	62.16	84.24	30.77
3	stella_en_1.5B_v5	1543	5.75	8192	131072	71.19	87.63	57.69	88.07	61.21	61.01	84.51	31.49
4	SFR-Embedding-2_R	7111	26.49	4096	32768	70.31	89.05	56.17	88.07	60.14	60.18	81.26	30.71
5	gte-Qwen2-7B-instruct	7613	28.36	3584	131072	70.24	86.58	56.92	85.79	61.42	60.25	83.04	31.35
6	stella_en_400M_v5	435	1.62	8192	8192	70.11	86.67	56.7	87.74	60.16	58.97	84.22	31.66
7	bge-multilingual-gemma2	9242	34.43	3584	8192	69.88	88.08	54.65	85.84	59.72	59.24	83.88	31.2
8	NV-Embed-v1	7851	29.25	4096	32768	69.32	87.35	52.8	86.91	60.54	59.36	82.84	31.2
9	voyage-large-2-instruct			1024	16000	68.23	81.49	53.35	89.25	60.09	58.28	84.31	30.84
10	Linq-Embed-Mistral	7111	26.49	4096	32768	68.17	80.2	51.42	88.35	60.29	60.19	84.97	30.98
11	SFR-Embedding-Mistral	7111	26.49	4096	32768	67.56	78.33	51.67	88.54	60.64	59	85.05	31.16

Rank	Model	Model Size (Million Parameters)	Memory Usage (GB, fp32)	Embedding Dimensions	Max Tokens	Average (56 datasets)	Classification Average (12 datasets)	Clustering Average (11 datasets)	PairClassification Average (3 datasets)	Reranking Average (4 datasets)	Retrieval Average (15 datasets)	STS Average (10 datasets)	Summarization Average (1 datasets)
12	Zeta-Alpha-E5-Mistral	7111	26.49	4096	32768	67.55	78.77	51.89	88.19	59.02	59.5	84.24	31.42
13	gte-Qwen1.5-7B-instruct	7099	26.45	4096	32768	67.35	79.6	55.83	87.41	60.13	56.24	82.42	31.46
14	gte-Qwen2-1.5B-instruct	1776	6.62	1536	131072	67.16	82.47	48.75	87.51	59.98	58.29	82.73	31.17
15	voyage-lite-02-instruct	1220	4.54	1024	4000	67.13	79.25	52.42	86.87	58.24	56.6	85.79	31.01
16	GritLM-7B	7240	26.97		4096	66.58	78.65	50.61	87.29	60.48	57.36	83.35	30.39
17	e5-mistral-7b-instruct	7111	26.49	4096	32768	66.4	77.37	50.26	88.42	60.21	56.87	84.62	31.53
18	text-embedding-004	1200	4.47	768	2048	66.31	81.17	47.48	87.61	58.9	55.7	85.07	32.63
19	TDTE					65.96	77.17	47.86	88.27	60.46	57.05	84.82	30.83
20	BinGSE-Meta-Llama-3-8B-Instruct					65.72	76.05	47.18

Total Datasets: 213
Total Languages: 113
Total Scores: 85285
Total Models: 447

其中许多模型（通常是开源模型）似乎已经在 MTEB 基准测试中进行了微调，从而产生了夸大的性能数字。尽管如此，一些开源模型（例如 E5）报告的性能是准确的。

MTEB 中的许多字段我们大多可以忽略。对于我们这些在现实世界中使用这些模型的人来说，最重要的字段是：

分数：我们应该关注的分数是 “average” 和 “retrieval average”。两者高度相关，因此专注于任何一个都有效。
序列长度告诉我们模型可以使用并压缩到单个嵌入中的数量标记。一般来说，我们不建议在单个 embedding 中塞入超过一个段落的权重 - 因此支持多达 512 个令牌的模型通常绰绰有余。
模型大小：模型的大小表示运行起来的难易程度。靠近 MTEB 顶部的所有型号都大小合理。最大的一个是 instructor-xl（需要 4.96GB 内存），我们可以轻松地在消费类硬件上运行它。

专注于这些列为我们提供了选择可能适合我们需求的模型所需的所有信息。考虑到这一点，我们选择了三个模型来在本文中展示 — 两个专有模型 Ada 002 和 embed-english-v3.0 — 以及一个微小但高性能的开源模型;e5-base-v2 的 v2 中。

三、下载测试数据并使用

要执行我们的比较，我们需要一个数据集。我们将使用来自 HuggingFace 数据集的预分块 AI ArXiv 数据集。

!pip install -qU datasets==2.14.6

from datasets import load_dataset

data = load_dataset(
    "jamescalam/ai-arxiv-chunked",
    split= "train"
)

这个数据集为我们提供了 ~42K 个文本块来嵌入，每个文本块大约一两个段落。

3.1 先决条件

每个模型的先决条件略有不同。OpenAI 和 Cohere 将这两个专有模型存储在 API 后面，因此它们的客户端库非常轻量级，我们只需要安装它们并获取它们各自的 API 密钥即可。

!pip install -qU \
    cohere==4.34 \
    openai==1.2.2

import os
import cohere
import openai

# initialize cohere
os.environ["COHERE_API_KEY"] = "your cohere api key"
co = cohere.Client()

# openai doesn't need to be initialized, but need to set api key
os.environ["OPENAI_API_KEY"] = "your openai api key"

E5 是一个本地模型，我们需要更多的代码和安装来运行它——包括像 PyTorch 这样相当繁重的库。该模型相对较轻，因此我们不需要繁重的 GPU 实例或大量的运行内存。然而，理想情况下，我们确实希望至少有一个 GPU 来运行它以获得更快的性能，但这并不是一个严格的要求，我们只能使用 CPU 进行管理。

!pip install -qU \
    torch==2.1.2 \
    transformers==4.25.0

import torch
from transformers import AutoModel, AutoTokenizer

# use GPU if available, on mac can use MPS
device = "cuda" if torch.cuda.is_available() else "cpu"

model_id = "intfloat/e5-base-v2"

# initialize tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id).to(device)
model.eval()

3.2 创建嵌入

为了创建我们的嵌入，我们将为每个模型创建一个嵌入函数。我们将字符串列表传递给这些嵌入函数，并期望返回 vector 嵌入列表。

3.2.1 API 嵌入

当然，我们专有的嵌入模型的代码更直接，所以让我们先介绍这两个：

# cohere embedding function
def embed(docs: list[str]) -> list[list[float]]:
    doc_embeds = co.embed(
        docs,
        input_type="search_document",
        model="embed-english-v3.0"
    )
    return doc_embeds.embeddings

# openai embedding function
def embed(docs: list[str]) -> list[list[float]]:
    res = openai.embeddings.create(
        input=docs,
        model="text-embedding-ada-002"
    )
    doc_embeds = [r.embedding for r in res.data]
    return doc_embeds

借助 Cohere 和 OpenAI，我们可以进行简单的 API 调用。除了 Cohere API 的 input_type 参数外，这里几乎没有什么需要注意的。input_type 定义当前输入是文档向量还是查询向量。我们定义此参数是为了支持提高非对称语义搜索的性能 — 在非对称语义搜索中，我们使用较小的文本块（即搜索查询）进行查询，并尝试检索较大的文本块（即几个句子或段落）。

3.2.2. 本地嵌入

E5 的工作方式类似于 Cohere 嵌入模型，但支持非对称搜索。但是，实现略有不同。我们不是通过参数指定输入是查询还是文档，而是将该信息作为输入文本的前缀。对于 query，我们加上 “query：” 的前缀，对于文档，我们加上 “passage：” （文档的另一个名称）。

def embed(docs: list[str]) -> list[list[float]]:
    docs = [f"passage: {d}" for d in docs]
    # tokenize
    tokens = tokenizer(
        docs, padding=True, max_length=512, truncation=True, return_tensors="pt"
    ).to(device)
    with torch.no_grad():
        # process with model for token-level embeddings
        out = model(**tokens)
        # mask padding tokens
        last_hidden = out.last_hidden_state.masked_fill(
            ~tokens["attention_mask"][..., None].bool(), 0.0
        )
        # create mean pooled embeddings
        doc_embeds = last_hidden.sum(dim=1) / \
            tokens["attention_mask"].sum(dim=1)[..., None]
    return doc_embeds.cpu().numpy()

在指定这些块是文档后，我们对它们进行标记化以给我们 tokens 参数。每个基于 transformer 的模型都需要一个标记化步骤。分词化是将人类可读的纯文本转换为 transformer 可读的输入，它只是一个整数列表，如 [0， 531， 81， 944， ...]，其中每个整数代表一个单词或子单词。

一旦我们有了我们的 token，我们就用 model（**tokens） 将它们输入到我们的模型中。由此，我们在 out 参数中得到我们的输出 logits（即预测）。

我们的一些 input token 是 padding token。这些用作占位符，以对齐我们通过每个模型层馈送的数组/张量的维度。默认情况下，我们忽略这些 token 生成的输出 logits，但在某些情况下（例如，使用嵌入模型），我们必须计算所有输出 logit 的平均值。如果我们在此计算中考虑 padding 标记生成的输出 logit，我们将降低嵌入质量。

为避免嵌入质量下降，我们必须屏蔽（即，通过设置为 None 来隐藏）填充标记生成的输出 logits。这就是 out.last_hidden_state.masked_fill 行的作用。

最后，我们准备好计算我们的单向量嵌入——我们通过均值池化来实现。均值汇聚意味着从所有输出 logit 向量中获取平均值以生成单个向量，我们将其存储在 doc_embeds 参数中。

从那里，我们使用 .cpu（）将 doc_embeds 从 GPU 移动到 CPU（如果我们使用 GPU），并使用 .numpy（） 将 doc_embeds的 PyTorch 张量转换为 Numpy 数组。

3.3 构建向量索引

定义所选 embed 函数后，我们可以使用相同的逻辑创建向量索引。我们定义一个 batch_size 并迭代我们的数据集以创建嵌入并将它们添加到名为 arr 的局部向量索引中。

from tqdm.auto import tqdm
import numpy as np

chunks = data["chunk"]
batch_size = 256

for i in tqdm(range(0, len(chunks), batch_size)):
    i_end = min(len(chunks), i+batch_size)
    chunk_batch = chunks[i:i_end]
    # embed current batch
    embed_batch = embed(chunk_batch)
    # add to existing np array if exists (otherwise create)
    if i == 0:
        arr = embed_batch.copy()
    else:
        arr = np.concatenate([arr, embed_batch.copy()])

在这里，我们可以测量两个指标 — 嵌入延迟和向量维数。在 Google Colab 上运行所有这些操作，我们看到为每个模型的整个数据集编制索引所花费的时间是：

型	批量大小	所用时间	矢量暗淡
嵌入式英语 v3.0	128	05:32	1024
文本嵌入 ADA 002	128	09:07	1536
intfloat/e5-base-v2	256	03:53	768

Ada 002 是这里最慢的方法。E5 是最快的，_但是_ 在 Google Colab 中的 V100 GPU 实例上运行 — API 模型不需要我们在 GPU 实例上运行代码。另一个考虑因素是存储要求。更高维度的向量的存储成本更高，并且这些成本会随着时间的推移而累积。

四、性能

在测试这些模型时，我们将看到相对相似的结果。我们使用的是杂乱的数据集，这更具挑战性 _但_ 也更现实。

Q1：为什么我应该使用 Llama 2？

注意：为简洁起见，我们在下面转述了结果。有关完整结果和代码 [Ada 002、Cohere embed v3、E5 base v2]，请参阅原始笔记本。