add stream output

app.py

@@ -1,5 +1,6 @@
 # app.py
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from threading import Thread
 import gradio as gr
 import torch
 
@@ -24,16 +25,32 @@ def chat(user_input, max_new_tokens=512):
     
     # encode the input prompt
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    #create streamer
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+
+    def generate():
+        model.generate(**inputs, max_new_tokens=max_new_tokens, streamer=streamer)
+
+    thread = Thread(target=generate)
+    thread.start()
+
+    generated_text = ""
+    for new_text in streamer:
+        generated_text += new_text
+        yield generated_text
+
+    thread.join()
     
     # generate response
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            pad_token_id=tokenizer.eos_token_id
-        )
-    response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[-1]:], skip_special_tokens=True)
-    return response
+    #with torch.no_grad():
+    #    outputs = model.generate(
+    #        **inputs,
+    #        max_new_tokens=max_new_tokens,
+    #        pad_token_id=tokenizer.eos_token_id
+    #    )
+    #response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[-1]:], skip_special_tokens=True)
+    #return response
 
 # Construct Gradio Interface
 interface = gr.Interface(
@@ -42,7 +59,7 @@ interface = gr.Interface(
         gr.Textbox(lines=5, label="输入你的问题"),
         gr.Slider(minimum=100, maximum=1024, step=50, label="生成长度")
     ],
-    outputs=gr.Textbox(label="模型回复"),
+    outputs=gr.Textbox(label="模型回复", stream=True),
     title="Ling-lite-1.5 MoE 模型 Demo",
     description="基于 [inclusionAI/Ling-lite-1.5](https://huggingface.co/inclusionAI/Ling-lite-1.5)  的对话式文本生成演示。",
     examples=[