Update app.py

app.py

@@ -19,6 +19,7 @@ import numpy as np
 from datetime import datetime
 import json
 import os
+import gc  # 用於記憶體清理
 
 # PEFT 相關的 import（LoRA 和 AdaLoRA）
 try:
@@ -37,8 +38,7 @@ except ImportError:
 # 檢查 GPU
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# 全域變數儲存最後訓練的模型路徑和 tokenizer
-LAST_MODEL_PATH = None
+_MODEL_PATH = None
 LAST_TOKENIZER = None
 LAST_TUNING_METHOD = None
 
@@ -157,6 +157,12 @@ def run_original_code_with_tuning(
     
     global LAST_MODEL_PATH, LAST_TOKENIZER, LAST_TUNING_METHOD
     
+    # ==================== 清空記憶體（訓練前） ====================
+    import gc
+    torch.cuda.empty_cache()
+    gc.collect()
+    print("🧹 記憶體已清空")
+    
     # ==================== 您的原始程式碼開始 ====================
     
     # 讀取上傳的檔案
@@ -466,6 +472,32 @@ def run_original_code_with_tuning(
         model.save_pretrained(save_dir)
         tokenizer.save_pretrained(save_dir)
     
+    # 儲存模型資訊到 JSON 檔案（用於預測頁面選擇）
+    model_info = {
+        'model_path': save_dir,
+        'tuning_method': tuning_method,
+        'best_metric': best_metric,
+        'best_metric_value': float(results[f'eval_{metric_map.get(best_metric, "f1")}']),
+        'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
+        'weight_multiplier': weight_multiplier,
+        'epochs': epochs
+    }
+    
+    # 讀取現有的模型列表
+    models_list_file = './saved_models_list.json'
+    if os.path.exists(models_list_file):
+        with open(models_list_file, 'r') as f:
+            models_list = json.load(f)
+    else:
+        models_list = []
+    
+    # 加入新模型資訊
+    models_list.append(model_info)
+    
+    # 儲存更新後的列表
+    with open(models_list_file, 'w') as f:
+        json.dump(models_list, f, indent=2)
+    
     # 儲存到全域變數供預測使用
     LAST_MODEL_PATH = save_dir
     LAST_TOKENIZER = tokenizer
@@ -477,40 +509,131 @@ def run_original_code_with_tuning(
     print("=" * 80)
     print(f"完成時間: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
     
+    # ==================== 清空記憶體（訓練後） ====================
+    del model
+    del trainer
+    torch.cuda.empty_cache()
+    gc.collect()
+    print("🧹 訓練後記憶體已清空")
+    
     # 加入所有資訊到結果中
     results['tuning_method'] = tuning_method
     results['best_metric'] = best_metric
     results['best_metric_value'] = results[f'eval_{metric_map.get(best_metric, "f1")}']
     results['baseline_results'] = baseline_results
+    results['model_path'] = save_dir
     
     return results
 
-def predict_text(text_input):
+def predict_text(model_choice, text_input):
     """
-    預測功能 - 使用最後訓練的模型
+    預測功能 - 支援選擇已訓練的模型，並同時顯示未微調和微調的預測結果
     """
-    global LAST_MODEL_PATH, LAST_TOKENIZER, LAST_TUNING_METHOD
     
-    if LAST_MODEL_PATH is None:
-        return "❌ 請先訓練模型"
+    if not text_input or text_input.strip() == "":
+        return "請輸入文本", "請輸入文本"
     
     try:
-        print(f"\n使用模型: {LAST_MODEL_PATH}")
+        # ==================== 未微調的 BERT 預測 ====================
+        print("\n使用未微調 BERT 預測...")
+        baseline_tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+        baseline_model = BertForSequenceClassification.from_pretrained(
+            "bert-base-uncased",
+            num_labels=2
+        ).to(device)
+        baseline_model.eval()
+        
+        # Tokenize 輸入（未微調）
+        baseline_inputs = baseline_tokenizer(
+            text_input,
+            truncation=True,
+            padding='max_length',
+            max_length=256,
+            return_tensors='pt'
+        ).to(device)
+        
+        # 預測（未微調）
+        with torch.no_grad():
+            baseline_outputs = baseline_model(**baseline_inputs)
+            baseline_probs = torch.nn.functional.softmax(baseline_outputs.logits, dim=-1)
+            baseline_pred_class = baseline_probs.argmax(-1).item()
+            baseline_confidence = baseline_probs[0][baseline_pred_class].item()
+        
+        baseline_result = "存活" if baseline_pred_class == 0 else "死亡"
+        baseline_prob_survive = baseline_probs[0][0].item()
+        baseline_prob_death = baseline_probs[0][1].item()
+        
+        baseline_output = f"""
+# 🔵 未微調 BERT 預測結果
+
+## 預測類別: **{baseline_result}**
+
+## 信心度: **{baseline_confidence:.1%}**
+
+## 機率分布:
+- 🟢 **存活機率**: {baseline_prob_survive:.2%}
+- 🔴 **死亡機率**: {baseline_prob_death:.2%}
+
+---
+**說明**: 此為原始 BERT 模型，未經任何領域資料訓練
+        """
+        
+        # 清空記憶體
+        del baseline_model
+        del baseline_tokenizer
+        torch.cuda.empty_cache()
+        
+        # ==================== 微調後的 BERT 預測 ====================
+        
+        if model_choice == "請先訓練模型":
+            finetuned_output = """
+# 🟢 微調 BERT 預測結果
+
+❌ 尚未訓練任何模型，請先在「模型訓練」頁面訓練模型
+            """
+            return baseline_output, finetuned_output
+        
+        # 解析選擇的模型路徑
+        model_path = model_choice.split(" | ")[0].replace("路徑: ", "")
+        
+        # 從 JSON 讀取模型資訊
+        with open('./saved_models_list.json', 'r') as f:
+            models_list = json.load(f)
+        
+        selected_model_info = None
+        for model_info in models_list:
+            if model_info['model_path'] == model_path:
+                selected_model_info = model_info
+                break
+        
+        if selected_model_info is None:
+            finetuned_output = f"""
+# 🟢 微調 BERT 預測結果
+
+❌ 找不到模型：{model_path}
+            """
+            return baseline_output, finetuned_output
+        
+        print(f"\n使用微調模型: {model_path}")
+        
+        # 載入 tokenizer
+        finetuned_tokenizer = BertTokenizer.from_pretrained(model_path)
         
         # 載入模型
-        if LAST_TUNING_METHOD in ["LoRA", "AdaLoRA"] and PEFT_AVAILABLE:
+        tuning_method = selected_model_info['tuning_method']
+        if tuning_method in ["LoRA", "AdaLoRA"] and PEFT_AVAILABLE:
             # 載入 PEFT 模型
             base_model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
-            model = PeftModel.from_pretrained(base_model, LAST_MODEL_PATH)
-            model = model.to(device)
+            finetuned_model = PeftModel.from_pretrained(base_model, model_path)
+            finetuned_model = finetuned_model.to(device)
         else:
             # 載入一般模型
-            model = BertForSequenceClassification.from_pretrained(LAST_MODEL_PATH).to(device)
+            finetuned_model = BertForSequenceClassification.from_pretrained(model_path).to(device)
         
-        model.eval()
+        finetuned_model.eval()
         
-        # Tokenize 輸入
-        inputs = LAST_TOKENIZER(
+        # Tokenize 輸入（微調）
+        finetuned_inputs = finetuned_tokenizer(
             text_input,
             truncation=True,
             padding='max_length',
@@ -518,41 +641,72 @@ def predict_text(text_input):
             return_tensors='pt'
         ).to(device)
         
-        # 預測
+        # 預測（微調）
         with torch.no_grad():
-            outputs = model(**inputs)
-            probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-            pred_class = probs.argmax(-1).item()
-            confidence = probs[0][pred_class].item()
+            finetuned_outputs = finetuned_model(**finetuned_inputs)
+            finetuned_probs = torch.nn.functional.softmax(finetuned_outputs.logits, dim=-1)
+            finetuned_pred_class = finetuned_probs.argmax(-1).item()
+            finetuned_confidence = finetuned_probs[0][finetuned_pred_class].item()
         
-        # 準備結果
-        result = "存活" if pred_class == 0 else "死亡"
-        prob_survive = probs[0][0].item()
-        prob_death = probs[0][1].item()
+        finetuned_result = "存活" if finetuned_pred_class == 0 else "死亡"
+        finetuned_prob_survive = finetuned_probs[0][0].item()
+        finetuned_prob_death = finetuned_probs[0][1].item()
         
-        output = f"""
-## 🔮 預測結果
+        finetuned_output = f"""
+# 🟢 微調 BERT 預測結果
 
-### 預測類別: **{result}**
+## 預測類別: **{finetuned_result}**
 
-### 信心度: **{confidence:.1%}**
+## 信心度: **{finetuned_confidence:.1%}**
 
-### 機率分布:
-- 🟢 **存活機率**: {prob_survive:.2%}
-- 🔴 **死亡機率**: {prob_death:.2%}
+## 機率分布:
+- 🟢 **存活機率**: {finetuned_prob_survive:.2%}
+- 🔴 **死亡機率**: {finetuned_prob_death:.2%}
 
+---
 ### 模型資訊:
-- 使用方法: {LAST_TUNING_METHOD}
-- 模型路徑: {LAST_MODEL_PATH}
+- **微調方法**: {selected_model_info['tuning_method']}
+- **最佳化指標**: {selected_model_info['best_metric']}
+- **訓練時間**: {selected_model_info['timestamp']}
+- **模型路徑**: {model_path}
 
 ---
 **注意**: 此預測僅供參考，實際醫療決策應由專業醫師判斷。
         """
         
-        return output
+        # 清空記憶體
+        del finetuned_model
+        del finetuned_tokenizer
+        torch.cuda.empty_cache()
+        
+        return baseline_output, finetuned_output
         
     except Exception as e:
-        return f"❌ 預測錯誤：{str(e)}"
+        import traceback
+        error_msg = f"❌ 預測錯誤：{str(e)}\n\n詳細錯誤訊息：\n{traceback.format_exc()}"
+        return error_msg, error_msg
+
+def get_available_models():
+    """
+    取得所有已訓練的模型列表
+    """
+    models_list_file = './saved_models_list.json'
+    if not os.path.exists(models_list_file):
+        return ["請先訓練模型"]
+    
+    with open(models_list_file, 'r') as f:
+        models_list = json.load(f)
+    
+    if len(models_list) == 0:
+        return ["請先訓練模型"]
+    
+    # 格式化模型選項
+    model_choices = []
+    for i, model_info in enumerate(models_list, 1):
+        choice = f"路徑: {model_info['model_path']} | 方法: {model_info['tuning_method']} | 時間: {model_info['timestamp']}"
+        model_choices.append(choice)
+    
+    return model_choices
 
 # ============================================================================
 # Gradio 介面部分 - 修改輸出為三個格子
@@ -851,11 +1005,24 @@ with gr.Blocks(title="BERT 完整訓練與預測平台", theme=gr.themes.Soft())
         gr.Markdown("""
         ### 使用訓練好的模型進行預測
         
-        在訓練頁面完成訓練後，可以在這裡輸入病歷文本進行預測。
+        選擇已訓練的模型，輸入病歷文本進行預測。會同時顯示未微調和微調模型的預測結果以供比較。
         """)
         
         with gr.Row():
             with gr.Column():
+                # 模型選擇下拉選單
+                model_dropdown = gr.Dropdown(
+                    label="選擇模型",
+                    choices=["請先訓練模型"],
+                    value="請先訓練模型",
+                    info="選擇要使用的已訓練模型"
+                )
+                
+                refresh_button = gr.Button(
+                    "🔄 重新整理模型列表",
+                    size="sm"
+                )
+                
                 text_input = gr.Textbox(
                     label="輸入病歷文本",
                     placeholder="請輸入患者的病歷描述（英文）...",
@@ -881,9 +1048,18 @@ with gr.Blocks(title="BERT 完整訓練與預測平台", theme=gr.themes.Soft())
                 )
             
             with gr.Column():
-                prediction_output = gr.Markdown(
-                    label="預測結果",
-                    value="請先完成模型訓練，然後輸入文本進行預測..."
+                gr.Markdown("### 預測結果比較")
+                
+                # 上框：未微調 BERT 預測結果
+                baseline_prediction_output = gr.Markdown(
+                    label="未微調 BERT",
+                    value="等待預測..."
+                )
+                
+                # 下框：微調 BERT 預測結果
+                finetuned_prediction_output = gr.Markdown(
+                    label="微調 BERT",
+                    value="等待預測..."
                 )
     
     with gr.Tab("📖 使用說明"):
@@ -936,7 +1112,7 @@ with gr.Blocks(title="BERT 完整訓練與預測平台", theme=gr.themes.Soft())
         outputs=[lora_params, adalora_params]
     )
     
-    # 設定按鈕動作 - 注意這裡改為三個輸出
+    # 設定訓練按鈕動作 - 注意這裡改為三個輸出
     train_button.click(
         fn=train_wrapper,
         inputs=[
@@ -963,10 +1139,21 @@ with gr.Blocks(title="BERT 完整訓練與預測平台", theme=gr.themes.Soft())
         outputs=[data_info_output, baseline_output, finetuned_output]  # 三個輸出
     )
     
+    # 重新整理模型列表按鈕
+    def refresh_models():
+        return gr.update(choices=get_available_models(), value=get_available_models()[0])
+    
+    refresh_button.click(
+        fn=refresh_models,
+        inputs=[],
+        outputs=[model_dropdown]
+    )
+    
+    # 預測按鈕動作 - 兩個輸出：未微調和微調
     predict_button.click(
         fn=predict_text,
-        inputs=[text_input],
-        outputs=[prediction_output]
+        inputs=[model_dropdown, text_input],
+        outputs=[baseline_prediction_output, finetuned_prediction_output]
     )
 
 if __name__ == "__main__":