164 lines
5.8 KiB
Python
164 lines
5.8 KiB
Python
import os
|
||
|
||
os.system('pip install transformers -U')
|
||
os.system('pip install modelscope -U')
|
||
os.system('pip install accelerate')
|
||
from threading import Thread
|
||
from typing import Iterator
|
||
|
||
import gradio as gr
|
||
import torch
|
||
from modelscope import AutoModelForCausalLM, AutoTokenizer
|
||
from transformers import TextIteratorStreamer
|
||
|
||
MAX_MAX_NEW_TOKENS = 2048
|
||
DEFAULT_MAX_NEW_TOKENS = 1024
|
||
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
|
||
|
||
# 系统提示
|
||
SYSTEM_PROMPT = """
|
||
作为助手，你的角色是通过系统的长期思考过程，深入探讨问题，然后提供最终精确的解决方案。
|
||
这需要通过分析、总结、探索、重新评估、反思、回溯和迭代等全面的思维循环来发展深思熟虑的思考过程。
|
||
请将你的回答结构化为两个主要部分：思考和解决方案。
|
||
在“思考”部分，使用指定的格式详细描述你的推理过程：
|
||
<think> {思考过程，每个步骤之间用“\n\n”分隔} </think>
|
||
|
||
每个步骤应包括详细的考虑事项，例如分析问题、总结相关发现、头脑风暴新想法、验证当前步骤的准确性、修正任何错误以及回顾之前的步骤。
|
||
在“解决方案”部分，根据“思考”部分的各种尝试、探索和反思，系统地呈现你认为正确的最终解决方案。
|
||
解决方案应保持逻辑准确、简洁表达，并详细说明达到结论所需的必要步骤，格式如下：
|
||
<answer> {最终格式化的、精确和清晰的解决方案} </answer>
|
||
"""
|
||
|
||
# 无论 GPU 是否可用，都加载模型和分词器
|
||
model_id = "daniel0527/qwen2.5-3b-dsitill"
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
model_id,
|
||
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
|
||
device_map="auto" if device == "cuda" else None
|
||
)
|
||
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
||
tokenizer.use_default_system_prompt = False
|
||
|
||
|
||
def generate(
|
||
message: str,
|
||
chat_history: list[tuple[str, str]],
|
||
system_prompt: str = SYSTEM_PROMPT,
|
||
max_new_tokens: int = 1024,
|
||
temperature: float = 0.6,
|
||
top_p: float = 0.9,
|
||
top_k: int = 50,
|
||
repetition_penalty: float = 1.2,
|
||
) -> Iterator[str]:
|
||
conversation = []
|
||
if system_prompt:
|
||
conversation.append({"role": "system", "content": system_prompt})
|
||
for user, assistant in chat_history:
|
||
conversation.extend([
|
||
{"role": "user", "content": user},
|
||
{"role": "assistant", "content": assistant}
|
||
])
|
||
conversation.append({"role": "user", "content": message})
|
||
|
||
# 生成对话模板字符串，并进行编码时添加 padding 以获得 attention_mask
|
||
input_str = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
|
||
inputs = tokenizer([input_str], return_tensors="pt", padding=True)
|
||
inputs = inputs.to(model.device)
|
||
|
||
# 调整 timeout 为 30 秒，必要时也可增加
|
||
streamer = TextIteratorStreamer(tokenizer, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
|
||
generate_kwargs = dict(
|
||
input_ids=inputs.input_ids,
|
||
attention_mask=inputs.attention_mask, # 传入 attention_mask
|
||
streamer=streamer,
|
||
max_new_tokens=max_new_tokens,
|
||
do_sample=True,
|
||
top_p=top_p,
|
||
top_k=top_k,
|
||
temperature=temperature,
|
||
repetition_penalty=repetition_penalty,
|
||
)
|
||
t = Thread(target=model.generate, kwargs=generate_kwargs)
|
||
t.start()
|
||
|
||
outputs = []
|
||
# 使用 try/except 捕获可能的 Empty 异常，确保生成过程结束后退出
|
||
try:
|
||
for text in streamer:
|
||
outputs.append(text)
|
||
yield "".join(outputs)
|
||
except Exception:
|
||
yield "".join(outputs)
|
||
|
||
|
||
# 创建 Gradio ChatInterface，不在此处添加示例问题
|
||
chat_interface = gr.ChatInterface(
|
||
fn=generate,
|
||
additional_inputs=[
|
||
gr.Textbox(label="系统提示", value=SYSTEM_PROMPT, lines=6),
|
||
gr.Slider(
|
||
label="最大生成标记数",
|
||
minimum=1,
|
||
maximum=MAX_MAX_NEW_TOKENS,
|
||
step=1,
|
||
value=DEFAULT_MAX_NEW_TOKENS,
|
||
),
|
||
gr.Slider(
|
||
label="温度",
|
||
minimum=0.1,
|
||
maximum=4.0,
|
||
step=0.1,
|
||
value=0.6,
|
||
),
|
||
gr.Slider(
|
||
label="Top-p（核采样）",
|
||
minimum=0.05,
|
||
maximum=1.0,
|
||
step=0.05,
|
||
value=0.9,
|
||
),
|
||
gr.Slider(
|
||
label="Top-k",
|
||
minimum=1,
|
||
maximum=1000,
|
||
step=1,
|
||
value=50,
|
||
),
|
||
gr.Slider(
|
||
label="重复惩罚",
|
||
minimum=1.0,
|
||
maximum=2.0,
|
||
step=0.05,
|
||
value=1.2,
|
||
),
|
||
],
|
||
stop_btn="停止",
|
||
)
|
||
|
||
# 定义示例问题列表（将在聊天框下方展示）
|
||
example_questions = [
|
||
"在患者做闭眼动作时，患侧眼球向外上方转动的现象被称为什么？",
|
||
"对于一位病情稳定且血压在正常范围内波动，且未发生类似发作的患者，为了帮助诊断，应该进行哪种刺激试验？?",
|
||
"肚子疼去那个科室？",
|
||
]
|
||
|
||
with gr.Blocks(css="style.css") as demo:
|
||
gr.Markdown(
|
||
"""<p align="center"><img src="https://modelscope.oss-cn-beijing.aliyuncs.com/resource/qwen.png" style="height: 80px"/></p>""")
|
||
gr.Markdown("""<center><font size=8>jishi-3b-r1-Chat Bot👾</center>""")
|
||
gr.Markdown("""<center><font size=4>jishi-3b-r1是医疗行业30亿规模的大模型。</center>""")
|
||
|
||
chat_interface.render()
|
||
|
||
# 在聊天组件下方添加示例问题
|
||
gr.Markdown("### 示例问题")
|
||
for question in example_questions:
|
||
gr.Markdown(f"- {question}")
|
||
|
||
if __name__ == "__main__":
|
||
# 如果需要创建公网链接，请设置 share=True
|
||
demo.queue(max_size=20).launch()
|
||
|
||
|