Tuchuanhuhuhu commited on
Commit
b5ddb7e
·
1 Parent(s): 9813f91

feat: 添加取消所有任务按钮

Browse files
Files changed (2) hide show
  1. ChuanhuChatbot.py +3 -1
  2. modules/train_func.py +9 -2
ChuanhuChatbot.py CHANGED
@@ -16,7 +16,7 @@ from modules.overwrites import *
16
  from modules.webui import *
17
  from modules.repo import *
18
  from modules.models.models import get_model
19
- from modules.train_func import handle_dataset_selection, handle_dataset_clear, upload_to_openai, start_training, get_training_status, add_to_models
20
 
21
  logging.getLogger("httpx").setLevel(logging.WARNING)
22
 
@@ -201,6 +201,7 @@ with gr.Blocks(theme=small_and_beautiful_theme) as demo:
201
  with gr.Column(variant="panel"):
202
  openai_train_status = gr.Markdown(label=i18n("训练状态"), value=i18n("未开始训练"))
203
  openai_status_refresh_btn = gr.Button(i18n("刷新状态"))
 
204
  add_to_models_btn = gr.Button(i18n("添加训练好的模型到模型列表"), interactive=False)
205
 
206
  with gr.Tab(label=i18n("高级")):
@@ -498,6 +499,7 @@ with gr.Blocks(theme=small_and_beautiful_theme) as demo:
498
  openai_start_train_btn.click(start_training, [openai_ft_file_id, openai_ft_suffix, openai_train_epoch_slider], [openai_train_status])
499
  openai_status_refresh_btn.click(get_training_status, [], [openai_train_status, add_to_models_btn])
500
  add_to_models_btn.click(add_to_models, [], [model_select_dropdown, status_display], show_progress=True)
 
501
 
502
  # Advanced
503
  max_context_length_slider.change(set_token_upper_limit, [current_model, max_context_length_slider], None)
 
16
  from modules.webui import *
17
  from modules.repo import *
18
  from modules.models.models import get_model
19
+ 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
20
 
21
  logging.getLogger("httpx").setLevel(logging.WARNING)
22
 
 
201
  with gr.Column(variant="panel"):
202
  openai_train_status = gr.Markdown(label=i18n("训练状态"), value=i18n("未开始训练"))
203
  openai_status_refresh_btn = gr.Button(i18n("刷新状态"))
204
+ openai_cancel_all_jobs_btn = gr.Button(i18n("取消所有任务"))
205
  add_to_models_btn = gr.Button(i18n("添加训练好的模型到模型列表"), interactive=False)
206
 
207
  with gr.Tab(label=i18n("高级")):
 
499
  openai_start_train_btn.click(start_training, [openai_ft_file_id, openai_ft_suffix, openai_train_epoch_slider], [openai_train_status])
500
  openai_status_refresh_btn.click(get_training_status, [], [openai_train_status, add_to_models_btn])
501
  add_to_models_btn.click(add_to_models, [], [model_select_dropdown, status_display], show_progress=True)
502
+ openai_cancel_all_jobs_btn.click(cancel_all_jobs, [], [openai_train_status], show_progress=True)
503
 
504
  # Advanced
505
  max_context_length_slider.change(set_token_upper_limit, [current_model, max_context_length_slider], None)
modules/train_func.py CHANGED
@@ -111,6 +111,13 @@ def handle_dataset_clear():
111
 
112
  def add_to_models():
113
  openai.api_key = os.getenv("OPENAI_API_KEY")
114
- succeeded_jobs = [job for job in openai.FineTuningJob.list(limit=10)["data"] if job["status"] == "succeeded"]
115
  presets.MODELS.extend([job["fine_tuned_model"] for job in succeeded_jobs])
116
- return gr.update(choices=presets.MODELS), f"成功添加了 {len(succeeded_jobs)} 个模型。"
 
 
 
 
 
 
 
 
111
 
112
  def add_to_models():
113
  openai.api_key = os.getenv("OPENAI_API_KEY")
114
+ succeeded_jobs = [job for job in openai.FineTuningJob.list()["data"] if job["status"] == "succeeded"]
115
  presets.MODELS.extend([job["fine_tuned_model"] for job in succeeded_jobs])
116
+ return gr.update(choices=presets.MODELS), f"成功添加了 {len(succeeded_jobs)} 个模型。"
117
+
118
+ def cancel_all_jobs():
119
+ openai.api_key = os.getenv("OPENAI_API_KEY")
120
+ jobs = [job for job in openai.FineTuningJob.list()["data"] if job["status"] not in ["cancelled", "succeeded"]]
121
+ for job in jobs:
122
+ openai.FineTuningJob.cancel(job["id"])
123
+ return f"成功取消了 {len(jobs)} 个训练任务。"