Update app.py

app.py

@@ -1,36 +1,34 @@
 # app.py
-from threading import Thread
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-import torch
 import spaces
+from threading import Thread
+import torch
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TextIteratorStreamer,
+)
 
-# ---------------------------------------------
+# ------------------------------
 # 1. 加载模型与 Tokenizer
-# ---------------------------------------------
-# 如果你的模型需要加速/量化等特殊配置，可在 from_pretrained() 中添加相应参数
-# 例如 device_map='auto' 或 trust_remote_code=True 等
+# ------------------------------
 model_name = "agentica-org/DeepScaleR-1.5B-Preview"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
 
-# 根据需要加上 .half()/.float()/.quantize() 等操作
-# 例如
-# model.half()
-# 或者
-# model = model.quantize(4/8)  # 如果你的模型和环境支持
+# 如果 tokenizer 没有设置 pad_token_id，则显式指定为 eos_token_id
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+
 
-# ---------------------------------------------
-# 2. 对话历史处理
-# ---------------------------------------------
+# ------------------------------
+# 2. 对话历史 -> Prompt 格式
+# ------------------------------
 def preprocess_messages(history):
     """
-    将所有的用户与回复消息拼成一个文本 prompt。
-    这里仅示例最简单的形式：
-       User: ...
-       Assistant: ...
-    最后再接上 "Assistant: " 用于提示模型继续生成。
-    你也可以修改为自己需要的对话模板。
+    将聊天记录拼成一个最简单的 Prompt。
+    你可以自定义更适合该模型的提示格式或特殊 Token。
     """
     prompt = ""
     for user_msg, assistant_msg in history:
@@ -38,27 +36,33 @@ def preprocess_messages(history):
             prompt += f"User: {user_msg}\n"
         if assistant_msg:
             prompt += f"Assistant: {assistant_msg}\n"
-
-    # 继续生成时，让模型再续写 "Assistant:"
+    # 继续生成时，提示 "Assistant:"
     prompt += "Assistant: "
     return prompt
 
-# ---------------------------------------------
-# 3. 预测函数
-# ---------------------------------------------
-@spaces.GPU()  # 使用 huggingface spaces 的 GPU 装饰器
+
+# ------------------------------
+# 3. 预测 / 推理函数
+# ------------------------------
+@spaces.GPU()  # 让 huggingface spaces 调用 GPU
 def predict(history, max_length, top_p, temperature):
     """
-    输入为 history（对话历史）和若干超参，输出流式生成的结果。
-    每生成一个 token，就通过 yield 返回给 Gradio，更新界面。
+    基于当前的 history 做文本生成。
+    使用 HF 提供的 TextIteratorStreamer 实现流式生成。
     """
     prompt = preprocess_messages(history)
 
-    # 组装输入
-    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        padding=True,       # 自动 padding
+        truncation=True,    # 超长截断
+        max_length=2048     # 你可根据显存大小或模型上限做调整
+    )
     input_ids = inputs["input_ids"].to(model.device)
+    attention_mask = inputs["attention_mask"].to(model.device)
 
