# 导入标准库
import json
import os
# 导入 PyTorch
import torch
# 导入日志
from base import logger
# 导入numpy
import numpy as np
# 导入 Transformers 库
from transformers import BertTokenizer, BertForSequenceClassification
from transformers import Trainer, TrainingArguments
# 导入train_test_split
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix
class QueryClassifier:
def __init__(self, model_path="bert_query_classifier"):
# 初始化模型路径
self.pre_trained_model_path = r'D:\code\workspace7\integrated_qa_system\rag_qa\models\bert-base-chinese'
self.model_path = model_path
# 加载 BERT 分词器
self.tokenizer = BertTokenizer.from_pretrained(self.pre_trained_model_path)
# 初始化模型
self.model = None
# 确定设备（GPU 或 CPU）
self.device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
# 记录设备信息
logger.info(f"使用设备: {self.device}")
# 定义标签映射
self.label_map = {"通用知识": 0, "专业咨询": 1}
# 加载模型
self.load_model()
def load_model(self):
# 检查模型路径是否存在
if os.path.exists(self.model_path):
# 加载预训练模型
self.model = BertForSequenceClassification.from_pretrained(self.model_path, num_labels=2)
# 将模型移到指定设备
self.model.to(self.device)
# 记录加载成功的日志
logger.info(f"加载模型: {self.model_path}")
else:
# 初始化新模型
self.model = BertForSequenceClassification.from_pretrained(self.pre_trained_model_path, num_labels=2)
# 将模型移到指定设备
self.model.to(self.device)
# 记录初始化模型的日志
logger.info("初始化新 BERT 模型")
def save_model(self):
"""保存模型"""
logger.info(f"开始保存模型....")
self.model.save_pretrained(self.model_path)
self.tokenizer.save_pretrained(self.model_path)
logger.info(f"模型保存至: {self.model_path}")
def preprocess_data(self, texts, labels):
"""预处理数据为 BERT 输入格式"""
encodings = self.tokenizer(
texts,
truncation=True,
padding=True,
max_length=128,
return_tensors="pt"
)
return encodings, [self.label_map[label] for label in labels]
def create_dataset(self, encodings, labels):
"""创建 PyTorch 数据集"""
class Dataset(torch.utils.data.Dataset):
def __init__(self, encodings, labels):
self.encodings = encodings
self.labels = labels
def __getitem__(self, idx):
item = {key: val[idx] for key, val in self.encodings.items()}
item["labels"] = torch.tensor(self.labels[idx])
return item
def __len__(self):
return len(self.labels)
return Dataset(encodings, labels)
def train_model(self, data_file="training_dataset_hybrid_5000.json"):
"""训练 BERT 分类模型"""
# 加载数据集
if not os.path.exists(data_file):
logger.error(f"数据集文件 {data_file} 不存在")
raise FileNotFoundError(f"数据集文件 {data_file} 不存在")
with open(data_file, "r", encoding="utf-8") as f:
data = [json.loads(value) for value in f.readlines()]
texts = [item["query"] for item in data]
labels = [item["label"] for item in data]
# 数据划分
train_texts, val_texts, train_labels, val_labels = train_test_split(texts, labels, test_size=0.2, random_state=42)
# 预处理
train_encodings, train_labels = self.preprocess_data(train_texts, train_labels)
val_encodings, val_labels = self.preprocess_data(val_texts, val_labels)
# 创建数据集
train_dataset = self.create_dataset(train_encodings, train_labels)
# print(f'train_dataset--》{train_dataset[0]}')
val_dataset = self.create_dataset(val_encodings, val_labels)
#
# 设置训练参数
training_args = TrainingArguments(
output_dir="./bert_results",
num_train_epochs=1,
per_device_train_batch_size=8,
per_device_eval_batch_size=8,
warmup_steps=500,
weight_decay=0.01,
logging_dir="./bert_logs",
logging_steps=10,
eval_strategy="epoch",
save_strategy="epoch",
load_best_model_at_end=True,
save_total_limit=1, # 只保存一个检查点，即最优的模型
metric_for_best_model="eval_loss",
fp16=False, # 禁用混合精度
)
# 初始化 Trainer
trainer = Trainer(
model=self.model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=val_dataset,
compute_metrics=self.compute_metrics
)
# 训练模型
logger.info("开始训练 BERT 模型...")
trainer.train()
self.save_model()
# 评估模型
self.evaluate_model(val_texts, val_labels)
def compute_metrics(self, eval_pred):
"""计算评估指标"""
logits, labels = eval_pred
predictions = np.argmax(logits, axis=-1)
accuracy = (predictions == labels).mean()
return {"accuracy": accuracy}
def evaluate_model(self, texts, labels):
"""评估模型性能"""
# 仅对 texts 进行分词，labels 已为数字
encodings = self.tokenizer(
texts,
truncation=True,
padding=True,
max_length=128,
return_tensors="pt"
)
dataset = self.create_dataset(encodings, labels)
trainer = Trainer(model=self.model)
predictions = trainer.predict(dataset)
pred_labels = np.argmax(predictions.predictions, axis=-1)
true_labels = labels # 直接使用数字标签
logger.info("分类报告:")
logger.info(classification_report(
true_labels,
pred_labels,
target_names=["通用知识", "专业咨询"]
))
logger.info("混淆矩阵:")
logger.info(confusion_matrix(true_labels, pred_labels))
def predict_category(self, query):
# 检查模型是否加载
if self.model is None:
# 模型未加载，记录错误
logger.error("模型未训练或加载")
# 默认返回通用知识
return "通用知识"
# 对查询进行编码
encoding = self.tokenizer(query, truncation=True, padding=True, max_length=128, return_tensors="pt")
# 将编码移到指定设备
encoding = {k: v.to(self.device) for k, v in encoding.items()}
# 不计算梯度，进行预测
self.model.eval()
with torch.no_grad():
# 获取模型输出
outputs = self.model(**encoding)
# 获取预测结果
prediction = torch.argmax(outputs.logits, dim=1).item()
# 根据预测结果返回类别
return "专业咨询" if prediction == 1 else "通用知识"
if __name__ == "__main__":
# 初始化分类器
classifier = QueryClassifier(model_path="../models/bert_query_classifier")
logger.info(f"classifier-->{classifier}")
# 训练模型，如果放开下面训练模型的代码，则会在models目录下生成一个模型，名称为：bert_query_classifier
# classifier.train_model(data_file='../classify_data/model_generic_5000.json')
# logger.info("模型训练完成")
# 示例预测
test_queries = [
"AI学科的课程大纲是什么",
"JAVA课程费用多少？",
"5乘以9等于多少？",
"AI培训有哪些老师？"
]
for query in test_queries:
category = classifier.predict_category(query)
print(f"查询: {query} -> 分类: {category}")