|
@@ -1,10 +1,16 @@
|
|
|
# 导入标准库
|
|
# 导入标准库
|
|
|
import json
|
|
import json
|
|
|
import os
|
|
import os
|
|
|
|
|
+import sys
|
|
|
|
|
+
|
|
|
|
|
+current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
|
+rag_qa_path = os.path.abspath(os.path.dirname(current_dir))
|
|
|
|
|
+project_root = os.path.abspath(os.path.dirname(rag_qa_path))
|
|
|
|
|
+sys.path.insert(0, project_root)
|
|
|
|
|
+
|
|
|
# 导入 PyTorch
|
|
# 导入 PyTorch
|
|
|
import torch
|
|
import torch
|
|
|
# 导入日志
|
|
# 导入日志
|
|
|
-import sys
|
|
|
|
|
from base.logger import logger
|
|
from base.logger import logger
|
|
|
# 导入numpy
|
|
# 导入numpy
|
|
|
import numpy as np
|
|
import numpy as np
|
|
@@ -15,13 +21,11 @@ from transformers import Trainer, TrainingArguments
|
|
|
from sklearn.model_selection import train_test_split
|
|
from sklearn.model_selection import train_test_split
|
|
|
from sklearn.metrics import classification_report, confusion_matrix
|
|
from sklearn.metrics import classification_report, confusion_matrix
|
|
|
|
|
|
|
|
-current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
|
-rag_qa_path = os.path.abspath(os.path.dirname(os.path.abspath(current_dir)))
|
|
|
|
|
-project_root = os.path.abspath(os.path.dirname(os.path.abspath(rag_qa_path)))
|
|
|
|
|
-sys.path.insert(0, project_root)
|
|
|
|
|
|
|
+default_model_path = os.path.join(current_dir, "bert_query_classifier")
|
|
|
|
|
+default_data_file = os.path.join(rag_qa_path, "classify_data", "model_generic_5000.json")
|
|
|
|
|
|
|
|
class QueryClassifier:
|
|
class QueryClassifier:
|
|
|
- def __init__(self, model_path='models/bert_query_classifier'):
|
|
|
|
|
|
|
+ def __init__(self, model_path=default_model_path):
|
|
|
# 加载bert
|
|
# 加载bert
|
|
|
self.pre_trained_model_path = f'{rag_qa_path}/models/bert-base-chinese'
|
|
self.pre_trained_model_path = f'{rag_qa_path}/models/bert-base-chinese'
|
|
|
# 模型训练以后保存的位置
|
|
# 模型训练以后保存的位置
|
|
@@ -104,7 +108,7 @@ class QueryClassifier:
|
|
|
|
|
|
|
|
return Dataset(encodings, labels)
|
|
return Dataset(encodings, labels)
|
|
|
|
|
|
|
|
- def train_model(self, data_file="training_dataset_hybrid_5000.json"):
|
|
|
|
|
|
|
+ def train_model(self, data_file=default_data_file):
|
|
|
"""训练 BERT 分类模型"""
|
|
"""训练 BERT 分类模型"""
|
|
|
# 加载数据集
|
|
# 加载数据集
|
|
|
if not os.path.exists(data_file):
|
|
if not os.path.exists(data_file):
|
|
@@ -141,7 +145,7 @@ class QueryClassifier:
|
|
|
# 设置训练参数
|
|
# 设置训练参数
|
|
|
training_args = TrainingArguments(
|
|
training_args = TrainingArguments(
|
|
|
# 设置模型和检查点保存的目录路径
|
|
# 设置模型和检查点保存的目录路径
|
|
|
- output_dir="./bert_results",
|
|
|
|
|
|
|
+ output_dir=os.path.join(current_dir, "bert_results"),
|
|
|
# 设置训练的总轮数为3轮
|
|
# 设置训练的总轮数为3轮
|
|
|
num_train_epochs=3,
|
|
num_train_epochs=3,
|
|
|
# 设置每个设备(GPU/CPU)上的训练批次大小为8
|
|
# 设置每个设备(GPU/CPU)上的训练批次大小为8
|
|
@@ -153,7 +157,7 @@ class QueryClassifier:
|
|
|
# 设置权重衰减系数为0.01,用于防止过拟合
|
|
# 设置权重衰减系数为0.01,用于防止过拟合
|
|
|
weight_decay=0.01,
|
|
weight_decay=0.01,
|
|
|
# 设置日志文件保存的目录路径
|
|
# 设置日志文件保存的目录路径
|
|
|
- logging_dir="./bert_logs",
|
|
|
|
|
|
|
+ logging_dir=os.path.join(current_dir, "bert_logs"),
|
|
|
# 设置每10个训练步骤记录一次日志
|
|
# 设置每10个训练步骤记录一次日志
|
|
|
logging_steps=10,
|
|
logging_steps=10,
|
|
|
# 设置评估策略为每个epoch结束后进行评估
|
|
# 设置评估策略为每个epoch结束后进行评估
|
|
@@ -246,10 +250,10 @@ class QueryClassifier:
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if __name__ == "__main__":
|
|
|
# 初始化分类器
|
|
# 初始化分类器
|
|
|
- classifier = QueryClassifier(model_path="bert_query_classifier")
|
|
|
|
|
|
|
+ classifier = QueryClassifier()
|
|
|
|
|
|
|
|
# 训练模型
|
|
# 训练模型
|
|
|
- classifier.train_model(data_file='../classify_data/model_generic_5000.json')
|
|
|
|
|
|
|
+ classifier.train_model()
|
|
|
# 示例预测
|
|
# 示例预测
|
|
|
test_queries = [
|
|
test_queries = [
|
|
|
"AI知识域的知识文档目录是什么",
|
|
"AI知识域的知识文档目录是什么",
|