replace app.py

app.py

@@ -1,73 +1,49 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
+# app.py
+from transformers import AutoTokenizer, AutoModelForCausalLM
 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()
 
-
-# 支持流式输出的聊天函数
-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()
+# 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
 
 # Construct Gradio Interface
 interface = gr.Interface(
-    fn=chat_stream,
+    fn=chat,
     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 AI助手",
+    title="Ling-lite-1.5 MoE 模型 Demo",
     description="基于 [inclusionAI/Ling-lite-1.5](https://huggingface.co/inclusionAI/Ling-lite-1.5)  的对话式文本生成演示。",
     examples=[
         ["介绍大型语言模型的基本概念", 512],
@@ -76,4 +52,4 @@ interface = gr.Interface(
 )
 
 # launch Gradion Service
-interface.launch()
+interface.launch()