使用 Transformers 生成文本嵌入

作者 Muhammad Asad Iqbal Khan 于 2025 年 5 月 15 日发布在 Hugging Face Transformers 0

文本嵌入是文本的数值表示，它能以机器能够理解和处理的方式捕获语义含义。这些嵌入革新了自然语言处理，使计算机能够比传统的词袋模型或独热编码方法更具意义地处理文本。

接下来，您将探索如何使用 Hugging Face Hub 的 Transformer 模型生成高质量的文本嵌入。具体来说，您将学习：

什么是文本嵌入
如何从 BERT 模型生成文本嵌入
如何生成更高质量的嵌入

通过我的书籍《Hugging Face Transformers中的NLP》，快速启动您的项目。它提供了带有工作代码的自学教程。

让我们开始吧！

使用 Transformers 生成文本嵌入
照片由 Greg Rivers 拍摄。部分权利保留。

概述

这篇博文分为三部分；它们是：

理解文本嵌入
生成嵌入的其他技术
如何获得高质量的文本嵌入？

理解文本嵌入

文本嵌入是使用数值向量来表示文本。表示文本的一种简单方法是找到词典中的所有单词，并为每个单词分配一个唯一的数字。然后，您可以将每个单词表示为独热向量，或将句子表示为词袋向量：每个位置中的数字表示单词在句子中出现的次数。

一个词典有成千上万个单词，所以独热向量太大且稀疏。密集向量，其中每个元素都是浮点数而不是布尔值，可以使其更紧凑。但是，向量中的每个元素应该是什么值？这不容易确定。但它可以被训练。例如 Word2Vec、GloVe 和 FastText。密集词向量的有趣之处在于，它们将语义相似的词放在向量空间中更靠近的位置。您可以使用该向量来衡量单词之间的语义相似性并执行单词数学运算，例如“国王 - 男人 + 女人 = 女王”。

再进一步，您想将句子表示为向量。这比简单地将句子中的单词的词向量相加更困难，因为您需要识别单词的上下文。例如，“bear”可以是动词或名词；词向量无法区分它们的差异，但这对于上下文很重要。将句子的语义含义表示为向量对于许多 NLP 任务都非常有用。

Transformer 模型可以通过一次处理整个单词序列来生成此类上下文嵌入。嵌入中单词的表示取决于其在文本中的上下文。这使得表示更加丰富，能够捕获多义性（具有多个含义的单词）等细微差别。

训练 Transformer 模型来生成嵌入在计算上是昂贵且困难的，因为它需要高质量的数据集和复杂的训练过程。幸运的是，如果我们只想创建向量来表示文本的语义含义，我们可以使用预训练模型来生成嵌入。

让我们看看如何使用预训练的 BERT 模型为句子生成嵌入，该模型以生成高质量的上下文嵌入而闻名。

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

# Load pre-trained model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased")

# Define some example sentences
sentences = [
    "The cat sat on the mat.",
    "The dog slept on the floor.",
    "I love natural language processing."
]

def get_embeddings(sentences, model, tokenizer):
    "Function to get embeddings for a batch of sentences"

    # Tokenize input and get model output
    encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors="pt")
    with torch.no_grad():
        model_output = model(**encoded_input)

    # Use the CLS token embedding as the sentence embedding
    sentence_embeddings = model_output.last_hidden_state[:, 0, :]

    # Convert torch tensor to numpy array for easier handling
    return sentence_embeddings.numpy()

# Get embeddings for our example sentences
embeddings = get_embeddings(sentences, model, tokenizer)
print(f"Embedding shape: {embeddings.shape}")
print(f"First 5 dimensions of the sentences' embeddings:\n{np.round(embeddings[:, :5], 3)}")

from transformers import AutoTokenizer, AutoModel

import torch

import numpy as np

# 加载预训练模型和分词器

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

model = AutoModel.from_pretrained("bert-base-uncased")

# 定义一些示例句子

sentences = [

"The cat sat on the mat.",

"The dog slept on the floor.",

"I love natural language processing."

]

def get_embeddings(sentences, model, tokenizer):

"Function to get embeddings for a batch of sentences"

# Tokenize input and get model output

encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors="pt")

with torch.no_grad():

model_output = model(**encoded_input)

# Use the CLS token embedding as the sentence embedding

sentence_embeddings = model_output.last_hidden_state[:, 0, :]

# Convert torch tensor to numpy array for easier handling

