specify the GPU device and support streaming output

app.py

@@ -1,55 +1,88 @@
-# app.py
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 import gradio as gr
 import torch
 
-# load model and tokenizer
+# 加载模型和 Tokenizer
 model_name = "inclusionAI/Ling-lite-1.5"
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype="auto",
-    device_map="auto",  
+    device_map=device,
     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}
-    ]
+# 自定义 Streamer 以兼容 Gradio 异步回调
+class GradioStreamer(TextStreamer):
+    def __init__(self, tokenizer, chatbot, skip_prompt: bool = True, skip_special_tokens: bool = True):
+        super().__init__(tokenizer, skip_prompt=skip_prompt, skip_special_tokens=skip_special_tokens)
+        self.chatbot = chatbot
+        self.current_text = ""
+
+    def put(self, value):
+        # 解码 token 并追加到当前文本
+        self.current_text += super().decode(value)
+        # 更新 Chatbot 最后一条消息
+        self.chatbot[-1][1] = self.current_text
+        yield self.chatbot
+
+    def end(self):
+        # 结束时也触发一次更新
+        yield self.chatbot
+
+
+# 定义异步聊天函数
+async def chat_stream(message, chat_history, max_new_tokens=512):
+    # 构造系统提示 + 历史记录 + 当前问题
+    messages = [{"role": "system", "content": "You are Ling, an assistant created by inclusionAI"}]
+    for user, bot in chat_history:
+        messages.append({"role": "user", "content": user})
+        messages.append({"role": "assistant", "content": bot})
+    messages.append({"role": "user", "content": message})
+
     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,
-    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",
-    description="基于 [inclusionAI/Ling-lite-1.5](https://huggingface.co/inclusionAI/Ling-lite-1.5)  的对话式文本生成演示。",
-    examples=[
-        ["介绍大型语言模型的基本概念", 512],
-        ["如何解决数学问题中的长上下文依赖？", 768]
-    ]
-)
-
-# launch Gradion Service
-interface.launch()
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+    # 初始化 Chatbot 历史并创建流式对象
+    chat_history.append([message, ""])
+    streamer = GradioStreamer(tokenizer, chat_history)
+
+    # 异步生成（注意：transformers 的 generate 目前还不是 async，但我们可以模拟）
+    generation_kwargs = {
+        "input_ids": inputs["input_ids"],
+        "attention_mask": inputs["attention_mask"],
+        "streamer": streamer,
+        "max_new_tokens": max_new_tokens,
+        "pad_token_id": tokenizer.eos_token_id,
+    }
+
+    # 在后台线程中运行模型生成
+    import threading
+    thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    # 逐步返回 chat_history
+    while thread.is_alive():
+        yield chat_history
+        await asyncio.sleep(0.01)
+
+    # 返回最终结果
+    yield chat_history
+
+
+# 构建 Gradio 界面
+with gr.Blocks(title="Ling-lite-1.5 MoE 模型 Demo") as demo:
+    chatbot = gr.Chatbot(bubble_full_width=False, label="Ling 聊天机器人")
+    interface = gr.ChatInterface(
+        fn=chat_stream,
+        additional_inputs=[
+            gr.Slider(minimum=100, maximum=1024, step=50, label="生成长度", value=512),
+        ],
+        chatbot=chatbot
+    )
+
+# 启动服务
+demo.launch()