featmodel): 添加保存模式选择功能
在模型管理页面中新增保存模式选择功能,用户可以通过下拉菜单选择不同的保存模式(如默认、合并16位、合并4位等)。同时,将保存模型的逻辑抽离到独立的`save_model.py`文件中,以提高代码的可维护性和复用性。
This commit is contained in:
@@ -30,6 +30,11 @@ def model_manage_page():
|
||||
with gr.Row():
|
||||
with gr.Column(scale=3):
|
||||
save_model_name_input = gr.Textbox(label="保存模型名称", placeholder="输入模型保存名称")
|
||||
save_method_dropdown = gr.Dropdown(
|
||||
choices=["default", "merged_16bit", "merged_4bit", "lora", "gguf", "gguf_q4_k_m", "gguf_f16"],
|
||||
label="保存模式",
|
||||
value="default"
|
||||
)
|
||||
with gr.Column(scale=1):
|
||||
save_button = gr.Button("保存模型", variant="secondary")
|
||||
|
||||
@@ -73,21 +78,12 @@ def model_manage_page():
|
||||
|
||||
unload_button.click(fn=unload_model, inputs=None, outputs=state_output)
|
||||
|
||||
def save_model(save_model_name):
|
||||
try:
|
||||
global model, tokenizer
|
||||
if model is None:
|
||||
return "没有加载的模型可保存"
|
||||
|
||||
save_path = os.path.join(models_dir, save_model_name)
|
||||
os.makedirs(save_path, exist_ok=True)
|
||||
model.save_pretrained(save_path)
|
||||
tokenizer.save_pretrained(save_path)
|
||||
return f"模型已保存到 {save_path}"
|
||||
except Exception as e:
|
||||
return f"保存模型时出错: {str(e)}"
|
||||
from train.save_model import save_model_to_dir
|
||||
|
||||
def save_model(save_model_name, save_method):
|
||||
return save_model_to_dir(save_model_name, models_dir, get_model(), get_tokenizer(), save_method)
|
||||
|
||||
save_button.click(fn=save_model, inputs=save_model_name_input, outputs=state_output)
|
||||
save_button.click(fn=save_model, inputs=[save_model_name_input, save_method_dropdown], outputs=state_output)
|
||||
|
||||
def refresh_model_list():
|
||||
try:
|
||||
|
Reference in New Issue
Block a user