Update app.py

app.py

@@ -1,57 +1,70 @@
-import gradio as gr
-from transformers import pipeline
-
-# --- 1. 設定模型路徑 (請替換成你自己的 Hugging Face Model ID 或 本地路徑) ---
-# 為了示範，這裡使用兩個公開的小模型來模擬 "原版" vs "微調版"
-base_model_id = "gpt2"  # 假設這是原模型
-tuned_model_id = "distilgpt2" # 假設這是訓練後模型 (這裡先用 distil 代替)
-
-# --- 2. 載入模型 (使用 pipeline 比較方便) ---
-print("正在載入模型，請稍候...")
-# 實務上建議使用 GPU (device=0)，如果是 Hugging Face Free Space 則用 CPU (device=-1)
-pipe_base = pipeline("text-generation", model=base_model_id)
-pipe_tuned = pipeline("text-generation", model=tuned_model_id)
-
-# --- 3. 定義推論函數 ---
-def compare_models(input_text):
-    # 設定生成參數 (可以根據需求調整 max_length, temperature 等)
-    gen_kwargs = {"max_length": 100, "num_return_sequences": 1}
-    
-    # 原模型生成
-    out_base = pipe_base(input_text, **gen_kwargs)[0]['generated_text']
-    
-    # 訓練後模型生成
-    out_tuned = pipe_tuned(input_text, **gen_kwargs)[0]['generated_text']
-    
-    return out_base, out_tuned
-
-# --- 4. 建置 Gradio 介面 ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚀 模型訓練成果比對系統")
-    gr.Markdown("輸入一段 Prompt，同時觀察原模型與訓練後模型的差異。")
-
-    with gr.Row():
-        input_box = gr.Textbox(label="輸入測試文字 (Prompt)", placeholder="例如：請寫一封信給客戶...", lines=3)
-
-    submit_btn = gr.Button("開始比對", variant="primary")
-
-    gr.Markdown("### 📊 比對結果")
-    
-    # 使用 Row 來讓兩個輸出框並排顯示
-    with gr.Row():
-        with gr.Column():
-            output_base = gr.Textbox(label="❌ 原模型輸出 (Base Model)", lines=10, interactive=False)
-        
-        with gr.Column():
-            output_tuned = gr.Textbox(label="✅ 訓練後模型 (Fine-tuned)", lines=10, interactive=False)
-
-    # 設定按鈕動作
-    submit_btn.click(
-        fn=compare_models, 
-        inputs=input_box, 
-        outputs=[output_base, output_tuned]
-    )
-
-# --- 5. 啟動 ---
-if __name__ == "__main__":
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+
+# --- 設定 ---
+base_model_id = "如果你用Llama3這裡填Llama3的路徑"  # 例如 "unsloth/llama-3-8b-bnb-4bit"
+adapter_model_id = "你的帳號/你的模型名稱-lora" # 剛剛在 Kaggle 推上去的那個 ID
+
+# --- 1. 載入模型 (記憶體優化版) ---
+print("正在載入基底模型...")
+tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+
+# 如果是免費 CPU Space，不要用 4bit (bitsandbytes 對 CPU 支援不好)，直接用 float32 或 bfloat16
+# 如果你有買 GPU Space，一定要加 load_in_4bit=True
+model = AutoModelForCausalLM.from_pretrained(
+    base_model_id,
+    device_map="auto",
+    torch_dtype=torch.float16, # CPU 建議用 float32，有 GPU 用 float16
+    low_cpu_mem_usage=True
+)
+
+print("正在掛載 LoRA Adapter...")
+# 這一步把訓練好的微調層掛上去
+model = PeftModel.from_pretrained(model, adapter_model_id)
+
+# --- 2. 定義推論邏輯 (核心技巧) ---
+def compare_inference(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    # 設定生成參數
+    gen_kwargs = {
+        "max_new_tokens": 150,
+        "do_sample": True,
+        "temperature": 0.7,
+        "top_p": 0.9
+    }
+
+    # === A. 生成：原模型 (Base Model) ===
+    # 使用 disable_adapter() 暫時關閉 LoRA，讓模型變回原本的樣子
+    with model.disable_adapter():
+        output_base_ids = model.generate(**inputs, **gen_kwargs)
+        output_base = tokenizer.decode(output_base_ids[0], skip_special_tokens=True)
+        # 簡單處理，只取 prompt 之後的文字
+        response_base = output_base.replace(prompt, "").strip()
+
+    # === B. 生成：訓練後模型 (Fine-tuned) ===
+    # 這裡正常生成，LoRA 會生效
+    output_tuned_ids = model.generate(**inputs, **gen_kwargs)
+    output_tuned = tokenizer.decode(output_tuned_ids[0], skip_special_tokens=True)
+    response_tuned = output_tuned.replace(prompt, "").strip()
+
+    return response_base, response_tuned
+
+# --- 3. 介面設計 ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# ⚔️ SFT 模型效果比對")
+    gr.Markdown(f"基底模型: `{base_model_id}` vs 微調權重: `{adapter_model_id}`")
+    
+    inp = gr.Textbox(label="輸入測試指令 (Prompt)", placeholder="例如：這家公司的財報重點是什麼？", lines=2)
+    btn = gr.Button("開始比對 (Generate)", variant="primary")
+    
+    with gr.Row():
+        out_base = gr.Textbox(label="原始模型 (Base)", lines=8)
+        out_tuned = gr.Textbox(label="訓練後模型 (SFT)", lines=8)
+    
+    btn.click(compare_inference, inputs=inp, outputs=[out_base, out_tuned])
+
+if __name__ == "__main__":
     demo.launch()