Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 from typing import List, Dict, Optional, Tuple
+import time
 
 # ==========================================
 # Helper: dtype map & loader with simple cache
@@ -21,7 +22,9 @@ def _dtype_from_name(name: str):
 
 
 def load_model_and_tokenizer(repo_id: str, device_map: str = "cpu", dtype_name: str = "auto"):
-    
+    """Load & cache (tokenizer, model) keyed by (repo_id, device_map, dtype). No low_cpu_mem_usage.
+    Prefer `dtype=...`; on TypeError fallback to `torch_dtype=` or omit.
+    """
     key = (repo_id, device_map, dtype_name)
     if key in _MODEL_CACHE:
         return _MODEL_CACHE[key]
@@ -48,7 +51,7 @@ def load_model_and_tokenizer(repo_id: str, device_map: str = "cpu", dtype_name:
 
 
 # ==========================================
-# Chat utilities
+# Chat utilities & logging helpers
 # ==========================================
 
 def messages_to_pairs(messages: List[Dict[str, str]]) -> List[Tuple[str, str]]:
@@ -72,33 +75,66 @@ def messages_to_pairs(messages: List[Dict[str, str]]) -> List[Tuple[str, str]]:
     return pairs
 
 
+def _ts() -> str:
+    return time.strftime("%H:%M:%S")
+
+
+def append_log(logs: str, msg: str) -> str:
+    line = f"[{_ts()}] {msg}\n"
+    return (logs + line) if logs else line
+
+
+# ==========================================
+# Model state helpers (reload only when repo_id changes)
+# ==========================================
+
+def ensure_model(model_state: Dict, repo_id: str, device_map: str, dtype_name: str, logs: str):
+    """Ensure a model is available in model_state.
+    Only (re)load when repo_id changes or model_state is empty.
+    device_map/dtype_name 变更不会触发重新加载（按你的要求）。
+    """
+    ms = model_state or {"repo_id": None, "tok": None, "model": None}
+    if ms.get("repo_id") != repo_id or ms.get("model") is None:
+        logs = append_log(logs, f"加载模型 {repo_id}（触发：repo 变更）…")
+        tok, mdl = load_model_and_tokenizer(repo_id, device_map=device_map, dtype_name=dtype_name)
+        ms = {"repo_id": repo_id, "tok": tok, "model": mdl}
+        logs = append_log(logs, "模型加载完成。")
+    else:
+        logs = append_log(logs, f"使用已加载模型 {repo_id}（缓存）")
+    return ms, ms["tok"], ms["model"], logs
+
+
+# ==========================================
+# Predict
+# ==========================================
+
 def predict(user_text: str,
             messages_state: List[Dict[str, str]],
             repo_id: str, device_map: str, dtype_name: str,
-            max_new_token: int, top_k: int):
-    
+            max_new_token: int, top_k: int,
+            logs_state: str,
+            model_state: Dict):
+    """Generator for streaming output + live logs.
+    Only reload when repo_id changes.
+    Expects custom model.chat(conversations, tokenizer, max_new_token=..., top_k=...).
+    """
     messages_state = messages_state or []
+    logs_state = logs_state or ""
 
-    # Append user message
+    # 1) Ensure model based on repo_id only
+    model_state, tokenizer, model, logs_state = ensure_model(model_state, repo_id, device_map, dtype_name, logs_state)
+
+    # 2) Append user & paint
     messages_state.append({"role": "user", "content": user_text or ""})
+    logs_state = append_log(logs_state, f"收到输入：{(user_text or '').strip()[:60]}")
 
-    # Build initial display and show immediately
     chat_display = messages_to_pairs(messages_state)
-    yield chat_display, messages_state
-
-    # Load model/tokenizer lazily
-    try:
-        tokenizer, model = load_model_and_tokenizer(repo_id, device_map=device_map, dtype_name=dtype_name)
-    except Exception as e:
-        err = f"[加载错误] {e}"
-        # Show error as assistant reply
-        chat_display[-1] = (chat_display[-1][0], err)
-        messages_state.append({"role": "assistant", "content": err})
-        yield chat_display, messages_state
-        return
+    yield chat_display, messages_state, logs_state, logs_state, model_state
 
-    # Inference
+    # 3) Inference
     try:
