update web UI, support rm predict #210

Former-commit-id: 92cc6b655dc91b94d5bf9d8618c3b57d5cf94333

src/llmtuner/webui/components/__init__.py

@@ -1,4 +1,5 @@
-from llmtuner.webui.components.eval import create_eval_tab
-from llmtuner.webui.components.infer import create_infer_tab
 from llmtuner.webui.components.top import create_top
 from llmtuner.webui.components.sft import create_sft_tab
+from llmtuner.webui.components.eval import create_eval_tab
+from llmtuner.webui.components.infer import create_infer_tab
+from llmtuner.webui.components.export import create_export_tab

src/llmtuner/webui/components/chatbot.py

@@ -22,13 +22,9 @@ def create_chat_box(
 
             with gr.Column(scale=1):
                 clear_btn = gr.Button()
-                max_new_tokens = gr.Slider(
-                    10, 2048, value=chat_model.generating_args.max_new_tokens, step=1, interactive=True
-                )
-                top_p = gr.Slider(0.01, 1, value=chat_model.generating_args.top_p, step=0.01, interactive=True)
-                temperature = gr.Slider(
-                    0.01, 1.5, value=chat_model.generating_args.temperature, step=0.01, interactive=True
-                )
+                max_new_tokens = gr.Slider(10, 2048, value=chat_model.generating_args.max_new_tokens, step=1)
+                top_p = gr.Slider(0.01, 1, value=chat_model.generating_args.top_p, step=0.01)
+                temperature = gr.Slider(0.01, 1.5, value=chat_model.generating_args.temperature, step=0.01)
 
     history = gr.State([])
 

src/llmtuner/webui/components/export.py

@@ -0,0 +1,34 @@
+from typing import Dict
+import gradio as gr
+from gradio.components import Component
+
+from llmtuner.webui.utils import export_model
+
+
+def create_export_tab(top_elems: Dict[str, Component]) -> Dict[str, Component]:
+    with gr.Row():
+        save_dir = gr.Textbox()
+        max_shard_size = gr.Slider(value=10, minimum=1, maximum=100)
+
+    export_btn = gr.Button()
+    info_box = gr.Textbox(show_label=False, interactive=False)
+
+    export_btn.click(
+        export_model,
+        [
+            top_elems["lang"],
+            top_elems["model_name"],
+            top_elems["checkpoints"],
+            top_elems["finetuning_type"],
+            max_shard_size,
+            save_dir
+        ],
+        [info_box]
+    )
+
+    return dict(
+        save_dir=save_dir,
+        max_shard_size=max_shard_size,
+        export_btn=export_btn,
+        info_box=info_box
+    )

src/llmtuner/webui/components/sft.py

@@ -57,7 +57,7 @@ def create_sft_tab(top_elems: Dict[str, Component], runner: Runner) -> Dict[str,
 
     with gr.Row():
         with gr.Column(scale=4):
-            output_dir = gr.Textbox(interactive=True)
+            output_dir = gr.Textbox()
 
             with gr.Box():
                 output_box = gr.Markdown()