app.py

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
# ==========================================
DTYPE_MAP = {
    "auto": "auto",
    "float32": torch.float32,
    "bfloat16": torch.bfloat16,
    "float16": torch.float16,
}

_MODEL_CACHE = {}


def _dtype_from_name(name: str):
    return DTYPE_MAP.get(name, "auto")


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]

    tok = AutoTokenizer.from_pretrained(repo_id)
    dtype_val = _dtype_from_name(dtype_name)
    common_kwargs = dict(trust_remote_code=True, device_map=device_map)

    model = None
    try:
        if dtype_name == "auto":
            model = AutoModelForCausalLM.from_pretrained(repo_id, dtype="auto", **common_kwargs)
        else:
            model = AutoModelForCausalLM.from_pretrained(repo_id, dtype=dtype_val, **common_kwargs)
    except TypeError:
        if dtype_name == "auto":
            model = AutoModelForCausalLM.from_pretrained(repo_id, **common_kwargs)
        else:
            model = AutoModelForCausalLM.from_pretrained(repo_id, torch_dtype=dtype_val, **common_kwargs)

    model.eval()
    _MODEL_CACHE[key] = (tok, model)
    return tok, model


# ==========================================
# Chat utilities & logging helpers
# ==========================================

def messages_to_pairs(messages: List[Dict[str, str]]) -> List[Tuple[str, str]]:
    """Convert role-based messages into (user, assistant) pairs for Gradio Chatbot."""
    pairs: List[Tuple[str, str]] = []
    i = 0
    while i < len(messages):
        msg = messages[i]
        if msg.get("role") == "user":
            user = msg.get("content", "")
            if i + 1 < len(messages) and messages[i + 1].get("role") == "assistant":
                assistant = messages[i + 1].get("content", "")
                pairs.append((user, assistant))
                i += 2
            else:
                pairs.append((user, ""))
                i += 1
        else:
            # Skip unexpected assistant-first cases
            i += 1
    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):

    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,
            logs_state: str,
            model_state: Dict):
    
    messages_state = messages_state or []
    logs_state = logs_state or ""

    # 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]}")

    chat_display = messages_to_pairs(messages_state)
    yield chat_display, messages_state, logs_state, logs_state, model_state

    # 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,
                tokenizer,
                max_new_token=int(max_new_token),
                top_k=int(top_k),
            )
        except TypeError:
            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, logs_state, logs_state, model_state

        messages_state.append({"role": "assistant", "content": str(output)})
        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, 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:

    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):
            chatbot = gr.Chatbot(label="对话", height=520)
            with gr.Row():
                user_box = gr.Textbox(label="输入", placeholder="请输入你的问题… (Shift+Enter 换行)", lines=3)
            with gr.Row():
                send_btn = gr.Button("发送", variant="primary")
                clear_btn = gr.Button("清空对话")
        with gr.Column(scale=1):
            gr.Markdown("## ⚙️ 设置")
            repo_dd = gr.Dropdown(
                label="模型仓库 (HF repo)",
                choices=["caixiaoshun/mini-moe", "caixiaoshun/mini-llm"],
                value="caixiaoshun/mini-moe",
            )
            device_dd = gr.Dropdown(label="device_map", choices=["cpu", "auto"], value="cpu")
            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, 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)

    # Clear input after send
    def _clear_text():
        return ""
    send_btn.click(_clear_text, inputs=None, outputs=user_box)
    user_box.submit(_clear_text, inputs=None, outputs=user_box)

    # 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