戒酒的李白
92 lines
3.6 KiB
Python
92 lines
3.6 KiB
Python
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}")
|