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()}")