戒酒的李白

The integration process and a complete use example are given

import torch
import numpy as np
from transformers.models.bert import BertTokenizer, BertModel
from MHA import MultiHeadAttentionLayer
from classifier import FinalClassifier
# 加载BERT模型并生成嵌入
def get_sentence_embeddings(sentences, bert_model_path, max_length=80):
"""使用BERT生成多个句子的嵌入"""
tokenizer = BertTokenizer.from_pretrained(bert_model_path)
model = BertModel.from_pretrained(bert_model_path)
embeddings = []
for sentence in sentences:
inputs = tokenizer(sentence, return_tensors="pt", padding="max_length", truncation=True, max_length=max_length)
with torch.no_grad():
outputs = model(**inputs)
embedding = outputs.last_hidden_state.cpu().numpy()
embeddings.append(embedding)
return np.vstack(embeddings) # 保持多句子输出格式一致
# 加载已经训练好的模型
def load_model(model_path):
print(f"加载模型 {model_path}...")
model = torch.load(model_path)
model.eval() # 设置为评估模式
return model
# 多句子的预测函数
def predict_sentences(sentences, model, bert_model_path, max_length=80):
# 检查是否为单个句子输入,如果是,将其包装为列表
if isinstance(sentences, str):
sentences = [sentences]
# 生成句子的BERT嵌入
embeddings = get_sentence_embeddings(sentences, bert_model_path, max_length)
# 转换为Tensor
embedding_tensors = torch.tensor(embeddings, dtype=torch.float32).squeeze(1) # 修改squeeze以适应多个句子
# 检查嵌入维度是否符合注意力层要求
embed_size = embedding_tensors.size(-1)
num_heads = 12
if embed_size % num_heads != 0:
raise ValueError(f"嵌入维度 {embed_size} 无法被注意力头数量 {num_heads} 整除")
# 加载多头注意力机制
attention_model = MultiHeadAttentionLayer(embed_size=embed_size, num_heads=num_heads)
predictions = []
with torch.no_grad():
for embedding_tensor in embedding_tensors:
attention_output = attention_model(embedding_tensor.unsqueeze(0), embedding_tensor.unsqueeze(0),
embedding_tensor.unsqueeze(0))
outputs = model(attention_output)
outputs = torch.mean(outputs, dim=1)
_, predicted = torch.max(outputs, 1) # 获取预测的类别
predictions.append(predicted.item())
return predictions
if __name__ == "__main__":
# 加载已经训练好的模型
model_path = './final_model.pt'
model = load_model(model_path)
# 需要预测的句子,可以输入单个句子或多个句子
sentences = ["这是一条待预测的句子",
"他在你面前骂黑鬼 印度屎屁尿背后就会根人家骂你中国猴子,这可能不是种族歧视这是素质太低",
"完美女朋友",
"在美国的亚裔就是一盘散沙。日裔看不起韩裔 韩裔仇视日裔 港澳台裔看不起大陆裔,大陆裔里面又歧视福建裔"] # 可以替换为单个句子或多个句子
# BERT模型路径
bert_model_path = './bert_model'
# 对句子进行预测
predicted_labels = predict_sentences(sentences, model, bert_model_path)
# 根据预测的label输出对应的文本
for i, label in enumerate(predicted_labels):
if label == 1:
print(f"句子: '{sentences[i]}' 预测结果: 不良言论")
elif label == 0:
print(f"句子: '{sentences[i]}' 预测结果: 正常言论")
else:
print(f"句子: '{sentences[i]}' 未知标签: {label}")
... ...