From cddf324a5ea995234dd64e0838446d593bc56b5c Mon Sep 17 00:00:00 2001 From: hiyouga Date: Fri, 4 Aug 2023 14:20:27 +0800 Subject: [PATCH] fix webui export model Former-commit-id: d6b14658d4dec58fac3fbe1ec1d214e8422084bc --- src/llmtuner/webui/components/export.py | 1 + src/llmtuner/webui/utils.py | 9 ++++++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/src/llmtuner/webui/components/export.py b/src/llmtuner/webui/components/export.py index 295fc617..1b18fca0 100644 --- a/src/llmtuner/webui/components/export.py +++ b/src/llmtuner/webui/components/export.py @@ -22,6 +22,7 @@ def create_export_tab(top_elems: Dict[str, "Component"]) -> Dict[str, "Component top_elems["model_name"], top_elems["checkpoints"], top_elems["finetuning_type"], + top_elems["template"], max_shard_size, save_dir ], diff --git a/src/llmtuner/webui/utils.py b/src/llmtuner/webui/utils.py index 8a4f8d5a..7b667c0f 100644 --- a/src/llmtuner/webui/utils.py +++ b/src/llmtuner/webui/utils.py @@ -91,7 +91,13 @@ def gen_plot(base_model: str, finetuning_type: str, output_dir: str) -> matplotl 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]: if not model_name: yield ALERTS["err_no_model"][lang] @@ -118,6 +124,7 @@ def save_model( model_name_or_path=model_name_or_path, checkpoint_dir=checkpoint_dir, finetuning_type=finetuning_type, + template=template, output_dir=save_dir )