diff --git a/assets/wechat.jpg b/assets/wechat.jpg index c8662c40..f7cb3a8e 100644 Binary files a/assets/wechat.jpg and b/assets/wechat.jpg differ diff --git a/src/llmtuner/hparams/finetuning_args.py b/src/llmtuner/hparams/finetuning_args.py index 5bdeba96..e4c4a37b 100644 --- a/src/llmtuner/hparams/finetuning_args.py +++ b/src/llmtuner/hparams/finetuning_args.py @@ -149,6 +149,10 @@ class FinetuningArguments(FreezeArguments, LoraArguments, RLHFArguments): default=None, metadata={"help": "Path to the directory to save the exported model."} ) + export_size: Optional[int] = field( + default=1, + metadata={"help": "The file shard size (in GB) of the exported model."} + ) plot_loss: Optional[bool] = field( default=False, metadata={"help": "Whether to plot the training loss after fine-tuning or not."} diff --git a/src/llmtuner/train/tuner.py b/src/llmtuner/train/tuner.py index 5833ee9e..a6e8b2d1 100644 --- a/src/llmtuner/train/tuner.py +++ b/src/llmtuner/train/tuner.py @@ -34,7 +34,7 @@ def run_exp(args: Optional[Dict[str, Any]] = None, callbacks: Optional[List["Tra raise ValueError("Unknown task.") -def export_model(args: Optional[Dict[str, Any]] = None, max_shard_size: Optional[str] = "10GB"): +def export_model(args: Optional[Dict[str, Any]] = None): model_args, _, finetuning_args, _ = get_infer_args(args) model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args) @@ -42,7 +42,7 @@ def export_model(args: Optional[Dict[str, Any]] = None, max_shard_size: Optional raise ValueError("Cannot export a GPTQ quantized model.") model.config.use_cache = True - model.save_pretrained(finetuning_args.export_dir, max_shard_size=max_shard_size) + model.save_pretrained(finetuning_args.export_dir, max_shard_size="{}GB".format(finetuning_args.export_size)) try: tokenizer.padding_side = "left" # restore padding side diff --git a/src/llmtuner/webui/chatter.py b/src/llmtuner/webui/chatter.py index d591a2e3..ddf80b27 100644 --- a/src/llmtuner/webui/chatter.py +++ b/src/llmtuner/webui/chatter.py @@ -90,6 +90,7 @@ class WebChatModel(ChatModel): lang = data[self.manager.get_elem_by_name("top.lang")] if self.demo_mode: + gr.Warning(ALERTS["err_demo"][lang]) yield ALERTS["err_demo"][lang] return diff --git a/src/llmtuner/webui/components/export.py b/src/llmtuner/webui/components/export.py index b78dc831..6ac6f3e6 100644 --- a/src/llmtuner/webui/components/export.py +++ b/src/llmtuner/webui/components/export.py @@ -40,18 +40,19 @@ def save_model( checkpoint_dir=",".join([get_save_dir(model_name, finetuning_type, ckpt) for ckpt in checkpoints]), finetuning_type=finetuning_type, template=template, - export_dir=export_dir + export_dir=export_dir, + export_size=max_shard_size ) yield ALERTS["info_exporting"][lang] - export_model(args, max_shard_size="{}GB".format(max_shard_size)) + export_model(args) yield ALERTS["info_exported"][lang] def create_export_tab(engine: "Engine") -> Dict[str, "Component"]: with gr.Row(): export_dir = gr.Textbox() - max_shard_size = gr.Slider(value=10, minimum=1, maximum=100) + max_shard_size = gr.Slider(value=1, minimum=1, maximum=100) export_btn = gr.Button() info_box = gr.Textbox(show_label=False, interactive=False)