wehub-resource-sync
355 lines
14 KiB
Python
355 lines
14 KiB
Python
#!/usr/bin/env python
|
|
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import re
|
|
|
|
import gradio as gr
|
|
import requests
|
|
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(logging.DEBUG)
|
|
console_handler = logging.StreamHandler()
|
|
console_handler.setLevel(logging.INFO)
|
|
formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
|
|
console_handler.setFormatter(formatter)
|
|
logger.addHandler(console_handler)
|
|
|
|
|
|
def setup_args():
|
|
"""Setup arguments."""
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--port", type=int, default=8073)
|
|
parser.add_argument("--api_key", type=str, default=None, help="Your API key")
|
|
parser.add_argument("--model", type=str, default="", help="Model name")
|
|
parser.add_argument("--title", type=str, default="PaddleNLP Chat", help="UI Title")
|
|
parser.add_argument("--sub_title", type=str, default="powered by paddlenlp team.", help="UI Sub Title")
|
|
parser.add_argument("--flask_port", type=int, default=None, help="The port of flask service")
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def create_src_slider(value, maximum):
|
|
return gr.Slider(
|
|
minimum=1,
|
|
maximum=maximum,
|
|
value=value,
|
|
step=1,
|
|
label="Max Src Length",
|
|
info="最大输入长度。",
|
|
)
|
|
|
|
|
|
def create_max_slider(value, maximum):
|
|
return gr.Slider(
|
|
minimum=1,
|
|
maximum=maximum,
|
|
value=value,
|
|
step=1,
|
|
label="Max Decoding Length",
|
|
info="生成结果的最大长度。",
|
|
)
|
|
|
|
|
|
def remove_think_tags(text):
|
|
"""
|
|
清除文本中 <think> 和 </think> 标签之间的所有字符。
|
|
|
|
Args:
|
|
text: 要处理的文本字符串。
|
|
|
|
Returns:
|
|
清除 <think> 和 </think> 标签之间内容的文本字符串。
|
|
"""
|
|
pattern = re.compile(r"\\<think\\>.*?\\<\\\/think\\>", re.DOTALL)
|
|
# 将匹配到的部分替换为空字符串
|
|
cleaned_text = pattern.sub("", text).strip()
|
|
return cleaned_text
|
|
|
|
|
|
def launch(args, default_params: dict = {}):
|
|
"""Launch chat UI with OpenAI API."""
|
|
|
|
def rollback(state):
|
|
"""Rollback context."""
|
|
context = state.setdefault("context", [])
|
|
# 回退时移除最后一次对话
|
|
if len(context) >= 2:
|
|
content = context[-2]["content"]
|
|
context = context[:-2]
|
|
state["context"] = context
|
|
shown_context = get_shown_context(context)
|
|
return content, shown_context, context, state
|
|
else:
|
|
gr.Warning("没有可撤回的对话历史")
|
|
return None, get_shown_context(context), context, state
|
|
|
|
def regen(state, top_k, top_p, temperature, repetition_penalty, max_tokens, src_length):
|
|
"""Regenerate response."""
|
|
context = state.setdefault("context", [])
|
|
if len(context) < 2:
|
|
gr.Warning("No chat history!")
|
|
shown_context = get_shown_context(context)
|
|
return None, shown_context, context, state
|
|
|
|
# 删除上一次回复，重新生成
|
|
context.pop()
|
|
user_turn = context.pop()
|
|
context.append({"role": "user", "content": user_turn["content"]})
|
|
context.append({"role": "assistant", "content": ""})
|
|
shown_context = get_shown_context(context)
|
|
return user_turn["content"], shown_context, context, state
|
|
|
|
def begin(content, state):
|
|
"""记录用户输入，并初始化 bot 回复为空。"""
|
|
context = state.setdefault("context", [])
|
|
|
|
if not content:
|
|
gr.Warning("Invalid inputs")
|
|
shown_context = get_shown_context(context)
|
|
return None, shown_context, context, state
|
|
|
|
context.append({"role": "user", "content": content})
|
|
context.append({"role": "assistant", "content": ""})
|
|
shown_context = get_shown_context(context)
|
|
return content, shown_context, context, state
|
|
|
|
def infer(content, state, top_k, top_p, temperature, repetition_penalty, max_tokens, src_length):
|
|
"""调用 OpenAI 接口生成回答，并以流式返回部分结果。"""
|
|
context = state.setdefault("context", [])
|
|
if not content:
|
|
gr.Warning("Invalid inputs")
|
|
shown_context = get_shown_context(context)
|
|
return None, shown_context, context, state
|
|
|
|
# 构造 OpenAI API 要求的 messages 格式
|
|
messages = []
|
|
for turn in context[:-1]:
|
|
messages.append({"role": turn["role"], "content": remove_think_tags(turn["content"])})
|
|
|
|
# 默认模型名称从参数中获取
|
|
model = getattr(args, "model", default_params.get("model", ""))
|
|
payload = {
|
|
"model": model,
|
|
"messages": messages,
|
|
"temperature": temperature,
|
|
"repetition_penalty": repetition_penalty,
|
|
"max_tokens": max_tokens,
|
|
"src_length": src_length,
|
|
"top_p": top_p,
|
|
"top_k": top_k,
|
|
"stream": True,
|
|
}
|
|
headers = {
|
|
# "Authorization": "Bearer " + args.api_key,
|
|
"Content-Type": "application/json"
|
|
}
|
|
url = f"http://0.0.0.0:{args.flask_port}/v1/chat/completions"
|
|
try:
|
|
res = requests.post(url, json=payload, headers=headers, stream=True)
|
|
except Exception as e:
|
|
gr.Warning(f"请求异常: {e}")
|
|
shown_context = get_shown_context(context)
|
|
yield None, shown_context, context, state
|
|
return
|
|
|
|
# 流式处理返回结果，实时更新最后一个对话记录（即 bot 回复）
|
|
for line in res.iter_lines():
|
|
if line:
|
|
try:
|
|
decoded_line = line.decode("utf-8").strip()
|
|
# OpenAI 流返回每行以 "data:" 开头
|
|
if decoded_line.startswith("data:"):
|
|
data_str = decoded_line[len("data:") :].strip()
|
|
if data_str == "[DONE]":
|
|
logger.info("Conversation round over.")
|
|
break
|
|
data_json = json.loads(data_str)
|
|
|
|
# delta 中可能包含部分回复内容
|
|
delta = data_json["choices"][0]["delta"].get("content", "")
|
|
if delta:
|
|
# Reformat <think> tags to show in chatbot
|
|
delta = delta.replace("<think>", r"\<think\>")
|
|
delta = delta.replace("</think>", r"\<\/think\>")
|
|
context[-1]["content"] += delta
|
|
shown_context = get_shown_context(context)
|
|
yield None, shown_context, context, state
|
|
else:
|
|
logger.error(f"{decoded_line}")
|
|
gr.Warning(f"{decoded_line}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"解析返回结果异常: {e}")
|
|
gr.Warning(f"解析返回结果异常: {e}")
|
|
continue
|
|
|
|
def get_shown_context(context):
|
|
"""将对话上下文转换为 gr.Chatbot 显示格式，每一对 [用户, 助手]"""
|
|
shown_context = []
|
|
# 每两项组成一对
|
|
for turn_idx in range(0, len(context), 2):
|
|
user_text = context[turn_idx]["content"]
|
|
bot_text = context[turn_idx + 1]["content"] if turn_idx + 1 < len(context) else ""
|
|
shown_context.append([user_text, bot_text])
|
|
return shown_context
|
|
|
|
with gr.Blocks(title="LLM", theme=gr.themes.Soft()) as block:
|
|
gr.Markdown(f"# {args.title} <font style='color: red !important' size=2>{args.sub_title}</font>")
|
|
with gr.Row():
|
|
with gr.Column(scale=1):
|
|
top_k = gr.Slider(
|
|
minimum=0,
|
|
maximum=100,
|
|
value=0,
|
|
step=1,
|
|
label="Top-k",
|
|
info="控制采样token个数。(不建议设置)",
|
|
)
|
|
top_p = gr.Slider(
|
|
minimum=0,
|
|
maximum=1,
|
|
value=default_params.get("top_p", 0.7),
|
|
step=0.05,
|
|
label="Top-p",
|
|
info="控制采样范围。",
|
|
)
|
|
temperature = gr.Slider(
|
|
minimum=0.05,
|
|
maximum=1.5,
|
|
value=default_params.get("temperature", 0.95),
|
|
step=0.05,
|
|
label="Temperature",
|
|
info="温度，控制生成随机性。",
|
|
)
|
|
repetition_penalty = gr.Slider(
|
|
minimum=0.1,
|
|
maximum=10,
|
|
value=default_params.get("repetition_penalty", 1.2),
|
|
step=0.05,
|
|
label="Repetition Penalty",
|
|
info="生成结果重复惩罚。(不建议设置)",
|
|
)
|
|
default_src_length = default_params.get("src_length", 128)
|
|
total_length = default_src_length + default_params.get("max_tokens", 50)
|
|
src_length = create_src_slider(default_src_length, total_length)
|
|
max_tokens = create_max_slider(max(total_length - default_src_length, 50), total_length)
|
|
|
|
def src_length_change_event(src_length_value, max_tokens_value):
|
|
return create_max_slider(
|
|
min(total_length - src_length_value, max_tokens_value),
|
|
total_length - src_length_value,
|
|
)
|
|
|
|
def max_tokens_change_event(src_length_value, max_tokens_value):
|
|
return create_src_slider(
|
|
min(total_length - max_tokens_value, src_length_value),
|
|
total_length - max_tokens_value,
|
|
)
|
|
|
|
src_length.change(src_length_change_event, inputs=[src_length, max_tokens], outputs=max_tokens)
|
|
max_tokens.change(max_tokens_change_event, inputs=[src_length, max_tokens], outputs=src_length)
|
|
with gr.Column(scale=4):
|
|
state = gr.State({})
|
|
# 这里修改 gr.Chatbot 组件，启用 Markdown 渲染并支持 LaTeX 展示
|
|
context_chatbot = gr.Chatbot(
|
|
label="Context",
|
|
render_markdown=True,
|
|
latex_delimiters=[
|
|
{"left": "$$", "right": "$$", "display": True},
|
|
{"left": "\\[", "right": "\\]", "display": True},
|
|
{"left": "$", "right": "$", "display": True},
|
|
],
|
|
)
|
|
utt_text = gr.Textbox(placeholder="请输入...", label="Content")
|
|
with gr.Row():
|
|
clear_btn = gr.Button("清空")
|
|
rollback_btn = gr.Button("撤回")
|
|
regen_btn = gr.Button("重新生成")
|
|
send_btn = gr.Button("发送")
|
|
with gr.Row():
|
|
raw_context_json = gr.JSON(label="Raw Context")
|
|
|
|
utt_text.submit(
|
|
begin,
|
|
inputs=[utt_text, state],
|
|
outputs=[utt_text, context_chatbot, raw_context_json, state],
|
|
queue=False,
|
|
api_name="chat",
|
|
).then(
|
|
infer,
|
|
inputs=[utt_text, state, top_k, top_p, temperature, repetition_penalty, max_tokens, src_length],
|
|
outputs=[utt_text, context_chatbot, raw_context_json, state],
|
|
)
|
|
|
|
clear_btn.click(
|
|
lambda _: (None, None, None, {}),
|
|
inputs=clear_btn,
|
|
outputs=[utt_text, context_chatbot, raw_context_json, state],
|
|
api_name="clear",
|
|
show_progress=False,
|
|
)
|
|
rollback_btn.click(
|
|
rollback,
|
|
inputs=[state],
|
|
outputs=[utt_text, context_chatbot, raw_context_json, state],
|
|
show_progress=False,
|
|
)
|
|
regen_btn.click(
|
|
regen,
|
|
inputs=[state, top_k, top_p, temperature, repetition_penalty, max_tokens, src_length],
|
|
outputs=[utt_text, context_chatbot, raw_context_json, state],
|
|
queue=False,
|
|
api_name="chat",
|
|
).then(
|
|
infer,
|
|
inputs=[utt_text, state, top_k, top_p, temperature, repetition_penalty, max_tokens, src_length],
|
|
outputs=[utt_text, context_chatbot, raw_context_json, state],
|
|
)
|
|
|
|
send_btn.click(
|
|
begin,
|
|
inputs=[utt_text, state],
|
|
outputs=[utt_text, context_chatbot, raw_context_json, state],
|
|
queue=False,
|
|
api_name="chat",
|
|
).then(
|
|
infer,
|
|
inputs=[utt_text, state, top_k, top_p, temperature, repetition_penalty, max_tokens, src_length],
|
|
outputs=[utt_text, context_chatbot, raw_context_json, state],
|
|
)
|
|
|
|
block.queue().launch(server_name="0.0.0.0", server_port=args.port, debug=True)
|
|
|
|
|
|
def main(args, default_params: dict = {}):
|
|
launch(args, default_params)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# 可以在 default_params 中设置默认参数，如 src_length, max_tokens, temperature, top_p 等
|
|
default_params = {
|
|
"src_length": 1024,
|
|
"max_tokens": 1024,
|
|
"temperature": 0.95,
|
|
"top_p": 0.7,
|
|
}
|
|
args = setup_args()
|
|
main(args, default_params)
|