Update app.py

app.py

@@ -1,9 +1,8 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig  
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 import torch
 import os
 
-# 預先定義 Hugging Face 模型
 MODEL_NAMES = {
     "DeepSeek-R1-Distill-Qwen-7B": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
     "DeepSeek-R1-Distill-Llama-8B": "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
@@ -13,34 +12,27 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 
 def load_model(model_path):
     tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, token=HF_TOKEN)
-    
-    # 先載入 config，手動刪除量化設定，防止 FP8 問題
     config = AutoConfig.from_pretrained(model_path, trust_remote_code=True, token=HF_TOKEN)
     if hasattr(config, "quantization_config"):
         del config.quantization_config  # 刪除量化配置，避免使用 FP8
 
     model = AutoModelForCausalLM.from_pretrained(
         model_path,
-        config=config,  # 使用已移除量化的 config
+        config=config,
         trust_remote_code=True,
         token=HF_TOKEN,
-        torch_dtype=torch.float16,  # 強制 FP16，避免 FP8
+        torch_dtype=torch.float16,
         device_map="auto",
     )
     return model, tokenizer
 
-
-# 預設載入 DeepSeek-R1
 current_model, current_tokenizer = load_model("deepseek-ai/DeepSeek-R1-Distill-Llama-8B")
 
-
 def chat(message, history, model_name):
-    """處理聊天訊息"""
     global current_model, current_tokenizer
 
-    # 若模型不同則切換
     if model_name != current_model:
-        current_model, current_tokenizer = load_model(model_name)
+        current_model, current_tokenizer = load_model(MODEL_NAMES[model_name])
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
     inputs = current_tokenizer(message, return_tensors="pt").to(device)
@@ -49,15 +41,24 @@ def chat(message, history, model_name):
 
     return response
 
-
 with gr.Blocks() as app:
     gr.Markdown("## Chatbot with DeepSeek Models")
 
     with gr.Row():
-        chat_interface = gr.ChatInterface(chat, type="messages", flagging_mode="manual", save_history=True)
+        chat_interface = gr.ChatInterface(
+            chat,
+            type="messages",
+            flagging_mode="manual",
+            save_history=True,
+        )
         model_selector = gr.Dropdown(
-            choices=list(MODEL_NAMES.keys()), value="DeepSeek-R1-Distill-Llama-8B", label="Select Model"
+            choices=list(MODEL_NAMES.keys()),
+            value="DeepSeek-R1-Distill-Llama-8B",
+            label="Select Model",
         )
 
-    chat_interface.append(model_selector)
-    app.launch()
+    # 使用 gr.Blocks 的布局功能來組織元件
+    app.add_component(chat_interface)
+    app.add_component(model_selector)
+
+app.launch()