-
Notifications
You must be signed in to change notification settings - Fork 0
/
web_demo_gradio.py
127 lines (109 loc) · 4.65 KB
/
web_demo_gradio.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
"""A simple web interactive chat demo based on gradio."""
import gradio as gr
from stylellm_chat import StyleLLMChat, QwenChatModel, StyleModel
def parse_text(text):
lines = text.split("\n")
lines = [line for line in lines if line != ""]
count = 0
for i, line in enumerate(lines):
if "```" in line:
count += 1
items = line.split("`")
if count % 2 == 1:
lines[i] = f'<pre><code class="language-{items[-1]}">'
else:
lines[i] = f"<br></code></pre>"
else:
if i > 0:
if count % 2 == 1:
line = line.replace("`", "\`")
line = line.replace("<", "<")
line = line.replace(">", ">")
line = line.replace(" ", " ")
line = line.replace("*", "*")
line = line.replace("_", "_")
line = line.replace("-", "-")
line = line.replace(".", ".")
line = line.replace("!", "!")
line = line.replace("(", "(")
line = line.replace(")", ")")
line = line.replace("$", "$")
lines[i] = "<br>" + line
text = "".join(lines)
return text
def run_app(sc: StyleLLMChat):
sc.clear_history()
with gr.Blocks() as demo:
gr.Markdown(
'<center><font size=6><a href="https://github.com/stylellm/stylellm_chat">StyleLLM-Chat</a></center>'
)
gr.Markdown(
'<center><font size=2>对话模型:%s 角色:%s</center>' % (
sc.chat_model.model_name_or_path, sc.chat_model.character)
)
gr.Markdown(
'<center><font size=2>风格模型:%s 角色:%s</center>' % (
sc.style_model.model_name_or_path, sc.style_model.character)
)
chatbot = gr.Chatbot()
with gr.Row():
with gr.Column(scale=4):
with gr.Column(scale=12):
user_input = gr.Textbox(
show_label=False,
placeholder="Input...",
lines=10,
container=False,
)
with gr.Column(min_width=32, scale=1):
submit_btn = gr.Button("🚀 Submit")
with gr.Column(scale=1):
empty_btn = gr.Button("🧹 Clear History")
max_new_tokens = gr.Slider(
0,
32768,
value=1024,
step=1.0,
label="Maximum length",
interactive=True,
)
top_p = gr.Slider(
0, 1, value=0.8, step=0.01, label="Top P", interactive=True
)
temperature = gr.Slider(
0.01, 1, value=0.6, step=0.01, label="Temperature", interactive=True
)
def user(query, dialog):
return "", dialog + [[parse_text(query), ""]]
def predict(history, _max_new_tokens, _top_p, _temperature):
messages = []
for idx, (user_msg, model_msg) in enumerate(history):
if idx == len(history) - 1 and not model_msg:
messages.append({"role": "user", "content": user_msg})
break
if user_msg:
messages.append({"role": "user", "content": user_msg})
if model_msg:
messages.append({"role": "assistant", "content": model_msg})
out = sc.generate(messages, max_new_tokens=_max_new_tokens, top_p=_top_p, temperature=_temperature)
history[-1][-1] = out
return history
submit_btn.click(
user, [user_input, chatbot], [user_input, chatbot], queue=False
).then(predict, [chatbot, max_new_tokens, top_p, temperature], chatbot)
user_input.submit(
user, [user_input, chatbot], [user_input, chatbot], queue=False
).then(predict, [chatbot, max_new_tokens, top_p, temperature], chatbot)
empty_btn.click(lambda: None, inputs=None, outputs=chatbot, queue=False)
demo.queue()
demo.launch(share=True)
return demo
def main():
sc = StyleLLMChat(
chat_model=QwenChatModel(model_name_or_path="Qwen/Qwen1.5-7B-Chat-AWQ", device="cuda"),
style_model=StyleModel(model_name_or_path="stylellm/ShuiHuZhuan-6b-AWQ", device="cuda",
scene="水浒传", character="鲁智深")
)
run_app(sc)
if __name__ == '__main__':
main()