return sentence_embeddings.numpy()

# Get embeddings for our example sentences

embeddings = get_embeddings(sentences, model, tokenizer)

print(f"Embedding shape: {embeddings.shape}")

print(f"First 5 dimensions of the sentences' embeddings:\n{np.round(embeddings[:, :5], 3)}")

在此示例中，您使用预训练的 BERT 模型为三个示例句子生成嵌入。您需要同时使用 BERT 的分词器和模型。分词器将句子分割成子词标记，模型生成上下文嵌入。分词器和模型是使用 transformers 库中的“auto-class”创建的。您只需要指定预训练模型的名称，bert-base-uncased。

基础 BERT 模型有 12 层和 768 个隐藏维度。它是大小写不敏感的，这意味着输入文本被视为大小写不敏感。由于模型具有 768 的隐藏维度，因此为每个句子生成的嵌入是 768 维的向量。

get_embeddings() 函数接受一个句子列表、一个模型和一个分词器，并为每个句子返回嵌入。它的工作方式很简单。但请注意，句子嵌入是从模型输出的第一个标记中提取的。

...
sentence_embeddings = model_output.last_hidden_state[:, 0, :]

1 2	... sentence_embeddings = model_output.last_hidden_state[:, 0, :]

第一个标记是 [CLS] 标记，这是分词器添加到每个句子开头的特殊标记。这是模型训练用于表示句子的内容。您可以将其视为整个句子的摘要。在分词器中，您设置了 truncation=True 以防止将过长的序列发送给模型。您还设置了 return_tensors=" pt" 以获取 PyTorch 张量，这是您的模型所期望的。

最后，在函数末尾，您将嵌入转换为 numpy 数组，将其与 PyTorch 分离并移回 CPU。上述代码的输出是：

Embedding shape: (3, 768)
First 5 dimensions of the sentences' embeddings:
[[-0.364 -0.053 -0.367 -0.03  -0.461]
 [-0.276 -0.043 -0.613  0.175 -0.309]
 [-0.042  0.043 -0.253 -0.35  -0.374]]

Embedding shape: (3, 768)

句子嵌入的前 5 个维度

[[-0.364 -0.053 -0.367 -0.03 -0.461]

[-0.276 -0.043 -0.613 0.175 -0.309]

[-0.042 0.043 -0.253 -0.35 -0.374]]

您可以验证每个上下文嵌入向量的长度是 768。

生成嵌入的其他技术

虽然使用 [CLS] 标记嵌入是一种常见方法，但它不是唯一的方法。

平均池化

回想一下，BERT 模型是一个 Transformer 模型，它接受令牌序列作为输入并生成序列作为输出。如果您可以使用 [CLS] 前缀标记进行嵌入，您还可以取所有输出标记的平均值。这就是平均池化技术。它可能提供更好的句子表示。

让我们看看如何修改之前的代码以使用平均池化。

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

# Load pre-trained model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased")

# Define some example sentences
sentences = [
    "The cat sat on the mat.",
    "The dog slept on the floor.",
    "I love natural language processing."
]

def get_embeddings(sentences, model, tokenizer):
    "Function to get embeddings for a batch of sentences with mean pooling"

    # Tokenize input and get model output
    encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors="pt")
    with torch.no_grad():
        model_output = model(**encoded_input)

    # Extract the attention mask and output sequence
    attention_mask = encoded_input["attention_mask"]
    output_seq = model_output.last_hidden_state

    # Mean pooling: take the average of all token embeddings
    mask = attention_mask.unsqueeze(-1).expand(output_seq.size()).float()
    sum_embeddings = (output_seq * mask).sum(1)
    sum_mask = torch.clamp(mask.sum(1), min=1e-9)
    mean_pooled = sum_embeddings / sum_mask

    # Convert torch tensor to numpy array for easier handling
    return mean_pooled.numpy()

# Get embeddings with mean pooling
embeddings = get_embeddings(sentences, model, tokenizer)
print(f"Embedding shape: {embeddings.shape}")
print(f"First 5 dimensions of the sentences' embeddings:\n{np.round(embeddings[:, :5], 3)}")

rom transformers import AutoTokenizer, AutoModel

import torch

import numpy as np

# 加载预训练模型和分词器

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

model = AutoModel.from_pretrained("bert-base-uncased")

# 定义一些示例句子

sentences = [

"The cat sat on the mat.",

"The dog slept on the floor.",

"I love natural language processing."

]

