【拥抱AI】如何使用BERT等预训练模型计算语义相似度

使用BERT等预训练模型计算语义相似度是一种非常有效的方法，可以捕捉句子之间的深层次语义关系。下面是一个详细的步骤指南，介绍如何使用BERT和Sentence-BERT来计算语义相似度。

1. 环境准备

1.1 安装必要的库

首先，确保你已经安装了必要的Python库。这里我们使用transformers库来加载BERT模型，使用sentence-transformers库来加载Sentence-BERT模型。

pip install transformers sentence-transformers

2. 加载预训练模型

2.1 加载BERT模型

BERT模型可以将输入文本编码为向量表示，但默认情况下，BERT模型的输出是每个token的向量表示，而不是整个句子的向量表示。为了得到句子级别的向量表示，我们可以使用池化操作（如取平均值或使用CLS标记）。

from transformers import BertTokenizer, BertModel
import torch

# 加载预训练的BERT模型和分词器
model_name = 'bert-base-uncased'
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertModel.from_pretrained(model_name)

# 输入文本
reference = "This is a test."
candidate = "This is a test."

# 编码为token IDs
input_ids = tokenizer([reference, candidate], padding=True, truncation=True, return_tensors='pt')

# 获取模型的输出
with torch.no_grad():
    outputs = model(**input_ids)

# 获取最后一层的隐藏状态
last_hidden_states = outputs.last_hidden_state

# 使用CLS标记的向量作为句子表示
cls_embeddings = last_hidden_states[:, 0, :]

# 计算余弦相似度
cosine_similarity = torch.nn.functional.cosine_similarity(cls_embeddings[0], cls_embeddings[1], dim=0)
print(f"BERT Semantic Similarity: {
     
     cosine_similarity.item()}")

2.2 加载Sentence-BERT模型

Sentence-BERT是BERT的一个变种，专门用于生成句子级别的向量表示。Sentence-BERT通过微调BERT模型，使其能够更有效地生成句子级别的向量表示。

from sentence_transformers import SentenceTransformer, util

# 加载预训练的Sentence-BERT模型
model = SentenceTransformer('paraphrase-MiniLM-L6-v2')

# 输入文本
reference = "This is a test."
candidate = "This is a test."

# 编码为向量
reference_embedding = model.encode(reference, convert_to_tensor=True)
candidate_embedding = model.encode(candidate, convert_to_tensor=True)

# 计算余弦相似度
similarity = util.pytorch_cos_sim(reference_embedding, candidate_embedding).item()
print(f"Sentence-BERT Semantic Similarity: {
     
     similarity}")

3. 详细步骤

3.1 数据预处理

1. 分词：将输入文本分词为token序列。
2. 编码：将token序列转换为模型可以接受的输入格式（如token IDs）。

# 输入文本
reference = "This is a test."
candidate = "This is a test."

# 编码为token IDs
input_ids = tokenizer([reference, candidate], padding=True, truncation=True, return_tensors='pt')

3.2 模型推理

1. 前向传播：将输入传递给模型，获取模型的输出。
2. 提取向量：从模型的输出中提取句子级别的向量表示。

# 获取模型的输出
with torch.no_grad():
    outputs = model(**input_ids)

# 获取最后一层的隐藏状态
last_hidden_states = outputs.last_hidden_state

# 使用CLS标记的向量作为句子表示
cls_embeddings = last_hidden_states[:, 0, :]

3.3 计算相似度

1. 余弦相似度：计算两个向量之间的余弦相似度，范围从-1到1，值越接近1表示相似度越高。

# 计算余弦相似度
cosine_similarity = torch.nn.functional.cosine_similarity(cls_embeddings[0], cls_embeddings[1], dim=0)
print(f"BERT Semantic Similarity: {
     
     cosine_similarity.item(