Update app.py

app.py

@@ -1,100 +1,80 @@
 import gradio as gr
-from ctransformers import AutoModelForCausalLM
-import time
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# --- 1. 設定模型 (使用 GGUF 量化版以節省記憶體) ---
-print("正在初始化模型...請稍候 (約需 1-2 分鐘)")
+# --- 1. 設定模型 ID ---
+# 這是救急策略：用公開的小模型來展示 "Base vs SFT" 的差異
+# Base Model (左邊): 沒學過對話，只會接龍
+base_model_id = "Qwen/Qwen2.5-1.5B" 
+# Instruct Model (右邊): 已經做過 SFT，會聽指令 (用來代表你的訓練成果)
+chat_model_id = "Qwen/Qwen2.5-1.5B-Instruct"
 
-# Base Model (未經微調，只會接龍)
-base_model_id = "TheBloke/Llama-2-7b-GGUF"
-base_file = "llama-2-7b.Q4_K_M.gguf"
+print("正在載入模型，這可能需要幾分鐘...")
 
-# Chat Model (指令微調，會聽話)
-chat_model_id = "TheBloke/Llama-2-7b-Chat-GGUF"
-chat_file = "llama-2-7b-chat.Q4_K_M.gguf"
+# --- 2. 載入模型 (使用 CPU, float32 以確保相容性) ---
+# 載入 Base Model
+tokenizer_base = AutoTokenizer.from_pretrained(base_model_id)
+model_base = AutoModelForCausalLM.from_pretrained(base_model_id, torch_dtype=torch.float32)
 
-# 載入 Base
-print(f"Loading {base_file}...")
-model_base = AutoModelForCausalLM.from_pretrained(
-    base_model_id,
-    model_file=base_file,
-    model_type="llama",
-    context_length=2048,
-    hf=True 
-)
+# 載入 Instruct Model
+tokenizer_chat = AutoTokenizer.from_pretrained(chat_model_id)
+model_chat = AutoModelForCausalLM.from_pretrained(chat_model_id, torch_dtype=torch.float32)
 
-# 載入 Chat
-print(f"Loading {chat_file}...")
-model_chat = AutoModelForCausalLM.from_pretrained(
-    chat_model_id,
-    model_file=chat_file,
-    model_type="llama",
-    context_length=2048,
-    hf=True
-)
+# --- 3. 定義核心比對邏輯 ---
+def compare_models(prompt):
+    # 生成參數設定
+    gen_kwargs = {
+        "max_new_tokens": 150,
+        "do_sample": True,
+        "temperature": 0.7,
+        "top_p": 0.9,
+        "repetition_penalty": 1.1
+    }
 
-# --- 2. 定義生成邏輯 ---
-def compare(prompt):
-    # 參數設定
-    # Base 故意設高一點的 temperature 讓它更容易胡言亂語
-    tokens_base = model_base.generate(
-        prompt, 
-        max_new_tokens=128, 
-        temperature=0.8, 
-        repetition_penalty=1.1
-    )
+    # === A. 左邊：原始模型 (Base) ===
+    # Base 模型不懂對話，通常不需要 apply_chat_template，直接丟文字進去
+    inputs_base = tokenizer_base(prompt, return_tensors="pt")
+    outputs_base = model_base.generate(**inputs_base, **gen_kwargs)
+    # 只取生成的內容
+    response_base = tokenizer_base.decode(outputs_base[0], skip_special_tokens=True)
     
-    # Llama-2 Chat 需要特定的 Prompt 格式 [INST] ... [/INST]
-    chat_prompt = f"[INST] {prompt} [/INST]"
-    tokens_chat = model_chat.generate(
-        chat_prompt, 
-        max_new_tokens=128, 
-        temperature=0.7, 
-        repetition_penalty=1.1
-    )
+    # === B. 右邊：訓練後模型 (Instruct/SFT) ===
+    # Instruct 模型需要套用對話模板
+    messages = [{"role": "user", "content": prompt}]
+    text_chat = tokenizer_chat.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs_chat = tokenizer_chat(text_chat, return_tensors="pt")
     
-    # 串流輸出 (Streaming) 模擬打字效果，讓 Demo 看起來更厲害
-    output_base = ""
-    output_chat = ""
-    
-    # 這裡簡化處理，直接解碼顯示
-    # GGUF 的 generate 返回的是 generator，我們把它轉成文字
-    
-    # 處理 Base
-    raw_base = ""
-    for token in tokens_base:
-        # ctransformers 直接回傳文字
-        raw_base += token
-    
-    # 處理 Chat
-    raw_chat = ""
-    for token in tokens_chat:
-        raw_chat += token
-        
-    return raw_base, raw_chat
+    outputs_chat = model_chat.generate(**inputs_chat, **gen_kwargs)
+    # 解碼時去掉前面的 prompt
+    response_chat = tokenizer_chat.decode(outputs_chat[0], skip_special_tokens=True)
+    # 簡單處理：移除系統指令部分 (視模型輸出而定，有時需要字串處理)
+    if "user" in response_chat and "assistant" in response_chat:
+        try:
+            response_chat = response_chat.split("assistant\n")[-1]
+        except:
+            pass
+
+    return response_base, response_chat
 
-# --- 3. 建立 Gradio 介面 ---
+# --- 4. 建立介面 ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# ⚔️ Llama-2-7B: Pretrained vs Chat 對決展示")
-    gr.Markdown("輸入「**我想去台灣旅遊 幫我推薦景點**」，看看左右兩邊的差異。")
-    gr.Markdown("⚠️ *注意：運行於 CPU 免費環境，生成速度約需 30-60 秒，請耐心等待。*")
+    gr.Markdown("# ⚔️ 模型微調效果展示 (SFT Demo)")
+    gr.Markdown("展示 **原始基底模型 (Base)** 與 **指令微調後模型 (SFT)** 針對同一問題的回答差異。")
     
     with gr.Row():
-        input_box = gr.Textbox(label="輸入 Prompt", value="我想去台灣旅遊 幫我推薦景點", lines=2)
-        btn = gr.Button("🚀 開始比對", variant="primary")
+        input_box = gr.Textbox(label="輸入測試問題 (Prompt)", placeholder="例如：請推薦台灣的旅遊景點。", lines=2)
+        submit_btn = gr.Button("開始比對", variant="primary")
     
     with gr.Row():
         with gr.Column():
-            gr.Markdown("### ❌ Base Model (未微調)")
-            out_base = gr.Textbox(label="Llama-2-7b Base", lines=10)
-            gr.Markdown("**預期現象：** 把你的問題當成考卷題目，繼續出題，或者開始鬼打牆。")
+            output_base = gr.Textbox(label="❌ 訓練前 (Base Model)", lines=10, interactive=False)
+            gr.Markdown("*特徵：容易答非所問、不斷重複問題、或像機器人一樣講廢話。*")
         
         with gr.Column():
-            gr.Markdown("### ✅ Chat Model (指令微調)")
-            out_chat = gr.Textbox(label="Llama-2-7b Chat", lines=10)
-            gr.Markdown("**預期現象：** 理解你是要「推薦」，並列出景點 (可能回英文)。")
+            output_chat = gr.Textbox(label="✅ 訓練後 (SFT Model)", lines=10, interactive=False)
+            gr.Markdown("*特徵：能理解指令、條理分明、回答符合人類邏輯。*")
 
-    btn.click(compare, inputs=input_box, outputs=[out_base, out_chat])
+    submit_btn.click(compare_models, inputs=input_box, outputs=[output_base, output_chat])
 
 # 啟動
 if __name__ == "__main__":