def get_embeddings(sentences, model, tokenizer):

"Function to get embeddings for a batch of sentences with mean pooling"

# Tokenize input and get model output

encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors="pt")

with torch.no_grad():

model_output = model(**encoded_input)

# Extract the attention mask and output sequence

attention_mask = encoded_input["attention_mask"]

output_seq = model_output.last_hidden_state

# Mean pooling: take the average of all token embeddings

mask = attention_mask.unsqueeze(-1).expand(output_seq.size()).float()

sum_embeddings = (output_seq * mask).sum(1)

sum_mask = torch.clamp(mask.sum(1), min=1e-9)

mean_pooled = sum_embeddings / sum_mask

# Convert torch tensor to numpy array for easier handling

return mean_pooled.numpy()

# Get embeddings with mean pooling

embeddings = get_embeddings(sentences, model, tokenizer)

print(f"Embedding shape: {embeddings.shape}")

print(f"First 5 dimensions of the sentences' embeddings:\n{np.round(embeddings[:, :5], 3)}")

关键区别在于 get_embeddings() 函数。首先，您利用了分词器输出中的注意力掩码。它是一个二进制张量，指示哪些标记是真实标记 (1)，哪些是填充标记 (0)。它的形状是 (batch size, sequence length)，但模型输出的形状是 (batch size, sequence length, hidden dimension)。因此，您使用 unsqueeze(-1) 在注意力掩码的末尾添加一个额外的维度，并将其扩展以匹配模型输出的形状。

然后，通过将模型输出序列与注意力掩码相乘来计算所有嵌入向量的总和，其中掩码值为 0，不会计入总和。总和在第二个维度（即 axis=1）上计算，对应于序列长度维度。

然后通过将总和除以掩码的总和来计算平均值。由于掩码是 1 或 0，因此掩码的总和表示序列中有多少个非填充元素。为避免除以零，您使用 torch.clamp() 来确保掩码的总和至少为 1e-9。

以上代码的输出是：

Embedding shape: (3, 768)
First 5 dimensions of the sentences' embeddings:
[[-0.182 -0.266 -0.219  0.211  0.285]
 [-0.056 -0.208 -0.281  0.223  0.417]
 [ 0.428  0.355 -0.182 -0.048  0.142]]

Embedding shape: (3, 768)

句子嵌入的前 5 个维度

[[-0.182 -0.266 -0.219 0.211 0.285]

[-0.056 -0.208 -0.281 0.223 0.417]

[ 0.428 0.355 -0.182 -0.048 0.142]]

平均池化方法被认为比仅使用 [CLS] 标记提供更好的句子嵌入，尤其适用于语义相似性和信息检索等任务。

使用 Sentence Transformers

BERT 是一个通用模型，上例中使用的是基础模型，它应该与不同的“头部”一起用于特定任务。例如，[CLS] 标记是在其原始论文中为分类任务提出的。因此，它可能不是生成句子嵌入的最佳选择。您可能会发现生成的嵌入向量不具备您期望的属性，例如句子之间的余弦相似性不能反映它们的语义相似性。

事实上，没有什么能阻止您对 BERT 模型或任何其他 Transformer 模型进行微调以生成更好的句子嵌入。但如果您不想经历麻烦，可以使用 Sentence Transformers 库，它提供了专门为生成高质量句子嵌入而微调的模型。它还托管了 Hugging Face Hub 的预训练模型。让我们看看如何使用它们。

Sentence Transformers 库提供了专门为生成高质量句子嵌入而微调的模型。这是一个独立的 Python 库。您可以使用以下命令安装它：

pip install sentence-transformers

1	pip install sentence-transformers

让我们看看如何使用它们。

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

# Define some example sentences
sentences = [
    "The cat sat on the mat.",
    "The dog slept on the floor.",
    "I love natural language processing."
]

# Load a pre-trained model and generate embeddings
model = SentenceTransformer("all-MiniLM-L6-v2")
embeddings = model.encode(sentences)

# Get embeddings with mean pooling
print(f"Embedding shape: {embeddings.shape}")
print(f"First 5 dimensions of the sentences' embeddings:\n{np.round(embeddings[:, :5], 3)}")

# Calculate cosine similarity between the first two sentences
similarity = cosine_similarity([embeddings[0]], [embeddings[1]])
print(f"Cosine similarity between '{sentences[0]}' and '{sentences[1]}': {np.round(similarity[0][0], 3)}")

