| 12345678910111213141516171819202122232425262728293031323334 |
- import torch
- from transformers import BertTokenizer, BertModel
- from torch.nn.functional import cosine_similarity, normalize
- # 加载预训练的BERT分词器和模型
- tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
- model = BertModel.from_pretrained('bert-base-chinese')
- # 准备文本
- texts = ["今天天气真好", "今天天气不错"]
- # 编码文本
- encoded_inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
- # 使用BERT模型获取文本的向量表示
- with torch.no_grad():
- outputs = model(**encoded_inputs)
- last_hidden_states = outputs.last_hidden_state
- # 通常取[CLS] token的隐藏状态作为整个句子的表示
- # 假设batch_size为文本的数量（在这个例子中是2）
- cls_vectors = last_hidden_states[:, 0, :] # 取每个序列的第一个token（即[CLS]）
- # 归一化向量
- cls_vectors_normalized = normalize(cls_vectors, p=2, dim=1)
- # 计算两段文本之间的余弦相似度
- # 注意：我们需要将向量reshape为(1, embedding_dim)来匹配cosine_similarity的输入要求
- similarity = cosine_similarity(cls_vectors_normalized[0:1], cls_vectors_normalized[1:2])
- # 打印相似度
- print(f"两段文本的相似度为: {similarity.item()}")
-
|
