app.py

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 預先定義 Hugging Face 模型
MODEL_NAMES = {
    "DeepSeek-V3": "deepseek-ai/DeepSeek-V3",
    "DeepSeek-R1": "deepseek-ai/DeepSeek-R1",
}


def load_model(model_path):
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_path, torch_dtype=torch.float16, trust_remote_code=True
    ).cuda()
    return model, tokenizer


# 預設載入 DeepSeek-V3
current_model, current_tokenizer = load_model("deepseek-ai/DeepSeek-V3")


def chat(message, history, model_name):
    """處理聊天訊息"""
    global current_model, current_tokenizer

    # 若模型不同則切換
    if model_name != current_model:
        current_model, current_tokenizer = load_model(model_name)

    inputs = current_tokenizer(message, return_tensors="pt").to("cuda")
    outputs = current_model.generate(**inputs, max_length=1024)
    response = current_tokenizer.decode(outputs[0], skip_special_tokens=True)

    return response


with gr.Blocks() as app:
    gr.Markdown("## Chatbot with DeepSeek Models")

    with gr.Row():
        chat_interface = gr.ChatInterface(chat, streaming=True, save_history=True)
        model_selector = gr.Dropdown(
            choices=list(MODEL_NAMES.keys()), value="DeepSeek-V3", label="Select Model"
        )

    chat_interface.append(model_selector)
    app.launch()