webui.py 3.65 KB
Newer Older
1
import argparse
2
3
4
5
import json
import os
import requests

6
7
8
9
10
11
12
13
14
import gradio as gr

from utils import DocAction

def parseArgs():
    parser = argparse.ArgumentParser()
    parser.add_argument("--http_host", default="0.0.0.0")
    parser.add_argument("--http_port", type=int, default=13666)
    return parser.parse_args()
15

16
def get_response(data, url):
17
    headers = {"Content-type": "application/json"}
18
    response = requests.post(url, json=data, headers=headers)
19
20
21
22
23
24
25
    response = json.loads(response.content)
    return response

def add_text(history, text):
    history = history + [(text, None)]
    return history, gr.update(value=None, interactive=True)

26

27
def add_file(history, files):
28
29
30
31
32
33
34
35
36
    files_string = "\n".join([os.path.basename(file.name) for file in files])

    doc_files = [file.name for file in files]
    data = {
        "doc_files": doc_files,
        "action": DocAction.ADD
    }
    response = get_response(data, update_url)["response"]
    history = history + [(files_string, response)]
37
38
    return history

39
40
41
42
43
44
def bot(history):    
    data = {
        "user_input": history[-1][0].strip()
    }
    response = get_response(data, gen_url)

45
46
47
48
49
50
51
    if response["error"] != "":
        raise gr.Error(response["error"])
    
    history[-1][1] = response["response"]
    yield history


52
53
54
55
56
57
58
59
60
61
62
def restart(chatbot, txt):
    # Reset the conversation state and clear the chat history
    data = {
        "doc_files": "",
        "action": DocAction.CLEAR
    }
    response = get_response(data, update_url)
    
    return gr.update(value=None), gr.update(value=None, interactive=True)


63
64
65
66
67
68
69
70
71
CSS = """
.contain { display: flex; flex-direction: column; height: 100vh }
#component-0 { height: 100%; }
#chatbot { flex-grow: 1; }
"""

header_html = """
<div style="background: linear-gradient(to right, #2a0cf4, #7100ed, #9800e6, #b600df, #ce00d9, #dc0cd1, #e81bca, #f229c3, #f738ba, #f946b2, #fb53ab, #fb5fa5); padding: 20px; text-align: left;">
    <h1 style="color: white;">ColossalQA</h1>
72
    <h4 style="color: white;">A powerful Q&A system with knowledge bases</h4>
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
</div>
"""

with gr.Blocks(css=CSS) as demo:
    html = gr.HTML(header_html)
    chatbot = gr.Chatbot(
        [],
        elem_id="chatbot",
        bubble_full_width=False,
        avatar_images=(
            (os.path.join(os.path.dirname(__file__), "img/avatar_user.png")),
            (os.path.join(os.path.dirname(__file__), "img/avatar_ai.png")),
        ),
    )
    with gr.Row():
88
89
        btn = gr.UploadButton("📁", file_types=["file"], file_count="multiple", size="sm")
        restart_btn = gr.Button(str("\u21BB"), elem_id="restart-btn", scale=1)
90
        txt = gr.Textbox(
91
            scale=8,
92
            show_label=False,
93
            placeholder="Enter text and press enter, or use 📁 to upload files, click \u21BB to clear loaded files and restart chat",
94
95
96
97
98
99
100
101
            container=True,
            autofocus=True,
        )

    txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(bot, chatbot, chatbot)
    # Clear the original textbox
    txt_msg.then(lambda: gr.update(value=None, interactive=True), None, [txt], queue=False) 
    # Click Upload Button: 1. upload files  2. send config to backend, initalize model 3. get response "conversation_ready" = True/False
102
    file_msg = btn.upload(add_file, [chatbot, btn], [chatbot], queue=False)
103

104
105
    # restart
    restart_msg = restart_btn.click(restart, [chatbot, txt], [chatbot, txt], queue=False)
106
107
108


if __name__ == "__main__":
109
110
111
112
113
    args = parseArgs()

    update_url = f"http://{args.http_host}:{args.http_port}/update"
    gen_url = f"http://{args.http_host}:{args.http_port}/generate"

114
115
    demo.queue()
    demo.launch(share=True)  # share=True will release a public link of the demo