from sentence_transformers import SentenceTransformer

from sklearn.metrics.pairwise import cosine_similarity

import numpy as np

# 定义一些示例句子

sentences = [

"The cat sat on the mat.",

"The dog slept on the floor.",

"I love natural language processing."

]

# Load a pre-trained model and generate embeddings

model = SentenceTransformer("all-MiniLM-L6-v2")

embeddings = model.encode(sentences)

# Get embeddings with mean pooling

print(f"Embedding shape: {embeddings.shape}")

print(f"First 5 dimensions of the sentences' embeddings:\n{np.round(embeddings[:, :5], 3)}")

# Calculate cosine similarity between the first two sentences

similarity = cosine_similarity([embeddings[0]], [embeddings[1]])

print(f"Cosine similarity between '{sentences[0]}' and '{sentences[1]}': {np.round(similarity[0][0], 3)}")

代码更短，因为 Sentence Transformers 库中的模型一步处理了分词和嵌入生成。请注意，Sentence Transformers 模型与可以通过 transformers 库实例化的模型不同。您必须确保模型名称受 Sentence Transformers 库支持，或者您可以从该库的文档中选择一些“原始”预训练模型。

示例中使用的模型是 all-MiniLM-L6-v2。它很小，因此运行速度更快，占用的内存也更少。它输出一个 384 维的嵌入。要说明专门的句子嵌入模型为何更好，您可以比较前两个句子之间的余弦相似性。

$$
\cos(\theta_{\mathbf{a}, \mathbf{b}}) = \frac{ \mathbf{a} \cdot \mathbf{b} }{ \vert \mathbf{a} \vert \vert \mathbf{b} \vert }
$$

Scikit-learn 将余弦相似性实现为 cosine_similarity() 函数，该函数接受两个矩阵作为输入，并计算两个矩阵之间每一行对的余弦相似性。因此，如果您有两个向量，您必须将它们分别封装在列表中。

以上代码的输出是：

Embedding shape: (3, 384)
First 5 dimensions of the sentences' embeddings:
[[ 0.13  -0.016 -0.037  0.058 -0.06 ]
 [ 0.01  -0.01  -0.039  0.14  -0.006]
 [ 0.039 -0.078  0.055  0.     0.036]]
Cosine similarity between 'The cat sat on the mat.' and 'The dog slept on the floor.': 0.408

Embedding shape: (3, 384)

句子嵌入的前 5 个维度

[[ 0.13 -0.016 -0.037 0.058 -0.06 ]

[ 0.01 -0.01 -0.039 0.14 -0.006]

[ 0.039 -0.078 0.055 0. 0.036]]

Cosine similarity between 'The cat sat on the mat.' and 'The dog slept on the floor.': 0.408

如果您想比较所有句子对之间的相似性，

...
print(cosine_similarity(embeddings, embeddings).round(3))

1 2	... print(cosine_similarity(embeddings, embeddings).round(3))

这将为您提供一个对称的 3x3 矩阵，所有对角线元素均为 1。非对角线元素是句子之间的余弦相似性。

如果您比较上面所有示例的嵌入结果，您会发现 Sentence Transformer 模型在第一对句子（0.408）的余弦相似性上比最后两个句子（-0.028）提供了更好的区分。相比之下，第一个示例（仅使用 [CLS] 标记）并未提供良好的区分（0.941 vs 0.792）。因此，您可以看到 Sentence Transformer 模型输出的质量更高。

如何获得高质量的文本嵌入？

所有句子嵌入都来自深度学习模型，尤其是 Transformer 模型。嵌入的质量在很大程度上取决于模型和训练数据的质量。

更大的模型，如 BERT 和 RoBERTa，通常比小型模型（如 DistilBERT）更好，它们以牺牲速度和内存使用来换取质量。对于特定任务训练或微调的模型，如果用于专业领域，也可能比通用模型提供更好的嵌入。例如，用医学领域语料库训练的模型可能比通用模型提供更好的医学文本嵌入。

另外，请注意分词器在嵌入质量中起着重要作用。Transformer 模型处理标记序列。将句子分割成保留语义含义的子词的分词器有助于模型生成更好的嵌入。在极端情况下，分词器可以将每个单独的字符作为标记发出，但这会在序列馈送到模型时丢失大量信息。词汇量更大的分词器，使标记更有可能是有意义的单词，将有助于模型理解上下文。但它也会使模型变大。