feat: 支持添加训练好的模型到配置文件里

ChuanhuChatbot.py

@@ -15,8 +15,8 @@ from modules.presets import *
 from modules.overwrites import *
 from modules.webui import *
 from modules.repo import *
+from modules.train_func import *
 from modules.models.models import get_model
-from modules.train_func import handle_dataset_selection, handle_dataset_clear, upload_to_openai, start_training, get_training_status, add_to_models, cancel_all_jobs
 
 logging.getLogger("httpx").setLevel(logging.WARNING)
 

modules/config.py

@@ -96,6 +96,10 @@ else:
     sensitive_id = config.get("sensitive_id", "")
     sensitive_id = os.environ.get("SENSITIVE_ID", sensitive_id)
 
+# 模型配置
+if "extra_models" in  config:
+    presets.MODELS.extend(config["extra_models"])
+    logging.info(f"已添加额外的模型：{config['extra_models']}")
 
 google_palm_api_key = config.get("google_palm_api_key", "")
 google_palm_api_key = os.environ.get(

modules/train_func.py

@@ -5,6 +5,7 @@ import traceback
 import openai
 import gradio as gr
 import ujson as json
+import commentjson
 
 import modules.presets as presets
 from modules.utils import get_file_hash
@@ -112,7 +113,18 @@ def handle_dataset_clear():
 def add_to_models():
     openai.api_key = os.getenv("OPENAI_API_KEY")
     succeeded_jobs = [job for job in openai.FineTuningJob.list()["data"] if job["status"] == "succeeded"]
-    presets.MODELS.extend([job["fine_tuned_model"] for job in succeeded_jobs])
+    extra_models = [job["fine_tuned_model"] for job in succeeded_jobs]
+    presets.MODELS.extend(extra_models)
+
+    with open('config.json', 'r') as f:
+        data = commentjson.load(f)
+    if 'extra_models' in data:
+        data['extra_models'].extend(extra_models)
+    else:
+        data['extra_models'] = extra_models
+    with open('config.json', 'w') as f:
+        commentjson.dump(data, f, indent=4)
+
     return gr.update(choices=presets.MODELS), f"成功添加了 {len(succeeded_jobs)} 个模型。"
 
 def cancel_all_jobs():