Update app.py

app.py

@@ -1,19 +1,21 @@
 import os
 
-# ---- Robust fix for OMP_NUM_THREADS (HF / K8s may set it to '7500m') ----
+# --- Robust fix: HF/K8s may set OMP_NUM_THREADS like "7500m" (invalid for libgomp) ---
 _raw_omp = os.getenv("OMP_NUM_THREADS", "")
 if not _raw_omp.isdigit():
     os.environ["OMP_NUM_THREADS"] = "1"
 
-
 import threading
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 
+# -------------------------
+# Config
+# -------------------------
 MODEL_ID = os.getenv("MODEL_ID", "Milkfish033/deepseek-r1-1.5b-merged")
 
-# 🔒 固定 system prompt（UI 不暴露）
+# 🔒 固定系统提示词（不在 UI 暴露）
 SYSTEM_PROMPT = "你是 Bello，一个友好的智能助手。请用清晰、简洁的中文回答用户问题。"
 
 theme = gr.themes.Soft()
@@ -35,7 +37,7 @@ footer { display: none !important; }
   padding: 12px;
 }
 
-/* 输入框边框 */
+/* 输入框边框更明显 */
 .chat-card textarea,
 .chat-card input {
   border: 1px solid #d1d5db !important;
@@ -46,25 +48,38 @@ footer { display: none !important; }
 /* 发送按钮圆角 */
 .chat-card button { border-radius: 14px !important; }
 
-/* 气泡样式（不同 gradio 版本 class 不同，多写点提高命中率） */
+/* 气泡样式：不同 gradio 版本 class 名不一，多写 selector 提升命中 */
 .chat-card .message.user,
-.chat-card .bubble.user {
+.chat-card .bubble.user,
+.chat-card [data-testid="chatbot"] .message.user,
+.chat-card [data-testid="chatbot"] .bubble.user {
   background: #eef2ff !important;
   border: 1px solid #e0e7ff !important;
   border-radius: 16px !important;
 }
 
-.chat-card .message.bot,
 .chat-card .message.assistant,
+.chat-card .message.bot,
+.chat-card .bubble.assistant,
 .chat-card .bubble.bot,
-.chat-card .bubble.assistant {
+.chat-card [data-testid="chatbot"] .message.assistant,
+.chat-card [data-testid="chatbot"] .bubble.assistant {
   background: #f8fafc !important;
   border: 1px solid #eef2f7 !important;
   border-radius: 16px !important;
 }
+
+/* 每条消息 spacing */
+.chat-card .message,
+.chat-card .bubble {
+  padding: 10px 12px !important;
+  margin: 8px 0 !important;
+}
 """
 
-# ---- Load model once ----
+# -------------------------
+# Load model once
+# -------------------------
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
@@ -74,20 +89,60 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
-
-def _build_prompt(history_pairs, user_msg: str) -> str:
+# -------------------------
+# History adapter (CRITICAL FIX)
+# -------------------------
+def history_to_messages(history) -> list[dict]:
     """
-    旧版 ChatInterface 的 history 是 [(user, bot), ...]
-    我们把它转成 messages，再用 chat_template 生成 prompt
+    兼容 Gradio ChatInterface 不同版本/不同轮次可能传入的 history 结构：
+    A) 旧格式: [(user, bot), ...]
+    B) 新格式: [{"role":"user"/"assistant","content":"..."}, ...]
+    C) 其它异常结构：尽量容错，不抛异常
     """
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-    for u, a in history_pairs:
-        if u:
-            messages.append({"role": "user", "content": u})
-        if a:
-            messages.append({"role": "assistant", "content": a})
+    msgs = [{"role": "system", "content": SYSTEM_PROMPT}]
+
+    if not history:
+        return msgs
+
+    # 1) 如果是 list，拿第一个元素判断结构
+    first = history[0]
+
+    # Case A: tuple/list pairs
+    if isinstance(first, (tuple, list)):
+        for item in history:
+            if not isinstance(item, (tuple, list)):
+                continue
+            # 有些情况下可能是 (user, bot, meta...) 长度>2
+            user = item[0] if len(item) > 0 else ""
+            bot = item[1] if len(item) > 1 else ""
+            if user:
+                msgs.append({"role": "user", "content": str(user)})
+            if bot:
+                msgs.append({"role": "assistant", "content": str(bot)})
+        return msgs
+
+    # Case B: dict messages
+    if isinstance(first, dict) and "role" in first:
+        for m in history:
+            if not isinstance(m, dict):
+                continue
+            role = m.get("role")
+            content = m.get("content", "")
+            if role in ("user", "assistant"):
+                msgs.append({"role": role, "content": str(content)})
+        return msgs
+
+    # Case C: unknown -> stringify
+    for item in history:
+        msgs.append({"role": "assistant", "content": str(item)})
+    return msgs
+
+
+def build_prompt(history, user_msg: str) -> str:
+    messages = history_to_messages(history)
     messages.append({"role": "user", "content": user_msg})
 
+    # 优先使用模型自带 chat template（deepseek 带 jinja template）
     if hasattr(tokenizer, "apply_chat_template"):
         try:
             return tokenizer.apply_chat_template(
@@ -98,20 +153,21 @@ def _build_prompt(history_pairs, user_msg: str) -> str:
         except Exception:
             pass
 
-    # fallback
+    # fallback：简单拼接
     prompt = f"System: {SYSTEM_PROMPT}\n"
-    for u, a in history_pairs:
-        prompt += f"User: {u}\nAssistant: {a}\n"
-    prompt += f"User: {user_msg}\nAssistant:"
+    for m in messages:
+        if m["role"] == "user":
+            prompt += f"User: {m['content']}\n"
+        elif m["role"] == "assistant":
+            prompt += f"Assistant: {m['content']}\n"
+    prompt += "Assistant:"
     return prompt
 
-
-def respond(message: str, history):
-    """
-    ✅ 兼容旧版 gradio.ChatInterface：fn(message, history) -> str 或 generator
-    history: List[Tuple[str, str]]
-    """
-    prompt = _build_prompt(history, message)
+# -------------------------
+# Generation (streaming)
+# -------------------------
+def respond(message, history, max_tokens=512, temperature=0.7, top_p=0.95):
+    prompt = build_prompt(history, message)
 
     inputs = tokenizer(prompt, return_tensors="pt")
     if torch.cuda.is_available():
@@ -120,16 +176,16 @@ def respond(message: str, history):
     streamer = TextIteratorStreamer(
         tokenizer,
         skip_special_tokens=True,
-        skip_prompt=True,     # ✅ 不回显 prompt（解决 <|User|> 问题）
+        skip_prompt=True,  # ✅ 不回显 prompt（解决 <|User|>...）
     )
 
     gen_kwargs = dict(
         **inputs,
         streamer=streamer,
-        max_new_tokens=512,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.95,
+        max_new_tokens=int(max_tokens),
+        do_sample=(float(temperature) > 0),
+        temperature=float(temperature),
+        top_p=float(top_p),
         pad_token_id=tokenizer.eos_token_id,
     )
 
@@ -141,20 +197,64 @@ def respond(message: str, history):
         out += piece
         yield out.strip()
 
-
-with gr.Blocks(theme=theme, css=css) as demo:
+# -------------------------
+# UI
+# -------------------------
+with gr.Blocks() as demo:
     with gr.Column(elem_classes=["page-wrap"]):
         gr.Markdown("# 我是 Bello，有什么能帮到您？")
 
         with gr.Column(elem_classes=["chat-card"]):
-            # ✅ 老版本不支持 type="messages"，不要传 type
-            gr.ChatInterface(
-                fn=respond,
-                title="",
-                description="",
+            # ✅ 不传 type="messages"（避免旧版本报错）
+            # ✅ 用额外输入手工加 sliders（兼容性更稳）
+            chatbot = gr.Chatbot(height=520)
+            msg = gr.Textbox(placeholder="请输入问题...", show_label=False)
+            send = gr.Button("发送")
+
+            with gr.Row():
+                max_tokens = gr.Slider(1, 2048, value=512, step=1, label="Max new tokens")
+                temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="Temperature")
+                top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
+
+            def _user_submit(user_message, chat_history):
+                # chat_history 是 [(user, bot), ...] 但有时会携带 meta，因此我们不强依赖结构
+                if chat_history is None:
+                    chat_history = []
+                chat_history = list(chat_history)
+                chat_history.append((user_message, ""))  # 先占位
+                return "", chat_history
+
+            def _bot_stream(chat_history, max_tokens, temperature, top_p):
+                # 取最后一条 user
+                if not chat_history:
+                    return chat_history
+                last_user = chat_history[-1][0]
+
+                # history 给模型：去掉最后一条占位（只传已完成的对话）
+                prior = chat_history[:-1]
+
+                # 用我们的 respond()（它能吃 tuple 或 dict messages）
+                gen = respond(last_user, prior, max_tokens=max_tokens, temperature=temperature, top_p=top_p)
+
+                partial = ""
+                for chunk in gen:
+                    partial = chunk
+                    chat_history[-1] = (last_user, partial)
+                    yield chat_history
+
+            # Enter 提交
+            msg.submit(_user_submit, [msg, chatbot], [msg, chatbot], queue=False).then(
+                _bot_stream, [chatbot, max_tokens, temperature, top_p], chatbot
+            )
+
+            # 点击按钮提交
+            send.click(_user_submit, [msg, chatbot], [msg, chatbot], queue=False).then(
+                _bot_stream, [chatbot, max_tokens, temperature, top_p], chatbot
             )
 
-demo.queue(default_concurrency_limit=4)
+# 并发先设低，稳定第一；确认稳定后你再调大
+demo.queue(default_concurrency_limit=1)
 
 if __name__ == "__main__":
-    demo.launch(ssr_mode=False)
+    # Gradio 6: theme/css 建议放到 launch()
+    demo.launch(ssr_mode=False, theme=theme, css=css)