modify app.py

app.py

@@ -1,49 +1,73 @@
-# app.py
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
 import gradio as gr
 import torch
 
-# load model and tokenizer
+# 加载模型和分词器
 model_name = "inclusionAI/Ling-lite-1.5"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype="auto",
-    device_map="auto",  
+    device_map="auto",
     trust_remote_code=True
 ).eval()
 
-# define chat function
-def chat(user_input, max_new_tokens=512):
-    # chat history
-    messages = [
-        {"role": "system", "content": "You are Ling, an assistant created by inclusionAI"},
-        {"role": "user", "content": user_input}
-    ]
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    
-    # encode the input prompt
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
-    # 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
+
+# 支持流式输出的聊天函数
+def chat_stream(message, history):
+    system_prompt = {"role": "system", "content": "You are Ling, an assistant created by inclusionAI"}
+    user_message = {"role": "user", "content": message}
+
+    # 构建消息历史
+    messages = [system_prompt] + history + [user_message]
+
+    # 应用 chat template
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    # 编码输入
+    inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+    # 设置 streamer
+    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+
+    # 生成参数
+    generate_kwargs = dict(
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        streamer=streamer,
+        max_new_tokens=512,
+        do_sample=True,
+        temperature=0.7,
+        pad_token_id=tokenizer.eos_token_id
+    )
+
+    # 在后台线程中启动生成
+    def generate():
+        model.generate(**generate_kwargs)
+
+    thread = Thread(target=generate)
+    thread.start()
+
+    # 逐步读取生成的内容
+    response = ""
+    for new_text in streamer:
+        response += new_text
+        yield response.strip()
 
 # Construct Gradio Interface
 interface = gr.Interface(
-    fn=chat,
+    fn=chat_stream,
     inputs=[
         gr.Textbox(lines=5, label="输入你的问题"),
         gr.Slider(minimum=100, maximum=1024, step=50, label="生成长度")
     ],
     outputs=gr.Textbox(label="模型回复"),
-    title="Ling-lite-1.5 MoE 模型 Demo",
+    title="Ling-lite-1.5 MoE AI助手",
     description="基于 [inclusionAI/Ling-lite-1.5](https://huggingface.co/inclusionAI/Ling-lite-1.5)  的对话式文本生成演示。",
     examples=[
         ["介绍大型语言模型的基本概念", 512],
@@ -52,4 +76,4 @@ interface = gr.Interface(
 )
 
 # launch Gradion Service
-interface.launch()
+interface.launch()