Spaces:
Sleeping
Sleeping
import spaces | |
from transformers import AutoModelForCausalLM, AutoTokenizer | |
import torch | |
import gradio as gr | |
# 初始化 | |
peft_model_id = "CMLM/ZhongJing-2-1_8b" | |
base_model_id = "Qwen/Qwen1.5-1.8B-Chat" | |
device = "cuda" | |
model = AutoModelForCausalLM.from_pretrained(base_model_id, device_map={"": device}).to(device) | |
model.load_adapter(peft_model_id) | |
tokenizer = AutoTokenizer.from_pretrained( | |
"CMLM/ZhongJing-2-1_8b", | |
padding_side="right", | |
trust_remote_code=True, | |
pad_token='' | |
) | |
#多轮对话 | |
def multi_turn_chat(question, chat_history=None): | |
if not isinstance(question, str): | |
raise ValueError("The question must be a string.") | |
if chat_history is None or chat_history == []: | |
chat_history = [{"role": "system", "content": "You are a helpful TCM medical assistant named 仲景中医大语言模型, created by 医哲未来 of Fudan University."}] | |
chat_history.append({"role": "user", "content": question}) | |
# Apply the chat template and prepare the input | |
inputs = tokenizer.apply_chat_template(chat_history, tokenize=False, add_generation_prompt=True) | |
model_inputs = tokenizer([inputs], return_tensors="pt").to(device) | |
try: | |
# Generate the response from the model | |
outputs = model.generate(model_inputs.input_ids, max_new_tokens=512) | |
generated_ids = outputs[:, model_inputs.input_ids.shape[-1]:] | |
response = tokenizer.decode(generated_ids[0], skip_special_tokens=True) | |
except Exception as e: | |
raise RuntimeError("Error in model generation: " + str(e)) | |
# Append the assistant's response to the chat history | |
chat_history.append({"role": "assistant", "content": response}) | |
# Format the chat history for output | |
formatted_history = [] | |
tempuser = "" | |
for entry in chat_history: | |
if entry['role'] == 'user': | |
tempuser = entry['content'] | |
elif entry['role'] == 'assistant': | |
formatted_history.append((tempuser, entry['content'])) | |
return formatted_history, chat_history | |
def clear_history(): | |
return [], [] | |
# 多轮界面 | |
with gr.Blocks() as multi_turn_interface: | |
chatbot = gr.Chatbot(label="仲景GPT-V2-1.8B 多轮对话") | |
state = gr.State([]) | |
with gr.Row(): | |
with gr.Column(scale=6): | |
user_input = gr.Textbox(label="输入", placeholder="输入你的问题") | |
with gr.Column(scale=6): | |
submit_button = gr.Button("发送") | |
submit_button.click(multi_turn_chat, [user_input, state], [chatbot, state]) | |
user_input.submit(multi_turn_chat, [user_input, state], [chatbot, state]) | |
clear_button = gr.Button("清除对话历史") | |
clear_button.click(clear_history, [], [chatbot, state]) | |
multi_turn_interface.launch() | |