support export size setting

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."}

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

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
 

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)