44 lines
1.1 KiB
Python
44 lines
1.1 KiB
Python
import gradio as gr
|
|
from typing import TYPE_CHECKING, Dict
|
|
|
|
from llmtuner.webui.utils import save_model
|
|
|
|
if TYPE_CHECKING:
|
|
from gradio.components import Component
|
|
from llmtuner.webui.engine import Engine
|
|
|
|
|
|
def create_export_tab(engine: "Engine") -> Dict[str, "Component"]:
|
|
elem_dict = dict()
|
|
|
|
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(
|
|
save_model,
|
|
[
|
|
engine.manager.get_elem("top.lang"),
|
|
engine.manager.get_elem("top.model_name"),
|
|
engine.manager.get_elem("top.model_path"),
|
|
engine.manager.get_elem("top.checkpoints"),
|
|
engine.manager.get_elem("top.finetuning_type"),
|
|
engine.manager.get_elem("top.template"),
|
|
max_shard_size,
|
|
save_dir
|
|
],
|
|
[info_box]
|
|
)
|
|
|
|
elem_dict.update(dict(
|
|
save_dir=save_dir,
|
|
max_shard_size=max_shard_size,
|
|
export_btn=export_btn,
|
|
info_box=info_box
|
|
))
|
|
|
|
return elem_dict
|