From 5bc280c539f3af620cd78b7e30ee43a8fb05f2db Mon Sep 17 00:00:00 2001
From: 997146918 <997146918@qq.com>
Date: Sat, 23 Aug 2025 18:20:50 +0800
Subject: [PATCH] =?UTF-8?q?=E5=AE=8C=E5=96=84=E6=89=93=E5=88=86=E6=95=B0?=
 =?UTF-8?q?=E6=8D=AE=E9=9B=86=E7=94=9F=E6=88=90?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 AITrain/main_controller.py | 26 ++++++++++----------------
 1 file changed, 10 insertions(+), 16 deletions(-)

diff --git a/AITrain/main_controller.py b/AITrain/main_controller.py
index 90f0770..b3b578a 100644
--- a/AITrain/main_controller.py
+++ b/AITrain/main_controller.py
@@ -282,7 +282,7 @@ def generate_training_dataset():
         output_dir = "./training_data"
         os.makedirs(output_dir, exist_ok=True)
         
-        print("请选择训练数据生成类型:")
+        print("请选择训练数据生成类型:（默认1）")
         print("1. 高质量对话数据集 (分数≥8.0)")
         print("2. 问题对话改进数据集 (分数<6.0)")
         print("3. 角色一致性训练集")
@@ -290,11 +290,11 @@ def generate_training_dataset():
         print("5. 完整对话质量数据集")
         
         choice = input("请输入选择 (1-5): ").strip()
-        
+        choice = int(choice) if choice.isdigit() else 1
         with sqlite3.connect(conv_mgr.db_path) as conn:
             training_data = []
             
-            if choice == '1':
+            if choice == 1:
                 # 高质量对话数据集
                 print("\n生成高质量对话数据集...")
                 cursor = conn.execute("""
@@ -317,7 +317,7 @@ def generate_training_dataset():
                 
                 output_file = f"{output_dir}/high_quality_dialogues_{datetime.now().strftime('%Y%m%d_%H%M')}.json"
                 
-            elif choice == '2':
+            elif choice == 2:
                 # 问题对话改进数据集
                 print("\n生成问题对话改进数据集...")
                 cursor = conn.execute("""
@@ -344,7 +344,7 @@ def generate_training_dataset():
                 
                 output_file = f"{output_dir}/improvement_dialogues_{datetime.now().strftime('%Y%m%d_%H%M')}.json"
                 
-            elif choice == '3':
+            elif choice == 3:
                 # 角色一致性训练集
                 print("\n生成角色一致性训练集...")
                 cursor = conn.execute("""
@@ -369,7 +369,7 @@ def generate_training_dataset():
                 
                 output_file = f"{output_dir}/character_consistency_{datetime.now().strftime('%Y%m%d_%H%M')}.json"
                 
-            elif choice == '4':
+            elif choice == 4:
                 # 创意性增强训练集
                 print("\n生成创意性增强训练集...")
                 cursor = conn.execute("""
@@ -394,7 +394,7 @@ def generate_training_dataset():
                 
                 output_file = f"{output_dir}/creativity_enhancement_{datetime.now().strftime('%Y%m%d_%H%M')}.json"
                 
-            elif choice == '5':
+            elif choice == 5:
                 # 完整对话质量数据集
                 print("\n生成完整对话质量数据集...")
                 cursor = conn.execute("""
@@ -1107,10 +1107,8 @@ def main():
         print("5. 启动双AI对话系统 (开启人工打分)")
         print("6. 系统状态检查")
         print("7. 查看对话评分统计")
-        print("8. 模型性能分析与优化")
-        print("9. 生成训练数据集")
-        print("10. 模型迭代优化")
-        print("11. 查看使用说明")
+        print("8. 生成训练数据集")
+        print("9. 模型迭代优化")
         print("0. 退出")
         print("="*50)
         
@@ -1142,13 +1140,9 @@ def main():
             show_scoring_statistics()
             
         elif choice == '8':
-            # 模型性能分析与优化 - 待实现
-            print("模型性能分析与优化功能开发中...")
-            
-        elif choice == '9':
             generate_training_dataset()
             
-        elif choice == '10':
+        elif choice == '9':
             run_model_optimization()
             
         elif choice == '11':