7a4388c928
在模型管理页面中新增保存模式选择功能，用户可以通过下拉菜单选择不同的保存模式（如默认、合并16位、合并4位等）。同时，将保存模型的逻辑抽离到独立的`save_model.py`文件中，以提高代码的可维护性和复用性。
48 lines
2.1 KiB
Python
48 lines
2.1 KiB
Python
import os
|
|
from global_var import get_model, get_tokenizer
|
|
|
|
def save_model_to_dir(save_model_name, models_dir, model, tokenizer, save_method="default"):
|
|
"""
|
|
保存模型到指定目录
|
|
:param save_model_name: 要保存的模型名称
|
|
:param models_dir: 模型保存的基础目录
|
|
:param model: 要保存的模型
|
|
:param tokenizer: 要保存的tokenizer
|
|
:param save_method: 保存模式选项
|
|
- "default": 默认保存方式
|
|
- "merged_16bit": 合并为16位
|
|
- "merged_4bit": 合并为4位
|
|
- "lora": 仅LoRA适配器
|
|
- "gguf": 保存为GGUF格式
|
|
- "gguf_q4_k_m": 保存为q4_k_m GGUF格式
|
|
- "gguf_f16": 保存为16位GGUF格式
|
|
:return: 保存结果消息或错误信息
|
|
"""
|
|
try:
|
|
if model is None:
|
|
return "没有加载的模型可保存"
|
|
|
|
save_path = os.path.join(models_dir, save_model_name)
|
|
os.makedirs(save_path, exist_ok=True)
|
|
|
|
if save_method == "default":
|
|
model.save_pretrained(save_path)
|
|
tokenizer.save_pretrained(save_path)
|
|
elif save_method == "merged_16bit":
|
|
model.save_pretrained_merged(save_path, tokenizer, save_method="merged_16bit")
|
|
elif save_method == "merged_4bit":
|
|
model.save_pretrained_merged(save_path, tokenizer, save_method="merged_4bit_forced")
|
|
elif save_method == "lora":
|
|
model.save_pretrained_merged(save_path, tokenizer, save_method="lora")
|
|
elif save_method == "gguf":
|
|
model.save_pretrained_gguf(save_path, tokenizer)
|
|
elif save_method == "gguf_q4_k_m":
|
|
model.save_pretrained_gguf(save_path, tokenizer, quantization_method="q4_k_m")
|
|
elif save_method == "gguf_f16":
|
|
model.save_pretrained_gguf(save_path, tokenizer, quantization_method="f16")
|
|
else:
|
|
return f"不支持的保存模式: {save_method}"
|
|
|
|
return f"模型已保存到 {save_path} (模式: {save_method})"
|
|
except Exception as e:
|
|
return f"保存模型时出错: {str(e)}" |