aurora

1
import gradio as gr
2
from typing import TYPE_CHECKING, Dict
3

4
from llmtuner.webui.components.chatbot import create_chat_box
5

6
if TYPE_CHECKING:
7
    from gradio.components import Component
8
    from llmtuner.webui.engine import Engine
9

10

11
def create_infer_tab(engine: "Engine") -> Dict[str, "Component"]:
12
    input_elems = engine.manager.get_base_elems()
13
    elem_dict = dict()
14

15
    with gr.Row():
16
        load_btn = gr.Button()
17
        unload_btn = gr.Button()
18

19
    info_box = gr.Textbox(show_label=False, interactive=False)
20
    elem_dict.update(dict(load_btn=load_btn, unload_btn=unload_btn, info_box=info_box))
21

22
    chat_box, chatbot, history, chat_elems = create_chat_box(engine, visible=False)
23
    elem_dict.update(dict(chat_box=chat_box, **chat_elems))
24

25
    load_btn.click(
26
        engine.chatter.load_model, input_elems, [info_box]
27
    ).then(
28
        lambda: gr.update(visible=engine.chatter.loaded), outputs=[chat_box]
29
    )
30

31
    unload_btn.click(
32
        engine.chatter.unload_model, input_elems, [info_box]
33
    ).then(
34
        lambda: ([], []), outputs=[chatbot, history]
35
    ).then(
36
        lambda: gr.update(visible=engine.chatter.loaded), outputs=[chat_box]
37
    )
38

39
    return elem_dict
40