0
0
Переглянути джерело

Update Milvus RAG query routing

zengli 1 місяць тому
батько
коміт
184e24be72

+ 2 - 0
.gitignore

@@ -11,6 +11,8 @@
 **/query_cls/
 **/bert_query_classifier/
 **/bert_results/
+**/bert_logs/
+**/tmp_trainer/
 **/notebooks/tokenizer.json
 **/notebooks/tokenizer_config.json
 **/notebooks/special_tokens_map.json

+ 15 - 11
4.基于Milvus库的问答系统/integrated_qa_system/rag_qa/core/query_classifier.py

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

+ 1 - 1
4.基于Milvus库的问答系统/integrated_qa_system/rag_qa/core/rag_system.py

@@ -139,7 +139,7 @@ class RAGSystem:
         #   如果查询属于“通用知识”类别,则直接使用 LLM 回答
         if query_category == "通用知识":
             logger.info("查询为通用知识,直接调用 LLM")
-             prompt_input = self.rag_prompt.format(
+            prompt_input = self.rag_prompt.format(
                 context="", question=query, history="", phone=config.CUSTOMER_SERVICE_PHONE
             )  #   不使用上下文
             try:

+ 26 - 1
4.基于Milvus库的问答系统/integrated_qa_system/rag_qa/core/strategy_selector.py

@@ -6,7 +6,7 @@ from base.config import config
 from base.logger import logger
 # 导入 OpenAI
 from openai import OpenAI
-
+import json
 
 class StrategySelector:
     def __init__(self):
@@ -73,6 +73,17 @@ class StrategySelector:
                         * 策略:回溯问题检索
 
             根据用户查询 {query},直接返回最适合的策略名称,例如 "直接检索"。不要输出任何分析过程或其他内容。
+            # 返回结果必须严格遵循以下json格式:
+            # {
+            #     "strategy": "<策略名称>"
+            # }
+            # 例如:
+            # {
+            #     "strategy": "直接检索"
+            # }
+            # 注意:策略名称只能在 ["直接检索", "假设问题检索", "子查询检索", "回溯问题检索"] 中选择。
+            # 多一个字都不行,少一个字也不行,必须严格按照要求返回策略名称。
+            # 如果你违反了我的要求,返回的结果将被视为无效。
             """
             ,
             input_variables=["query"],
@@ -83,6 +94,20 @@ class StrategySelector:
         #   调用 LLM 获取检索策略
         strategy = self.call_dashscope(self.strategy_prompt_template.format(query=query)).strip()
         logger.info(f"为查询 '{query}' 选择的检索策略:{strategy}")
+        #   验证返回的策略是否在允许的列表中
+        valid_strategies = ["直接检索", "假设问题检索", "子查询检索", "回溯问题检索"]
+        json_response = None
+        try:
+            json_response = json.loads(strategy)
+            strategy = json_response.get("strategy", "直接检索")
+        except json.JSONDecodeError:
+            logger.error(f"解析策略选择结果失败,返回的内容不是有效的 JSON: {strategy}")
+            strategy = "直接检索"
+
+        if strategy not in valid_strategies:
+            logger.warning(f"返回的策略 '{strategy}' 不在允许的列表中,默认使用 '直接检索'")
+            strategy = "直接检索"
+
         return strategy
 
 if __name__ == '__main__':