Merge branch 'main' of https://huggingface.co/spaces/unausagi/chatbot

app.py

@@ -1,15 +1,15 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 import torch
 import os
 
-# 預先定義 Hugging Face 模型
 MODEL_NAMES = {
-    "DeepSeek-V3": "deepseek-ai/DeepSeek-V3",
-    "DeepSeek-R1": "deepseek-ai/DeepSeek-R1",
+    "DeepSeek-R1-Distill-Qwen-7B": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
+    "DeepSeek-R1-Distill-Llama-8B": "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
 }
 
 HF_TOKEN = os.getenv("HF_TOKEN")
+<<<<<<< HEAD
 
 
 def load_model(model_path):
@@ -19,36 +19,60 @@ def load_model(model_path):
     model = AutoModelForCausalLM.from_pretrained(
         model_path, torch_dtype=torch.float16, trust_remote_code=True, token=HF_TOKEN
     ).cuda()
-    return model, tokenizer
+=======
 
+def load_model(model_path):
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, token=HF_TOKEN)
+    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True, token=HF_TOKEN)
+    if hasattr(config, "quantization_config"):
+        del config.quantization_config  # 刪除量化配置，避免使用 FP8
 
-# 預設載入 DeepSeek-V3
-current_model, current_tokenizer = load_model("deepseek-ai/DeepSeek-V3")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        config=config,
+        trust_remote_code=True,
+        token=HF_TOKEN,
+        torch_dtype=torch.float16,
+        device_map="auto",
+    )
+>>>>>>> c4a529954b9cc11b7b8c2dcdac51f3ef3a83c616
+    return model, tokenizer
 
+# 初始化預設模型
+current_model_name = "DeepSeek-R1-Distill-Llama-8B"
+current_model, current_tokenizer = load_model(MODEL_NAMES[current_model_name])
 
 def chat(message, history, model_name):
-    """處理聊天訊息"""
-    global current_model, current_tokenizer
+    global current_model, current_tokenizer, current_model_name
 
-    # 若模型不同則切換
-    if model_name != current_model:
-        current_model, current_tokenizer = load_model(model_name)
+    # 檢查是否需要更換模型
+    if model_name != current_model_name:
+        current_model_name = model_name
+        current_model, current_tokenizer = load_model(MODEL_NAMES[model_name])
 
-    inputs = current_tokenizer(message, return_tensors="pt").to("cuda")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    inputs = current_tokenizer(message, return_tensors="pt").to(device)
     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"
-        )
+    model_selector = gr.Dropdown(
+        choices=list(MODEL_NAMES.keys()),
+        value=current_model_name,
+        label="Select Model",
+    )
+
+    chat_interface = gr.ChatInterface(
+        fn=lambda message, history: chat(message, history, model_selector.value),
+        type="messages",
+        flagging_mode="manual",
+        save_history=True,
+    )
+
+    model_selector.change(fn=lambda model_name: None, inputs=[model_selector], outputs=[])
 
-    chat_interface.append(model_selector)
-    app.launch()
+app.launch()

requirements.txt

@@ -1,2 +1,8 @@
+<<<<<<< HEAD
 transformers>=4.37.2
-torch>=2.1.0
+torch>=2.1.0
+=======
+transformers>=4.38.0
+torch>=2.1.0
+accelerate
+>>>>>>> c4a529954b9cc11b7b8c2dcdac51f3ef3a83c616