Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import threading
 
 MODEL_NAME = "my2000cup/Gaia-Petro-LLM"
 
@@ -21,29 +22,35 @@ def build_prompt(history, system_message, user_message):
         if assistant:
             messages.append({"role": "assistant", "content": assistant})
     messages.append({"role": "user", "content": user_message})
-    # 如果你有chat模板支持，推荐用apply_chat_template
     if hasattr(tokenizer, "apply_chat_template"):
         prompt = tokenizer.apply_chat_template(
             messages, tokenize=False, add_generation_prompt=True
         )
     else:
-        # fallback: 简单拼接
         prompt = "\n".join([f"{m['role']}: {m['content']}" for m in messages]) + "\nassistant:"
     return prompt
 
 def respond(message, history, system_message, max_tokens, temperature, top_p):
     prompt = build_prompt(history, system_message, message)
     inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
-    output = model.generate(
+
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    # 在新线程中异步生成
+    generation_kwargs = dict(
         **inputs,
+        streamer=streamer,
         max_new_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
         do_sample=True,
-        pad_token_id=tokenizer.eos_token_id
+        pad_token_id=tokenizer.eos_token_id,
     )
-    response = tokenizer.decode(output[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
-    yield response
+    gen_thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
+    gen_thread.start()
+    output = ""
+    for new_text in streamer:
+        output += new_text
+        yield output
 
 demo = gr.ChatInterface(
     respond,