+        logs_state = append_log(logs_state, f"开始推理：max_new_token={int(max_new_token)}, top_k={int(top_k)}")
+        yield chat_display, messages_state, logs_state, logs_state, model_state
         try:
             output = model.chat(
                 messages_state,
@@ -107,35 +143,53 @@ def predict(user_text: str,
                 top_k=int(top_k),
             )
         except TypeError:
-            # Fallback to minimal signature if custom signature not supported in current build
             output = model.chat(messages_state, tokenizer)
 
         partial = ""
         for ch in str(output):
             partial += ch
             chat_display[-1] = (chat_display[-1][0], partial)
-            yield chat_display, messages_state
+            yield chat_display, messages_state, logs_state, logs_state, model_state
 
-        # Finalize state
         messages_state.append({"role": "assistant", "content": str(output)})
-        yield chat_display, messages_state
+        logs_state = append_log(logs_state, f"推理完成，输出长度 {len(str(output))} 字符。")
+        yield chat_display, messages_state, logs_state, logs_state, model_state
     except Exception as e:
         err = f"[推理错误] {e}"
+        logs_state = append_log(logs_state, err)
         chat_display[-1] = (chat_display[-1][0], err)
         messages_state.append({"role": "assistant", "content": err})
-        yield chat_display, messages_state
+        yield chat_display, messages_state, logs_state, logs_state, model_state
 
 
 def clear_chat():
     return [], []  # chatbot pairs, messages_state
 
 
+def clear_logs_fn():
+    return "", ""  # logs_box text, logs_state
+
+
+def preload_on_repo_change(repo_id: str, device_map: str, dtype_name: str, logs_state: str, model_state: Dict):
+    """当仓库切换时，预加载模型并写日志。"""
+    logs_state = logs_state or ""
+    model_state, _, _, logs_state = ensure_model(model_state, repo_id, device_map, dtype_name, logs_state)
+    return logs_state, model_state
+
+
 # ==========================================
 # Gradio UI
 # ==========================================
 with gr.Blocks(title="mini-moe Chat (Gradio)") as demo:
+    gr.Markdown("""
+    # 🤖 mini-moe Chat UI (Gradio)
+    仅在 **repo 变更** 时重新加载模型；设备/精度变更不会触发重新加载（按你的要求）。
+    右侧含 **日志面板**，实时显示加载与推理步骤；**不使用 system prompt**。
+    """)
 
     messages_state = gr.State([])  # 保存 role/content 历史
+    logs_state = gr.State("")     # 保存日志文本
+    model_state = gr.State({"repo_id": None, "tok": None, "model": None})  # 当前已加载模型
 
     with gr.Row():
         with gr.Column(scale=2):
@@ -156,10 +210,15 @@ with gr.Blocks(title="mini-moe Chat (Gradio)") as demo:
             dtype_dd = gr.Dropdown(label="精度 (dtype/torch_dtype)", choices=["auto", "float32", "bfloat16", "float16"], value="auto")
             max_new_num = gr.Number(label="max_new_token", value=256, precision=0)
             top_k_num = gr.Number(label="top_k", value=5, precision=0)
+            with gr.Accordion("📜 日志 (展开查看)", open=False):
+                logs_box = gr.Textbox(label="运行日志", lines=12, interactive=False)
+                log_clear_btn = gr.Button("清空日志")
 
     # Events: send / submit
-    send_evt_inputs = [user_box, messages_state, repo_dd, device_dd, dtype_dd, max_new_num, top_k_num]
-    send_evt_outputs = [chatbot, messages_state]
+    send_evt_inputs = [
+        user_box, messages_state, repo_dd, device_dd, dtype_dd, max_new_num, top_k_num, logs_state, model_state
+    ]
+    send_evt_outputs = [chatbot, messages_state, logs_box, logs_state, model_state]
 
     send_btn.click(predict, inputs=send_evt_inputs, outputs=send_evt_outputs)
     user_box.submit(predict, inputs=send_evt_inputs, outputs=send_evt_outputs)
@@ -173,6 +232,14 @@ with gr.Blocks(title="mini-moe Chat (Gradio)") as demo:
     # Clear chat
     clear_btn.click(clear_chat, inputs=None, outputs=[chatbot, messages_state])
 
+    # Clear logs
+    log_clear_btn.click(clear_logs_fn, inputs=None, outputs=[logs_box, logs_state])
+
+    # Preload on repo change (only reload on repo change)
+    repo_dd.change(preload_on_repo_change,
+                   inputs=[repo_dd, device_dd, dtype_dd, logs_state, model_state],
+                   outputs=[logs_box, model_state])
+
 
 if __name__ == "__main__":
     demo.queue().launch()  # set share=True if you want a public link