fix webui export model

Former-commit-id: d6b14658d4dec58fac3fbe1ec1d214e8422084bc
This commit is contained in:
hiyouga 2023-08-04 14:20:27 +08:00
parent 0328c0e07c
commit cddf324a5e
2 changed files with 9 additions and 1 deletions

View File

@ -22,6 +22,7 @@ def create_export_tab(top_elems: Dict[str, "Component"]) -> Dict[str, "Component
top_elems["model_name"], top_elems["model_name"],
top_elems["checkpoints"], top_elems["checkpoints"],
top_elems["finetuning_type"], top_elems["finetuning_type"],
top_elems["template"],
max_shard_size, max_shard_size,
save_dir save_dir
], ],

View File

@ -91,7 +91,13 @@ def gen_plot(base_model: str, finetuning_type: str, output_dir: str) -> matplotl
def save_model( def save_model(
lang: str, model_name: str, checkpoints: List[str], finetuning_type: str, max_shard_size: int, save_dir: str lang: str,
model_name: str,
checkpoints: List[str],
finetuning_type: str,
template: str,
max_shard_size: int,
save_dir: str
) -> Generator[str, None, None]: ) -> Generator[str, None, None]:
if not model_name: if not model_name:
yield ALERTS["err_no_model"][lang] yield ALERTS["err_no_model"][lang]
@ -118,6 +124,7 @@ def save_model(
model_name_or_path=model_name_or_path, model_name_or_path=model_name_or_path,
checkpoint_dir=checkpoint_dir, checkpoint_dir=checkpoint_dir,
finetuning_type=finetuning_type, finetuning_type=finetuning_type,
template=template,
output_dir=save_dir output_dir=save_dir
) )