-    # 使用 TextIteratorStreamer 来实现流式输出
+    # 流式输出器
     streamer = TextIteratorStreamer(
         tokenizer=tokenizer,
         timeout=60,
@@ -68,103 +72,98 @@ def predict(history, max_length, top_p, temperature):
 
     generate_kwargs = {
         "input_ids": input_ids,
-        "max_new_tokens": max_length,
+        "attention_mask": attention_mask,
+        "max_new_tokens": max_length,   # 新生成的 token 数
         "do_sample": True,
         "top_p": top_p,
         "temperature": temperature,
         "repetition_penalty": 1.2,
         "streamer": streamer,
-        # 如果需要自定义一些特殊 token 或其他参数可在此补充
-        # "eos_token_id": ...
     }
 
-    # 启动一个线程去执行 generate，然后主线程读取流式输出
+    # 在后台线程中执行 generate，主线程循环读取新 token
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
 
-    # history[-1][1] 存放当前最新的 assistant 回复，因此不断累加
+    # 将最新生成的 token 依次拼接到 history[-1][1]
     partial_output = ""
     for new_token in streamer:
         partial_output += new_token
         history[-1][1] = partial_output
         yield history
 
-# ---------------------------------------------
-# 4. 搭建 Gradio 界面
-# ---------------------------------------------
+
+# ------------------------------
+# 4. Gradio UI
+# ------------------------------
 def main():
     with gr.Blocks() as demo:
-        gr.HTML("<h1 align='center'>DeepScaleR-1.5B-Preview Chat Demo</h1>")
+        gr.HTML("<h1 align='center'>DeepScaleR-1.5B Chat Demo</h1>")
 
-        # 聊天窗口
         chatbot = gr.Chatbot()
 
         with gr.Row():
             with gr.Column(scale=2):
                 user_input = gr.Textbox(
-                    show_label=True, 
-                    placeholder="请输入您的问题...", 
+                    show_label=True,
+                    placeholder="请输入您的问题...",
                     label="User Input"
                 )
                 submitBtn = gr.Button("Submit")
-                emptyBtn = gr.Button("Clear History")
+                clearBtn = gr.Button("Clear History")
             with gr.Column(scale=1):
                 max_length = gr.Slider(
-                    minimum=0, 
-                    maximum=32000,  # 根据模型能力自行调整
-                    value=512, 
-                    step=1, 
-                    label="Max New Tokens", 
+                    minimum=0,
+                    maximum=1024,  # 可根据需要调大/调小
+                    value=512,
+                    step=1,
+                    label="Max New Tokens",
                     interactive=True
                 )
                 top_p = gr.Slider(
-                    minimum=0, 
-                    maximum=1, 
-                    value=0.8, 
-                    step=0.01, 
-                    label="Top P", 
+                    minimum=0,
+                    maximum=1,
+                    value=0.8,
+                    step=0.01,
+                    label="Top P",
                     interactive=True
                 )
                 temperature = gr.Slider(
-                    minimum=0.01, 
-                    maximum=2.0, 
-                    value=0.7, 
-                    step=0.01, 
-                    label="Temperature", 
+                    minimum=0.0,
+                    maximum=2.0,
+                    value=0.7,
+                    step=0.01,
+                    label="Temperature",
                     interactive=True
                 )
 
-        # 用于将用户输入插入到 chatbot 历史中
+        # 用户点击 Submit 时，先将输入添加到 history，然后再调用 predict 生成
         def user(query, history):
             return "", history + [[query, ""]]
 
-        # Submit: 
-        # 1) user() -> 新增一条 (user输入,"") 的对话记录
-        # 2) predict() -> 基于更新后的 history 进行生成
         submitBtn.click(
-            fn=user, 
-            inputs=[user_input, chatbot], 
+            fn=user,
+            inputs=[user_input, chatbot],
             outputs=[user_input, chatbot],
-            queue=False
+            queue=False  # 不排队
         ).then(
-            fn=predict, 
-            inputs=[chatbot, max_length, top_p, temperature], 
+            fn=predict,
+            inputs=[chatbot, max_length, top_p, temperature],
             outputs=chatbot
         )
 
-        # Clear: 清空对话历史
+        # 清空聊天记录
         def clear_history():
             return [], []
-        emptyBtn.click(
-            fn=clear_history, 
-            inputs=[], 
-            outputs=[chatbot, user_input], 
-            queue=False
-        )
 
-        # 可选：让 Gradio 自动对排队请求进行调度
-        demo.queue()
+        clearBtn.click(fn=clear_history, inputs=[], outputs=[chatbot, user_input], queue=False)
+
+        # 可选：启用队列防止并发冲突
+        demo.queue(concurrency_count=1)
         demo.launch()
 
+# ------------------------------
+# 入口
+# ------------------------------
 if __name__ == "__main__":
     main()