Update app.py

app.py

@@ -1,50 +1,40 @@
 import gradio as gr
 import torch
-import spaces
 from threading import Thread
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 
 MODEL_ID = "NoesisLab/Spartacus-1B-Instruct"
 
-# 静态加载 Tokenizer (不占 GPU)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-model = None
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    trust_remote_code=True,
+)
 
+import spaces
 @spaces.GPU
 def respond(message, history):
-    global model
-    # ZeroGPU 核心逻辑：在装饰器函数内初始化并移动到 CUDA
-    if model is None:
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID,
-            torch_dtype=torch.float16, 
-            trust_remote_code=True,
-        ).to("cuda")
-
-    # Gradio 5.x 的 history 已经是 [{'role': 'user', 'content': '...'}, ...] 格式
-    # 直接拼接到 messages 即可
     messages = [{"role": "system", "content": "You are Spartacus, a helpful assistant."}]
-    messages.extend(history)
+    for msg in history:
+        messages.append({"role": msg["role"], "content": msg["content"]})
     messages.append({"role": "user", "content": message})
 
     input_ids = tokenizer.apply_chat_template(
-        messages, 
-        add_generation_prompt=True, 
-        return_tensors="pt"
-    ).to("cuda")
+        messages, add_generation_prompt=True, return_tensors="pt"
+    ).to(model.device)
 
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    
+
     generate_kwargs = dict(
         input_ids=input_ids,
         streamer=streamer,
-        max_new_tokens=2048,
-        temperature=0.6,
-        top_p=0.95,
+        temperature=0.5,
+        top_p=0.9,
         do_sample=True,
     )
 
-    # 启动异步生成线程
     thread = Thread(target=model.generate, kwargs=generate_kwargs)
     thread.start()
 
@@ -53,13 +43,11 @@ def respond(message, history):
         response += token
         yield response
 
-# 使用最新版的配置参数
+
 demo = gr.ChatInterface(
     fn=respond,
-    type="messages", # 这需要 gradio>=5.0.0
     title="Spartacus Chat",
     description="Chat with NoesisLab/Spartacus-1B-Instruct",
-    examples=["Who are you?", "Explain the concept of Noesis."],
 )
 
 if __name__ == "__main__":