Update app.py

app.py

@@ -3,10 +3,12 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import gradio as gr
 
-# Initialize
+# 初始化
 peft_model_id = "CMLM/ZhongJing-2-1_8b"
 base_model_id = "Qwen/Qwen1.5-1.8B-Chat"
-model = AutoModelForCausalLM.from_pretrained(base_model_id, device_map="auto")
+
+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",
@@ -15,30 +17,9 @@ tokenizer = AutoTokenizer.from_pretrained(
     pad_token=''
 )
 
-# Single turn chat
-@spaces.GPU
-def single_turn_chat(question):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-    
-    prompt = f"Question: {question}"
-    messages = [
-        {"role": "system", "content": "You are a helpful TCM medical assistant named 仲景中医大语言模型, created by 医哲未来 of Fudan University."},
-        {"role": "user", "content": prompt}
-    ]
-    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    model_inputs = tokenizer([input_text], return_tensors="pt").to(device)
-    generated_ids = model.generate(model_inputs.input_ids, max_new_tokens=512)
-    generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    return response
-
-# Multi-turn chat
+#多轮对话
 @spaces.GPU
 def multi_turn_chat(question, chat_history=None):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-    
     if not isinstance(question, str):
         raise ValueError("The question must be a string.")
     
@@ -76,16 +57,7 @@ def multi_turn_chat(question, chat_history=None):
 def clear_history():
     return [], []
 
-# Single turn interface
-single_turn_interface = gr.Interface(
-    fn=single_turn_chat,
-    inputs=["text"],
-    outputs="text",
-    title="仲景GPT-V2-1.8B 单轮对话",
-    description="博极医源，精勤不倦。Unlocking the Wisdom of Traditional Chinese Medicine with AI."
-)
-
-# Multi-turn interface
+# 多轮界面
 with gr.Blocks() as multi_turn_interface:
     chatbot = gr.Chatbot(label="仲景GPT-V2-1.8B 多轮对话")
     state = gr.State([])
@@ -97,7 +69,8 @@ with gr.Blocks() as multi_turn_interface:
     
     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])
 
-single_turn_interface.launch()
 multi_turn_interface.launch()