Update app.py

app.py

@@ -1,53 +1,66 @@
 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 = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    trust_remote_code=True,
-)
 
+# 建议将模型定义为全局变量，但在函数内初始化，避免重复加载
+model = None
 
+@spaces.GPU # 必须加在生成函数上
 def respond(message, history):
+    global model
+    # 只有在第一次运行且 GPU 环境就绪时加载模型
+    if model is None:
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            torch_dtype=torch.float16,
+            # device_map="auto" 在 ZeroGPU 中容易报错，建议直接 .to("cuda")
+            trust_remote_code=True,
+        ).to("cuda")
+
     messages = [{"role": "system", "content": "You are Spartacus, a helpful assistant."}]
     for msg in history:
-        messages.append({"role": msg["role"], "content": msg["content"]})
+        # 注意：Gradio 的 history 结构可能需要适配
+        role = "assistant" if msg['role'] == 'assistant' else "user"
+        messages.append({"role": 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(model.device)
+    ).to("cuda")
 
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-
+    
     generate_kwargs = dict(
         input_ids=input_ids,
         streamer=streamer,
-        temperature=0.5,
-        top_p=0.9,
+        max_new_tokens=2048, # 建议加上限制，防止超时
+        temperature=0.6,
+        top_p=0.95,
         do_sample=True,
     )
 
+    # 在 ZeroGPU 中使用 Thread 也是允许的
     thread = Thread(target=model.generate, kwargs=generate_kwargs)
     thread.start()
 
     response = ""
     for token in streamer:
-        response += token
-        yield response
-
+        yield token # 直接 yield 增量或累计文本
 
 demo = gr.ChatInterface(
     fn=respond,
+    type="messages", # 建议明确指定 Gradio 的消息类型
     title="Spartacus Chat",
     description="Chat with NoesisLab/Spartacus-1B-Instruct",
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()