基于Bert的相关书名检索推荐算法

# -*- coding: utf-8 -*-
import pandas as pd
from sklearn.neural_network import BernoulliRBM
from transformers import BertTokenizer, BertModel
import torch
from sklearn.metrics.pairwise import cosine_similarity
# 步骤1：读取CSV文件，并选择用户指定的属性列
data =pd.read_csv(r'中国.csv',encoding='gbk') # 替换为实际的文件路径和文件名
new_documents = pd.DataFrame({'书名': ["太白金星有点不太高兴"]})
datas=pd.concat([data,new_documents],ignore_index=True)
#print(datas)
text_data = data['书名'].astype(str)  # 获取指定属性列并转换为字符串类型
# print(text_data)
# 步骤2：使用BERT tokenizer对文本数据进行tokenization
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')  # 加载预训练的BERT tokenizer
tokens = tokenizer.batch_encode_plus(
    text_data.tolist(),
    add_special_tokens=True,
    max_length=16,  # 根据需求自定义最大长度
    padding='max_length',
    truncation=True,
    return_tensors='pt')
# 步骤2：使用BERT模型生成词向量
model = BertModel.from_pretrained('bert-base-uncased')  # 加载预训练的BERT模型，可根据选择相应型
input_ids = tokens['input_ids']
attention_mask = tokens['attention_mask']
with torch.no_grad():
    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
embeddings = outputs.last_hidden_state  # 取得BERT模型最后一层的输出作为词向量表示
reshaped_embeddings = embeddings.reshape(embeddings.shape[0], -1)
print(reshaped_embeddings.shape)
# 步骤7：使用特征表示进行信息检索
new_documents = ["我在北京"]
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')  # 加载预训练的BERT tokenizer
tokens = tokenizer.batch_encode_plus(
    new_documents,
    add_special_tokens=True,
    max_length=16,  # 根据需求自定义最大长度
    padding='max_length',
    truncation=True,
    return_tensors='pt')
# 步骤5：使用BERT模型生成词向量
model = BertModel.from_pretrained('bert-base-uncased')  # 加载预训练的BERT模型，可根据选择相应型
input_ids = tokens['input_ids']
attention_mask = tokens['attention_mask']
with torch.no_grad():
    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
embedding1 = outputs.last_hidden_state
embedding11= embedding1.reshape(embedding1.shape[0], -1)  # 取得BERT模型最后一层的输出作为词向量表示
# # 输出词向量表示
print(embedding11.shape)
# 步骤6：使用RBM进行特征提取和训练
# new_rbm_features = rbm.transform(embedding11)
similarity_matrix = cosine_similarity(reshaped_embeddings, embedding11)
# # 输出相似度矩阵
# print("Similarity Matrix:")
# print(similarity_matrix)
pp=pd.DataFrame(similarity_matrix,columns=["相似度"])
new=pd.concat([data,pp],ignore_index=True)
new.to_csv("chinese.csv",index=False)
da=new["相似度"].astype(float)
# print(da)
max_value, max_index = da.max(), da.idxmax()
print("最大的相似度为：", max_value)
print("最大相似度的序号数为：", max_index)