Update app.py

app.py

@@ -2,17 +2,11 @@ import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import gradio as gr
-import os
 
-os.environ['CUDA_VISIBLE_DEVICES'] = "0,1"
-USE_CUDA = torch.cuda.is_available()
-device_ids_parallel = [0]
-device = torch.device("cuda:{}".format(device_ids_parallel[0]) if USE_CUDA else "cpu")
-
-# 初始化
+# 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").to(device)
+model = AutoModelForCausalLM.from_pretrained(base_model_id, device_map="auto")
 model.load_adapter(peft_model_id)
 tokenizer = AutoTokenizer.from_pretrained(
     "CMLM/ZhongJing-2-1_8b",
@@ -21,9 +15,12 @@ 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."},
@@ -36,9 +33,12 @@ def single_turn_chat(question):
     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,7 +76,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"],
@@ -85,7 +85,7 @@ single_turn_interface = gr.Interface(
     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([])
@@ -100,3 +100,4 @@ with gr.Blocks() as multi_turn_interface:
 
 single_turn_interface.launch()
 multi_turn_interface.launch()
+