app.py

import gradio as gr
import pandas as pd
import torch
from datasets import Dataset, DatasetDict
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer,
    DataCollatorWithPadding
)
from peft import (
    LoraConfig, 
    AdaLoraConfig,
    AdaptionPromptConfig,
    PromptTuningConfig,
    PrefixTuningConfig,
    get_peft_model, 
    TaskType, 
    PeftModel
)
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from sklearn.utils import resample
import numpy as np
import json
from datetime import datetime
import os
import gc
from huggingface_hub import login

# ==================== 全域變數 ====================
LAST_MODEL_PATH = None
LAST_TOKENIZER = None
MAX_LENGTH = 512

# ==================== HF Token 登入 ====================
print("🔐 檢查 Hugging Face Token...")
if "HF_TOKEN" in os.environ:
    try:
        login(token=os.environ["HF_TOKEN"])
        print("✅ 已使用 HF Token 登入")
    except Exception as e:
        print(f"⚠️ Token 登入失敗: {e}")
else:
    print("⚠️ 未找到 HF_TOKEN,可能無法下載 Llama 模型")

# 檢測設備
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"🖥️ 使用設備: {device}")

# ==================== 核心訓練函數(你的原始邏輯 - 完全不動) ====================
def run_llama_training(
    file_path,
    model_name,
    target_samples,
    use_class_weights,
    num_epochs,
    batch_size,
    learning_rate,
    tuning_method,
    lora_r,
    lora_alpha,
    lora_dropout,
    lora_target_modules,
    adalora_init_r,
    adalora_target_r,
    adalora_alpha,
    adalora_tinit,
    adalora_tfinal,
    adalora_delta_t,
    adapter_reduction_factor,
    prompt_tuning_num_tokens,
    prefix_tuning_num_tokens,
    best_metric,
    # 【新增】二次微調參數
    is_second_finetuning=False,
    base_model_path=None
):
    """
    你的原始 Llama 訓練邏輯
    """
    
    global LAST_MODEL_PATH, LAST_TOKENIZER
    
    # ==================== 清空記憶體(訓練前) ====================
    torch.cuda.empty_cache()
    gc.collect()
    print("🧹 記憶體已清空")
    
    # ==================== 1. 載入數據 ====================
    training_type = "二次微調" if is_second_finetuning else "第一次微調"
    
    print("\n" + "="*80)
    print(f"🦙 Llama NBCD {training_type} - {tuning_method} 方法")
    print("="*80)
    print(f"開始時間: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"訓練類型: {training_type}")
    print(f"微調方法: {tuning_method}")
    if is_second_finetuning:
        print(f"基礎模型: {base_model_path}")
    print("="*80)
    
    print("📂 載入訓練數據...")
    df = pd.read_csv(file_path)
    print(f"✅ 成功載入 {len(df)} 筆數據")
    
    # 自動偵測文本和標籤欄位
    text_col = None
    label_col = None
    
    # 支持的文本欄位名稱
    if 'Text' in df.columns:
        text_col = 'Text'
    elif 'text' in df.columns:
        text_col = 'text'
    
    # 支持的標籤欄位名稱
    if 'Label' in df.columns:
        label_col = 'Label'
    elif 'label' in df.columns:
        label_col = 'label'
    
    if text_col is None or label_col is None:
        raise ValueError(
            f"❌ 無法偵測到正確的欄位名稱!\n"
            f"📋 您的 CSV 欄位: {list(df.columns)}\n\n"
            f"✅ 請使用以下欄位名稱:\n"
            f"   文本欄位: 'Text' 或 'text'\n"
            f"   標籤欄位: 'Label' 或 'label'"
        )
    
    print(f"   ✅ 偵測到文本欄位: '{text_col}'")
    print(f"   ✅ 偵測到標籤欄位: '{label_col}'")
    
    # 統一重命名為標準欄位名
    df = df.rename(columns={text_col: 'Text', label_col: 'nbcd'})
    
    print(f"   原始 Class 0: {(df['nbcd']==0).sum()} 筆")
    print(f"   原始 Class 1: {(df['nbcd']==1).sum()} 筆")
    
    # ==================== 2. 資料平衡處理 ====================
    print("\n⚖️ 執行資料平衡...")
    
    df_class_0 = df[df['nbcd'] == 0]
    df_class_1 = df[df['nbcd'] == 1]
    
    target_n = int(target_samples)
    
    # 欠採樣 Class 0
    if len(df_class_0) > target_n:
        df_class_0_balanced = resample(df_class_0, n_samples=target_n, random_state=42, replace=False)
        print(f"✅ Class 0 欠採樣: {len(df_class_0)} → {len(df_class_0_balanced)} 筆")
    else:
        df_class_0_balanced = df_class_0
        print(f"⚠️ Class 0 樣本數不足,保持 {len(df_class_0)} 筆")
    
    # 過採樣 Class 1
    if len(df_class_1) < target_n:
        df_class_1_balanced = resample(df_class_1, n_samples=target_n, random_state=42, replace=True)
        print(f"✅ Class 1 過採樣: {len(df_class_1)} → {len(df_class_1_balanced)} 筆")
    else:
        df_class_1_balanced = df_class_1
        print(f"⚠️ Class 1 樣本數充足,保持 {len(df_class_1)} 筆")
    
    df_balanced = pd.concat([df_class_0_balanced, df_class_1_balanced])
    df_balanced = df_balanced.sample(frac=1, random_state=42).reset_index(drop=True)
    
    print(f"\n📊 平衡後數據:")
    print(f"   總樣本數: {len(df_balanced)} 筆")
    print(f"   Class 0: {(df_balanced['nbcd']==0).sum()} 筆")
    print(f"   Class 1: {(df_balanced['nbcd']==1).sum()} 筆")
    
    # ==================== 3. 計算類別權重 ====================
    if use_class_weights:
        print("\n⚖️ 計算類別權重...")
        class_counts = df_balanced['nbcd'].value_counts().sort_index()
        total = len(df_balanced)
        num_classes = 2
        
        class_weight_0 = total / (num_classes * class_counts[0])
        class_weight_1 = total / (num_classes * class_counts[1])
        class_weights = torch.tensor([class_weight_0, class_weight_1], dtype=torch.float32)
        
        print(f"✅ 類別權重計算完成:")
        print(f"   Class 0 權重: {class_weight_0:.4f}")
        print(f"   Class 1 權重: {class_weight_1:.4f}")
        
        if device == "cuda":
            class_weights = class_weights.to(device)
    else:
        class_weights = None
        print("\n⚠️ 未使用類別權重")
    
    # ==================== 4. 分割數據 ====================
    print("\n✂️ 分割訓練集和測試集...")
    train_df, test_df = train_test_split(
        df_balanced,
        test_size=0.2,
        stratify=df_balanced['nbcd'],
        random_state=42
    )
    print(f"✅ 訓練集: {len(train_df)} 筆 (Class 0: {(train_df['nbcd']==0).sum()}, Class 1: {(train_df['nbcd']==1).sum()})")
    print(f"✅ 測試集: {len(test_df)} 筆 (Class 0: {(test_df['nbcd']==0).sum()}, Class 1: {(test_df['nbcd']==1).sum()})")
    
    dataset = DatasetDict({
        'train': Dataset.from_pandas(train_df[['Text', 'nbcd']]),
        'test': Dataset.from_pandas(test_df[['Text', 'nbcd']])
    })
    
    # ==================== 5. 載入模型和 Tokenizer ====================
    print("\n🤖 載入 Llama 模型和 Tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id
    
    # ==================== 6. 載入未微調的基礎模型 (Baseline) ====================
    print("\n📦 載入未微調的基礎模型 (Baseline)...")
    baseline_model = AutoModelForSequenceClassification.from_pretrained(
        model_name,
        num_labels=2,
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        device_map="auto" if device == "cuda" else None
    )
    baseline_model.config.pad_token_id = tokenizer.pad_token_id
    print("✅ Baseline 模型載入完成")
    
    # ==================== 7. 載入要微調的模型 ====================
    print("\n🔧 載入用於微調的模型...")
    
    # 【新增】二次微調邏輯
    if is_second_finetuning and base_model_path:
        print(f"📦 載入第一次微調模型: {base_model_path}")
        
        # 讀取第一次模型資訊
        with open('./saved_llama_models_list.json', 'r') as f:
            models_list = json.load(f)
        
        base_model_info = None
        for model_info in models_list:
            if model_info['model_path'] == base_model_path:
                base_model_info = model_info
                break
        
        if base_model_info is None:
            raise ValueError(f"找不到基礎模型資訊: {base_model_path}")
        
        base_tuning_method = base_model_info['tuning_method']
        print(f"   第一次微調方法: {base_tuning_method}")
        
        # 根據第一次的方法載入模型
        if base_tuning_method in ["LoRA", "AdaLoRA", "Adapter", "Prompt Tuning"]:
            # 載入 PEFT 模型
            base_bert = AutoModelForSequenceClassification.from_pretrained(
                model_name,
                num_labels=2,
                torch_dtype=torch.float16 if device == "cuda" else torch.float32
            )
            base_model = PeftModel.from_pretrained(base_bert, base_model_path)
            print(f"   ✅ 已載入 {base_tuning_method} 模型")
        else:
            # 載入一般模型 (BitFit)
            base_model = AutoModelForSequenceClassification.from_pretrained(
                base_model_path,
                num_labels=2,
                torch_dtype=torch.float16 if device == "cuda" else torch.float32
            )
            print(f"   ✅ 已載入 BitFit 模型")
        
        if device == "cuda":
            base_model = base_model.to(device)
        
        print(f"   ⚠️ 注意:二次微調將使用與第一次相同的方法 ({base_tuning_method})")
        
        # 二次微調時強制使用相同方法
        tuning_method = base_tuning_method
        
    else:
        # 【原始邏輯】第一次微調:從純 Llama 開始
        base_model = AutoModelForSequenceClassification.from_pretrained(
            model_name,
            num_labels=2,
            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
            device_map="auto" if device == "cuda" else None
        )
    
    base_model.config.pad_token_id = tokenizer.pad_token_id
    print("✅ 基礎模型載入完成")
    
    # ==================== 8. 配置微調方法 ====================
    print(f"\n🔧 配置 {tuning_method}...")
    
    if tuning_method == "LoRA":
        # LoRA 配置 - 使用完整參數
        target_modules_map = {
            "query,value": ["q_proj", "v_proj"],
            "query,key,value": ["q_proj", "k_proj", "v_proj"],
            "all": ["q_proj", "k_proj", "v_proj", "o_proj"]
        }
        
        peft_config = LoraConfig(
            task_type=TaskType.SEQ_CLS,
            r=int(lora_r),
            lora_alpha=int(lora_alpha),
            lora_dropout=float(lora_dropout),
            target_modules=target_modules_map.get(lora_target_modules, ["q_proj", "v_proj"]),
            bias="none"
        )
        print(f"✅ LoRA 配置完成")
        print(f"   LoRA rank (r): {lora_r}")
        print(f"   LoRA alpha: {lora_alpha}")
        print(f"   LoRA dropout: {lora_dropout}")
        print(f"   目標模組: {lora_target_modules}")
    
    elif tuning_method == "AdaLoRA":
        # AdaLoRA 配置 - 使用獨立參數
        try:
            peft_config = AdaLoraConfig(
                task_type=TaskType.SEQ_CLS,
                inference_mode=False,
                r=int(adalora_target_r),
                lora_alpha=int(adalora_alpha),
                lora_dropout=0.1,
                target_modules=["q_proj", "v_proj"],
                # AdaLoRA 特定參數
                init_r=int(adalora_init_r),
                target_r=int(adalora_target_r),
                tinit=int(adalora_tinit),
                tfinal=int(adalora_tfinal),
                deltaT=int(adalora_delta_t),
            )
            print(f"✅ AdaLoRA 配置完成")
            print(f"   初始 rank: {adalora_init_r}")
            print(f"   目標 rank: {adalora_target_r}")
            print(f"   Alpha: {adalora_alpha}")
            print(f"   Tinit: {adalora_tinit}, Tfinal: {adalora_tfinal}")
            print(f"   Delta T: {adalora_delta_t}")
            print(f"   自適應秩調整: 啟用")
        except Exception as e:
            print(f"⚠️ AdaLoRA 配置失敗,回退到 LoRA: {e}")
            peft_config = LoraConfig(
                task_type=TaskType.SEQ_CLS,
                r=int(adalora_target_r),
                lora_alpha=int(adalora_alpha),
                lora_dropout=0.1,
                target_modules=["q_proj", "v_proj"],
                bias="none"
            )
        
    elif tuning_method == "Adapter":
        # Adapter (Bottleneck Adapters)
        peft_config = AdaptionPromptConfig(
            task_type=TaskType.SEQ_CLS,
            adapter_len=10,
            adapter_layers=30,
            reduction_factor=int(adapter_reduction_factor)
        )
        print(f"✅ Adapter 配置完成")
        print(f"   Reduction factor: {adapter_reduction_factor}")
        
    elif tuning_method == "Prompt Tuning":
        # Soft Prompt Tuning
        peft_config = PromptTuningConfig(
            task_type=TaskType.SEQ_CLS,
            num_virtual_tokens=int(prompt_tuning_num_tokens),
            prompt_tuning_init="TEXT",
            prompt_tuning_init_text="Classify if the following text indicates NBCD:",
            tokenizer_name_or_path=model_name
        )
        print(f"✅ Prompt Tuning 配置完成")
        print(f"   Virtual tokens: {prompt_tuning_num_tokens}")
        
    elif tuning_method == "Prefix Tuning":
        # Prefix Tuning - 可能有兼容性問題,但仍然嘗試
        print(f"⚠️ Prefix Tuning 在某些環境可能有兼容性問題")
        print(f"   如果遇到錯誤,建議使用 Prompt Tuning 替代")
        
        try:
            # 先禁用模型的緩存功能
            base_model.config.use_cache = False
            
            peft_config = PrefixTuningConfig(
                task_type=TaskType.SEQ_CLS,
                num_virtual_tokens=int(prefix_tuning_num_tokens),
                prefix_projection=False,
                inference_mode=False
            )
            print(f"✅ Prefix Tuning 配置完成")
            print(f"   Virtual tokens: {prefix_tuning_num_tokens}")
            print(f"   已禁用緩存")
        except Exception as e:
            print(f"❌ Prefix Tuning 配置失敗: {e}")
            raise ValueError(
                f"Prefix Tuning 配置失敗,原因: {e}\n"
                f"建議使用 Prompt Tuning 作為替代方案"
            )
        
    elif tuning_method == "BitFit":
        # BitFit: 只訓練 bias 參數 - 完全修復版
        model = base_model
        
        # 凍結所有參數
        for param in model.parameters():
            param.requires_grad = False
        
        # 只解凍 bias 和 分類頭
        trainable_params_list = []
        for name, param in model.named_parameters():
            if 'bias' in name or 'score' in name or 'classifier' in name:
                param.requires_grad = True
                trainable_params_list.append(name)
        
        print(f"✅ BitFit 配置完成")
        print(f"   僅訓練 bias 和分類頭參數")
        print(f"   可訓練參數: {', '.join(trainable_params_list[:5])}...")
    
    # 應用 PEFT 配置(BitFit 除外)
    if tuning_method != "BitFit":
        model = get_peft_model(base_model, peft_config)
        
        # Prefix Tuning 額外設置
        if tuning_method == "Prefix Tuning":
            model.config.use_cache = False
    
    # 計算可訓練參數
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in model.parameters())
    print(f"   可訓練參數: {trainable_params:,} / {total_params:,} ({trainable_params/total_params*100:.2f}%)")
    
    # ==================== 9. 預處理數據 ====================
    print("\n📄 預處理數據...")
    
    def preprocess_function(examples):
        return tokenizer(
            examples['Text'],
            truncation=True,
            padding='max_length',
            max_length=MAX_LENGTH
        )
    
    tokenized_dataset = dataset.map(preprocess_function, batched=True, remove_columns=['Text'])
    tokenized_dataset = tokenized_dataset.rename_column("nbcd", "labels")
    print("✅ 數據預處理完成")
    
    # ==================== 10. 評估指標函數 ====================
    def compute_metrics(eval_pred):
        predictions, labels = eval_pred
        predictions = np.argmax(predictions, axis=1)
        
        accuracy = accuracy_score(labels, predictions)
        precision, recall, f1, _ = precision_recall_fscore_support(
            labels, predictions, average='binary', zero_division=0
        )
        
        # 計算混淆矩陣以得到 sensitivity 和 specificity
        from sklearn.metrics import confusion_matrix
        cm = confusion_matrix(labels, predictions)
        
        if cm.shape == (2, 2):
            tn, fp, fn, tp = cm.ravel()
            sensitivity = tp / (tp + fn) if (tp + fn) > 0 else 0  # 敏感度 = Recall
            specificity = tn / (tn + fp) if (tn + fp) > 0 else 0  # 特異性
        else:
            sensitivity = 0
            specificity = 0
        
        return {
            'accuracy': accuracy,
            'precision': precision,
            'recall': recall,
            'f1': f1,
            'sensitivity': sensitivity,
            'specificity': specificity
        }
    
    # ==================== 11. 評估 Baseline 模型 ====================
    # 【僅第一次微調時執行】
    if not is_second_finetuning:
        print("\n" + "="*70)
        print("📊 評估未微調的 Baseline 模型...")
        print("="*70)
        
        baseline_trainer = Trainer(
            model=baseline_model,
            args=TrainingArguments(
                output_dir="./temp_baseline_llama",
                per_device_eval_batch_size=int(batch_size),
                bf16=(device == "cuda"),
                report_to="none"
            ),
            tokenizer=tokenizer,
            data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
            compute_metrics=compute_metrics
        )
        
        baseline_test_results = baseline_trainer.evaluate(eval_dataset=tokenized_dataset['test'])
        
        print("\n📋 Baseline 模型 - 測試集結果:")
        print(f"  Accuracy:    {baseline_test_results['eval_accuracy']:.4f}")
        print(f"  Precision:   {baseline_test_results['eval_precision']:.4f}")
        print(f"  Recall:      {baseline_test_results['eval_recall']:.4f}")
        print(f"  F1 Score:    {baseline_test_results['eval_f1']:.4f}")
        print(f"  Sensitivity: {baseline_test_results['eval_sensitivity']:.4f}")
        print(f"  Specificity: {baseline_test_results['eval_specificity']:.4f}")
        
        # 清空 baseline 模型記憶體
        del baseline_model
        del baseline_trainer
        torch.cuda.empty_cache()
        gc.collect()
    else:
        # 二次微調不評估 baseline
        baseline_test_results = None
        del baseline_model
        torch.cuda.empty_cache()
        gc.collect()
    
    # ==================== 12. 自定義 Trainer ====================
    if use_class_weights:
        class WeightedTrainer(Trainer):
            def __init__(self, *args, class_weights=None, **kwargs):
                super().__init__(*args, **kwargs)
                self.class_weights = class_weights
            
            def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
                labels = inputs.pop("labels")
                outputs = model(**inputs)
                logits = outputs.logits
                
                loss_fct = torch.nn.CrossEntropyLoss(weight=self.class_weights)
                loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
                
                return (loss, outputs) if return_outputs else loss
        
        TrainerClass = WeightedTrainer
    else:
        TrainerClass = Trainer
    
    # ==================== 13. 訓練配置 ====================
    print("\n" + "="*70)
    print("⚙️ 配置微調訓練器...")
    print("="*70)
    
    # 指標映射
    metric_map = {
        "f1": "f1",
        "accuracy": "accuracy",
        "precision": "precision",
        "recall": "recall",
        "sensitivity": "sensitivity",
        "specificity": "specificity"
    }
    
    training_label = "second" if is_second_finetuning else "first"
    output_dir = f'./llama_nbcd_{tuning_method.lower().replace(" ", "_")}_{training_label}_{datetime.now().strftime("%Y%m%d_%H%M%S")}'
    
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=int(num_epochs),
        per_device_train_batch_size=int(batch_size),
        per_device_eval_batch_size=int(batch_size),
        learning_rate=float(learning_rate),
        weight_decay=0.01,
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model=metric_map.get(best_metric, "recall"),
        logging_dir=f"{output_dir}/logs",
        logging_steps=10,
        bf16=(device == "cuda"),
        gradient_accumulation_steps=2,
        warmup_steps=50,
        report_to="none",
        seed=42
    )
    
    if use_class_weights:
        trainer = TrainerClass(
            model=model,
            args=training_args,
            train_dataset=tokenized_dataset['train'],
            eval_dataset=tokenized_dataset['test'],
            tokenizer=tokenizer,
            data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
            compute_metrics=compute_metrics,
            class_weights=class_weights
        )
    else:
        trainer = TrainerClass(
            model=model,
            args=training_args,
            train_dataset=tokenized_dataset['train'],
            eval_dataset=tokenized_dataset['test'],
            tokenizer=tokenizer,
            data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
            compute_metrics=compute_metrics
        )
    
    # ==================== 14. 開始訓練 ====================
    print("\n" + "="*70)
    print(f"🚀 開始{training_type}訓練...")
    print("="*70 + "\n")
    
    start_time = datetime.now()
    train_result = trainer.train()
    end_time = datetime.now()
    duration = (end_time - start_time).total_seconds() / 60
    
    print("\n" + "="*70)
    print(f"✅ 訓練完成!")
    print(f"   耗時: {duration:.1f} 分鐘")
    print("="*70)
    
    # ==================== 15. 評估微調後的模型 ====================
    print("\n" + "="*70)
    print(f"📊 評估{training_type}後的模型...")
    print("="*70)
    
    finetuned_test_results = trainer.evaluate(eval_dataset=tokenized_dataset['test'])
    
    print(f"\n📋 {training_type}模型 - 測試集結果:")
    print(f"  Accuracy:    {finetuned_test_results['eval_accuracy']:.4f}")
    print(f"  Precision:   {finetuned_test_results['eval_precision']:.4f}")
    print(f"  Recall:      {finetuned_test_results['eval_recall']:.4f}")
    print(f"  F1 Score:    {finetuned_test_results['eval_f1']:.4f}")
    print(f"  Sensitivity: {finetuned_test_results['eval_sensitivity']:.4f}")
    print(f"  Specificity: {finetuned_test_results['eval_specificity']:.4f}")
    
    # ==================== 16. 保存模型和結果 ====================
    print("\n💾 保存模型和結果...")
    trainer.save_model()
    tokenizer.save_pretrained(output_dir)
    
    # 儲存模型資訊到 JSON 檔案
    metric_key = 'eval_' + metric_map.get(best_metric, "recall")
    model_info = {
        'model_path': output_dir,
        'model_name': model_name,
        'tuning_method': tuning_method,
        'training_type': training_type,
        'best_metric': best_metric,
        'best_metric_value': float(finetuned_test_results[metric_key]),
        'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
        'target_samples': target_samples,
        'epochs': num_epochs,
        'batch_size': batch_size,
        'learning_rate': learning_rate,
        'lora_r': lora_r if tuning_method in ["LoRA", "AdaLoRA"] else None,
        'lora_alpha': lora_alpha if tuning_method in ["LoRA", "AdaLoRA"] else None,
        'is_second_finetuning': is_second_finetuning,
        'base_model_path': base_model_path if is_second_finetuning else None
    }
    
    # 讀取現有的模型列表
    models_list_file = './saved_llama_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 = output_dir
    LAST_TOKENIZER = tokenizer
    
    print(f"✅ 模型已儲存至: {output_dir}")
    
    # ==================== 清空記憶體(訓練後) ====================
    del model
    del trainer
    torch.cuda.empty_cache()
    gc.collect()
    print("🧹 訓練後記憶體已清空")
    
    # 準備返回結果
    results = {
        'baseline_results': baseline_test_results,
        'finetuned_results': finetuned_test_results,
        'model_path': output_dir,
        'duration': duration,
        'best_metric': best_metric,
        'model_name': model_name,
        'tuning_method': tuning_method,
        'training_type': training_type,
        'is_second_finetuning': is_second_finetuning
    }
    
    return results

# ==================== Gradio Wrapper 函數 ====================
def train_first_wrapper(
    file,
    model_name,
    target_samples,
    use_class_weights,
    num_epochs,
    batch_size,
    learning_rate,
    tuning_method,
    lora_r,
    lora_alpha,
    lora_dropout,
    lora_target_modules,
    adalora_init_r,
    adalora_target_r,
    adalora_alpha,
    adalora_tinit,
    adalora_tfinal,
    adalora_delta_t,
    adapter_reduction_factor,
    prompt_tuning_num_tokens,
    prefix_tuning_num_tokens,
    best_metric
):
    """第一次微調的包裝函數"""
    
    if file is None:
        return "請上傳 CSV 檔案", "", ""
    
    try:
        # 呼叫訓練函數
        results = run_llama_training(
            file_path=file.name,
            model_name=model_name,
            target_samples=target_samples,
            use_class_weights=use_class_weights,
            num_epochs=num_epochs,
            batch_size=batch_size,
            learning_rate=learning_rate,
            tuning_method=tuning_method,
            lora_r=lora_r,
            lora_alpha=lora_alpha,
            lora_dropout=lora_dropout,
            lora_target_modules=lora_target_modules,
            adalora_init_r=adalora_init_r,
            adalora_target_r=adalora_target_r,
            adalora_alpha=adalora_alpha,
            adalora_tinit=adalora_tinit,
            adalora_tfinal=adalora_tfinal,
            adalora_delta_t=adalora_delta_t,
            adapter_reduction_factor=adapter_reduction_factor,
            prompt_tuning_num_tokens=prompt_tuning_num_tokens,
            prefix_tuning_num_tokens=prefix_tuning_num_tokens,
            best_metric=best_metric,
            is_second_finetuning=False
        )
        
        baseline_results = results['baseline_results']
        finetuned_results = results['finetuned_results']
        
        # 第一格:資料資訊
        data_info = f"""
# 📊 資料資訊 (第一次微調)

## 🔧 訓練配置
- **模型**: {results['model_name']}
- **微調方法**: {results['tuning_method']}
- **最佳化指標**: {results['best_metric']}
- **訓練時長**: {results['duration']:.1f} 分鐘

## ⚙️ 訓練參數
- **目標樣本數**: {target_samples} 筆/類別
- **使用類別權重**: {'是' if use_class_weights else '否'}
- **訓練輪數**: {num_epochs}
- **批次大小**: {batch_size}
- **學習率**: {learning_rate}

✅ 第一次微調完成!可進行二次微調或預測!
        """
        
        # 第二格:未微調 Llama
        baseline_output = f"""
# 🔵 未微調 Llama (Baseline)
## 未經訓練

### 📈 評估指標

| 指標 | 數值 |
|------|------|
| **Accuracy** | {baseline_results['eval_accuracy']:.4f} |
| **Precision** | {baseline_results['eval_precision']:.4f} |
| **Recall** | {baseline_results['eval_recall']:.4f} |
| **F1 Score** | {baseline_results['eval_f1']:.4f} |
| **Sensitivity** | {baseline_results['eval_sensitivity']:.4f} |
| **Specificity** | {baseline_results['eval_specificity']:.4f} |
        """
        
        # 第三格:微調後 Llama
        finetuned_output = f"""
# 🟢 第一次微調 Llama
## {results['tuning_method']}

### 📈 評估指標

| 指標 | 數值 |
|------|------|
| **Accuracy** | {finetuned_results['eval_accuracy']:.4f} |
| **Precision** | {finetuned_results['eval_precision']:.4f} |
| **Recall** | {finetuned_results['eval_recall']:.4f} |
| **F1 Score** | {finetuned_results['eval_f1']:.4f} |
| **Sensitivity** | {finetuned_results['eval_sensitivity']:.4f} |
| **Specificity** | {finetuned_results['eval_specificity']:.4f} |
        """
        
        return data_info, baseline_output, finetuned_output
        
    except Exception as e:
        import traceback
        error_msg = f"❌ 錯誤:{str(e)}\n\n詳細錯誤訊息:\n{traceback.format_exc()}"
        return error_msg, "", ""

def train_second_wrapper(
    base_model_choice,
    file,
    target_samples,
    use_class_weights,
    num_epochs,
    batch_size,
    learning_rate,
    best_metric
):
    """二次微調的包裝函數"""
    
    if base_model_choice == "請先進行第一次微調":
        return "請先在「第一次微調」頁面訓練模型", ""
    
    if file is None:
        return "請上傳新的訓練數據 CSV 檔案", ""
    
    try:
        # 解析基礎模型路徑
        base_model_path = base_model_choice
        
        # 讀取第一次模型資訊
        with open('./saved_llama_models_list.json', 'r') as f:
            models_list = json.load(f)
        
        base_model_info = None
        for model_info in models_list:
            if model_info['model_path'] == base_model_path:
                base_model_info = model_info
                break
        
        if base_model_info is None:
            return "找不到基礎模型資訊", ""
        
        # 使用第一次的參數(二次微調不更改方法)
        tuning_method = base_model_info['tuning_method']
        model_name = base_model_info['model_name']
        
        # 獲取第一次的 PEFT 參數
        lora_r = base_model_info.get('lora_r', 16)
        lora_alpha = base_model_info.get('lora_alpha', 32)
        lora_dropout = 0.1
        lora_target_modules = "query,value"
        adalora_init_r = 12
        adalora_target_r = 8
        adalora_alpha = 32
        adalora_tinit = 0
        adalora_tfinal = 0
        adalora_delta_t = 1
        adapter_reduction_factor = 16
        prompt_tuning_num_tokens = 20
        prefix_tuning_num_tokens = 30
        
        results = run_llama_training(
            file_path=file.name,
            model_name=model_name,
            target_samples=target_samples,
            use_class_weights=use_class_weights,
            num_epochs=num_epochs,
            batch_size=batch_size,
            learning_rate=learning_rate,
            tuning_method=tuning_method,
            lora_r=lora_r,
            lora_alpha=lora_alpha,
            lora_dropout=lora_dropout,
            lora_target_modules=lora_target_modules,
            adalora_init_r=adalora_init_r,
            adalora_target_r=adalora_target_r,
            adalora_alpha=adalora_alpha,
            adalora_tinit=adalora_tinit,
            adalora_tfinal=adalora_tfinal,
            adalora_delta_t=adalora_delta_t,
            adapter_reduction_factor=adapter_reduction_factor,
            prompt_tuning_num_tokens=prompt_tuning_num_tokens,
            prefix_tuning_num_tokens=prefix_tuning_num_tokens,
            best_metric=best_metric,
            is_second_finetuning=True,
            base_model_path=base_model_path
        )
        
        finetuned_results = results['finetuned_results']
        
        data_info = f"""
# 📊 二次微調結果

## 🔧 訓練配置
- **基礎模型**: {base_model_path}
- **微調方法**: {results['tuning_method']} (繼承自第一次)
- **最佳化指標**: {results['best_metric']}
- **最佳指標值**: {finetuned_results['eval_' + results['best_metric']]:.4f}
- **訓練時長**: {results['duration']:.1f} 分鐘

## ⚙️ 訓練參數
- **目標樣本數**: {target_samples} 筆/類別
- **使用類別權重**: {'是' if use_class_weights else '否'}
- **訓練輪數**: {num_epochs}
- **批次大小**: {batch_size}
- **學習率**: {learning_rate}

✅ 二次微調完成!可進行預測!
        """
        
        finetuned_output = f"""
# 🟢 二次微調 Llama
## {results['tuning_method']}

### 📈 評估指標

| 指標 | 數值 |
|------|------|
| **Accuracy** | {finetuned_results['eval_accuracy']:.4f} |
| **Precision** | {finetuned_results['eval_precision']:.4f} |
| **Recall** | {finetuned_results['eval_recall']:.4f} |
| **F1 Score** | {finetuned_results['eval_f1']:.4f} |
| **Sensitivity** | {finetuned_results['eval_sensitivity']:.4f} |
| **Specificity** | {finetuned_results['eval_specificity']:.4f} |
        """
        
        return data_info, finetuned_output
        
    except Exception as e:
        import traceback
        error_msg = f"❌ 錯誤:{str(e)}\n\n詳細錯誤訊息:\n{traceback.format_exc()}"
        return error_msg, ""

# ==================== 新增:新數據測試函數 ====================

def test_on_new_data(test_file_path, baseline_choice, first_choice, second_choice):
    """
    在新測試數據上比較三個模型的表現:
    1. 純 Llama (baseline)
    2. 第一次微調模型
    3. 第二次微調模型
    """
    
    print("\n" + "=" * 80)
    print("📊 新數據測試 - 三模型比較")
    print("=" * 80)
    
    # 載入測試數據
    df_test = pd.read_csv(test_file_path)
    
    # 自動偵測欄位
    text_col = 'Text' if 'Text' in df_test.columns else 'text'
    label_col = 'Label' if 'Label' in df_test.columns else 'label'
    
    df_clean = pd.DataFrame({
        'text': df_test[text_col],
        'label': df_test[label_col]
    })
    df_clean = df_clean.dropna()
    
    print(f"\n測試數據:")
    print(f"  總筆數: {len(df_clean)}")
    print(f"  Class 0: {sum(df_clean['label']==0)} 筆")
    print(f"  Class 1: {sum(df_clean['label']==1)} 筆")
    
    # 準備測試數據
    test_dataset = Dataset.from_pandas(df_clean[['text', 'label']])
    
    # 評估函數
    def evaluate_model(model, tokenizer, model_name_str, dataset_name):
        model.eval()
        
        # 確保 tokenizer 有 pad_token
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
            tokenizer.pad_token_id = tokenizer.eos_token_id
        
        # 確保模型配置也有 pad_token_id
        if hasattr(model, 'config'):
            model.config.pad_token_id = tokenizer.pad_token_id
        
        def preprocess_function(examples):
            return tokenizer(examples['text'], truncation=True, padding='max_length', max_length=MAX_LENGTH)
        
        test_tokenized = test_dataset.map(preprocess_function, batched=True)
        
        trainer_args = TrainingArguments(
            output_dir='./temp_test',
            per_device_eval_batch_size=32,
            report_to="none"
        )
        
        def compute_metrics_test(eval_pred):
            predictions, labels = eval_pred
            predictions = np.argmax(predictions, axis=1)
            
            accuracy = accuracy_score(labels, predictions)
            precision, recall, f1, _ = precision_recall_fscore_support(
                labels, predictions, average='binary', zero_division=0
            )
            
            from sklearn.metrics import confusion_matrix
            cm = confusion_matrix(labels, predictions)
            
            if cm.shape == (2, 2):
                tn, fp, fn, tp = cm.ravel()
                sensitivity = tp / (tp + fn) if (tp + fn) > 0 else 0
                specificity = tn / (tn + fp) if (tn + fp) > 0 else 0
            else:
                sensitivity = 0
                specificity = 0
                tn = fp = fn = tp = 0
            
            return {
                'accuracy': accuracy,
                'precision': precision,
                'recall': recall,
                'f1': f1,
                'sensitivity': sensitivity,
                'specificity': specificity,
                'tp': int(tp),
                'tn': int(tn),
                'fp': int(fp),
                'fn': int(fn)
            }
        
        trainer = Trainer(
            model=model,
            args=trainer_args,
            compute_metrics=compute_metrics_test,
            data_collator=DataCollatorWithPadding(tokenizer=tokenizer)
        )
        
        predictions_output = trainer.predict(test_tokenized)
        
        results = {
            'accuracy': predictions_output.metrics['test_accuracy'],
            'precision': predictions_output.metrics['test_precision'],
            'recall': predictions_output.metrics['test_recall'],
            'f1': predictions_output.metrics['test_f1'],
            'sensitivity': predictions_output.metrics['test_sensitivity'],
            'specificity': predictions_output.metrics['test_specificity'],
            'tp': predictions_output.metrics['test_tp'],
            'tn': predictions_output.metrics['test_tn'],
            'fp': predictions_output.metrics['test_fp'],
            'fn': predictions_output.metrics['test_fn']
        }
        
        print(f"\n✅ {dataset_name} 評估完成")
        
        del trainer
        torch.cuda.empty_cache()
        gc.collect()
        
        return results
    
    all_results = {}
    
    # 1. 評估純 Llama
    if baseline_choice == "評估純 Llama":
        print("\n" + "-" * 80)
        print("1️⃣ 評估純 Llama (Baseline)")
        print("-" * 80)
        
        # 獲取模型名稱
        if first_choice != "請選擇":
            with open('./saved_llama_models_list.json', 'r') as f:
                models_list = json.load(f)
            for model_info in models_list:
                if model_info['model_path'] == first_choice:
                    model_name = model_info['model_name']
                    break
        else:
            model_name = "meta-llama/Llama-3.2-1B"
        
        baseline_tokenizer = AutoTokenizer.from_pretrained(model_name)
        if baseline_tokenizer.pad_token is None:
            baseline_tokenizer.pad_token = baseline_tokenizer.eos_token
            baseline_tokenizer.pad_token_id = baseline_tokenizer.eos_token_id
        
        baseline_model = AutoModelForSequenceClassification.from_pretrained(
            model_name,
            num_labels=2,
            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
            device_map="auto" if device == "cuda" else None
        )
        baseline_model.config.pad_token_id = baseline_tokenizer.pad_token_id
        
        all_results['baseline'] = evaluate_model(baseline_model, baseline_tokenizer, model_name, "純 Llama")
        del baseline_model, baseline_tokenizer
        torch.cuda.empty_cache()
    else:
        all_results['baseline'] = None
    
    # 2. 評估第一次微調模型
    if first_choice != "請選擇":
        print("\n" + "-" * 80)
        print("2️⃣ 評估第一次微調模型")
        print("-" * 80)
        
        # 讀取模型資訊
        with open('./saved_llama_models_list.json', 'r') as f:
            models_list = json.load(f)
        
        first_model_info = None
        for model_info in models_list:
            if model_info['model_path'] == first_choice:
                first_model_info = model_info
                break
        
        if first_model_info:
            tuning_method = first_model_info['tuning_method']
            model_name = first_model_info['model_name']
            
            first_tokenizer = AutoTokenizer.from_pretrained(first_choice)
            if first_tokenizer.pad_token is None:
                first_tokenizer.pad_token = first_tokenizer.eos_token
                first_tokenizer.pad_token_id = first_tokenizer.eos_token_id
            
            if tuning_method in ["LoRA", "AdaLoRA", "Adapter", "Prompt Tuning"]:
                base_model = AutoModelForSequenceClassification.from_pretrained(
                    model_name,
                    num_labels=2,
                    torch_dtype=torch.float16 if device == "cuda" else torch.float32
                )
                first_model = PeftModel.from_pretrained(base_model, first_choice)
                if device == "cuda":
                    first_model = first_model.to(device)
            else:
                first_model = AutoModelForSequenceClassification.from_pretrained(
                    first_choice,
                    num_labels=2,
                    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
                    device_map="auto" if device == "cuda" else None
                )
            
            all_results['first'] = evaluate_model(first_model, first_tokenizer, model_name, "第一次微調模型")
            del first_model, first_tokenizer
            torch.cuda.empty_cache()
        else:
            all_results['first'] = None
    else:
        all_results['first'] = None
    
    # 3. 評估第二次微調模型
    if second_choice != "請選擇":
        print("\n" + "-" * 80)
        print("3️⃣ 評估第二次微調模型")
        print("-" * 80)
        
        # 讀取模型資訊
        with open('./saved_llama_models_list.json', 'r') as f:
            models_list = json.load(f)
        
        second_model_info = None
        for model_info in models_list:
            if model_info['model_path'] == second_choice:
                second_model_info = model_info
                break
        
        if second_model_info:
            tuning_method = second_model_info['tuning_method']
            model_name = second_model_info['model_name']
            
            second_tokenizer = AutoTokenizer.from_pretrained(second_choice)
            if second_tokenizer.pad_token is None:
                second_tokenizer.pad_token = second_tokenizer.eos_token
                second_tokenizer.pad_token_id = second_tokenizer.eos_token_id
            
            if tuning_method in ["LoRA", "AdaLoRA", "Adapter", "Prompt Tuning"]:
                base_model = AutoModelForSequenceClassification.from_pretrained(
                    model_name,
                    num_labels=2,
                    torch_dtype=torch.float16 if device == "cuda" else torch.float32
                )
                second_model = PeftModel.from_pretrained(base_model, second_choice)
                if device == "cuda":
                    second_model = second_model.to(device)
            else:
                second_model = AutoModelForSequenceClassification.from_pretrained(
                    second_choice,
                    num_labels=2,
                    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
                    device_map="auto" if device == "cuda" else None
                )
            
            all_results['second'] = evaluate_model(second_model, second_tokenizer, model_name, "第二次微調模型")
            del second_model, second_tokenizer
            torch.cuda.empty_cache()
        else:
            all_results['second'] = None
    else:
        all_results['second'] = None
    
    print("\n" + "=" * 80)
    print("✅ 新數據測試完成")
    print("=" * 80)
    
    return all_results

def test_new_data_wrapper(test_file, baseline_choice, first_choice, second_choice):
    """新數據測試的包裝函數"""
    
    if test_file is None:
        return "請上傳測試數據 CSV 檔案", "", ""
    
    try:
        all_results = test_on_new_data(
            test_file.name,
            baseline_choice,
            first_choice,
            second_choice
        )
        
        # 格式化輸出
        outputs = []
        
        # 1. 純 Llama
        if all_results['baseline']:
            r = all_results['baseline']
            baseline_output = f"""
# 🔵 純 Llama (Baseline)

| 指標 | 數值 |
|------|------|
| **F1 Score** | {r['f1']:.4f} |
| **Accuracy** | {r['accuracy']:.4f} |
| **Precision** | {r['precision']:.4f} |
| **Recall** | {r['recall']:.4f} |
| **Sensitivity** | {r['sensitivity']:.4f} |
| **Specificity** | {r['specificity']:.4f} |

### 混淆矩陣
|  | 預測:Class 0 | 預測:Class 1 |
|---|-----------|-----------|
| **實際:Class 0** | TN={r['tn']} | FP={r['fp']} |
| **實際:Class 1** | FN={r['fn']} | TP={r['tp']} |
            """
        else:
            baseline_output = "未選擇評估純 Llama"
        outputs.append(baseline_output)
        
        # 2. 第一次微調
        if all_results['first']:
            r = all_results['first']
            first_output = f"""
# 🟢 第一次微調模型

| 指標 | 數值 |
|------|------|
| **F1 Score** | {r['f1']:.4f} |
| **Accuracy** | {r['accuracy']:.4f} |
| **Precision** | {r['precision']:.4f} |
| **Recall** | {r['recall']:.4f} |
| **Sensitivity** | {r['sensitivity']:.4f} |
| **Specificity** | {r['specificity']:.4f} |

### 混淆矩陣
|  | 預測:Class 0 | 預測:Class 1 |
|---|-----------|-----------|
| **實際:Class 0** | TN={r['tn']} | FP={r['fp']} |
| **實際:Class 1** | FN={r['fn']} | TP={r['tp']} |
            """
        else:
            first_output = "未選擇第一次微調模型"
        outputs.append(first_output)
        
        # 3. 第二次微調
        if all_results['second']:
            r = all_results['second']
            second_output = f"""
# 🟡 第二次微調模型

| 指標 | 數值 |
|------|------|
| **F1 Score** | {r['f1']:.4f} |
| **Accuracy** | {r['accuracy']:.4f} |
| **Precision** | {r['precision']:.4f} |
| **Recall** | {r['recall']:.4f} |
| **Sensitivity** | {r['sensitivity']:.4f} |
| **Specificity** | {r['specificity']:.4f} |

### 混淆矩陣
|  | 預測:Class 0 | 預測:Class 1 |
|---|-----------|-----------|
| **實際:Class 0** | TN={r['tn']} | FP={r['fp']} |
| **實際:Class 1** | FN={r['fn']} | TP={r['tp']} |
            """
        else:
            second_output = "未選擇第二次微調模型"
        outputs.append(second_output)
        
        return outputs[0], outputs[1], outputs[2]
        
    except Exception as e:
        import traceback
        error_msg = f"❌ 錯誤:{str(e)}\n\n詳細錯誤訊息:\n{traceback.format_exc()}"
        return error_msg, "", ""

# ==================== 預測函數 ====================
def predict_text(model_choice, text_input):
    """
    預測功能 - 支持選擇已訓練的模型,並同時顯示未微調和微調的預測結果
    """
    
    if not text_input or text_input.strip() == "":
        return "請輸入文本", "請輸入文本"
    
    try:
        # ==================== 未微調的 Llama 預測 ====================
        print("\n使用未微調 Llama 預測...")
        
        # 載入 tokenizer
        if model_choice != "請先訓練模型":
            # 從選擇中解析模型名稱
            model_path = model_choice.split(" | ")[0].replace("路徑: ", "")
            
            # 從 JSON 讀取模型資訊
            with open('./saved_llama_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:
                return "找不到模型資訊", "找不到模型資訊"
            
            model_name = selected_model_info['model_name']
            baseline_tokenizer = AutoTokenizer.from_pretrained(model_name)
        else:
            baseline_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
            model_name = "meta-llama/Llama-3.2-1B"
        
        if baseline_tokenizer.pad_token is None:
            baseline_tokenizer.pad_token = baseline_tokenizer.eos_token
            baseline_tokenizer.pad_token_id = baseline_tokenizer.eos_token_id
        
        baseline_model = AutoModelForSequenceClassification.from_pretrained(
            model_name,
            num_labels=2,
            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
            device_map="auto" if device == "cuda" else None
        )
        baseline_model.config.pad_token_id = baseline_tokenizer.pad_token_id
        baseline_model.eval()
        
        # Tokenize 輸入(未微調)
        baseline_inputs = baseline_tokenizer(
            text_input,
            return_tensors="pt",
            truncation=True,
            max_length=MAX_LENGTH
        )
        if device == "cuda":
            baseline_inputs = {k: v.to(baseline_model.device) for k, v in baseline_inputs.items()}
        
        # 預測(未微調)
        with torch.no_grad():
            baseline_outputs = baseline_model(**baseline_inputs)
            baseline_probs = torch.nn.functional.softmax(baseline_outputs.logits, dim=-1)
            baseline_pred_class = torch.argmax(baseline_probs, dim=-1).item()
            baseline_confidence = baseline_probs[0][baseline_pred_class].item()
        
        baseline_result = "NBCD = 0" if baseline_pred_class == 0 else "NBCD = 1"
        baseline_prob_class0 = baseline_probs[0][0].item()
        baseline_prob_class1 = baseline_probs[0][1].item()
        
        baseline_output = f"""
# 🔵 未微調 Llama 預測結果

## 預測類別: **{baseline_result}**

## 信心度: **{baseline_confidence:.1%}**

## 機率分布:
- **Class 0 機率**: {baseline_prob_class0:.2%}
- **Class 1 機率**: {baseline_prob_class1:.2%}

---
**說明**: 此為原始 Llama 模型,未經任何領域資料訓練
        """
        
        # 清空記憶體
        del baseline_model
        del baseline_tokenizer
        torch.cuda.empty_cache()
        
        # ==================== 微調後的 Llama 預測 ====================
        
        if model_choice == "請先訓練模型":
            finetuned_output = """
# 🟢 微調 Llama 預測結果

❌ 尚未訓練任何模型,請先在「模型訓練」頁面訓練模型
            """
            return baseline_output, finetuned_output
        
        print(f"\n使用微調模型: {model_path}")
        
        # 載入 tokenizer
        finetuned_tokenizer = AutoTokenizer.from_pretrained(model_path)
        if finetuned_tokenizer.pad_token is None:
            finetuned_tokenizer.pad_token = finetuned_tokenizer.eos_token
            finetuned_tokenizer.pad_token_id = finetuned_tokenizer.eos_token_id
        
        # 載入 PEFT 模型(根據微調方法)
        base_model = AutoModelForSequenceClassification.from_pretrained(
            model_name,
            num_labels=2,
            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
            device_map="auto" if device == "cuda" else None
        )
        
        # 根據微調方法載入模型
        tuning_method = selected_model_info.get('tuning_method', 'LoRA')
        
        if tuning_method == "BitFit":
            # BitFit 直接載入完整模型
            finetuned_model = AutoModelForSequenceClassification.from_pretrained(
                model_path,
                num_labels=2,
                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
                device_map="auto" if device == "cuda" else None
            )
        else:
            # 其他方法使用 PEFT
            finetuned_model = PeftModel.from_pretrained(base_model, model_path)
            
            # Prefix Tuning 需要禁用緩存
            if tuning_method == "Prefix Tuning":
                finetuned_model.config.use_cache = False
        
        finetuned_model.config.pad_token_id = finetuned_tokenizer.pad_token_id
        finetuned_model.eval()
        
        # Tokenize 輸入(微調)
        finetuned_inputs = finetuned_tokenizer(
            text_input,
            return_tensors="pt",
            truncation=True,
            max_length=MAX_LENGTH
        )
        if device == "cuda":
            finetuned_inputs = {k: v.to(finetuned_model.device) for k, v in finetuned_inputs.items()}
        
        # 預測(微調)
        with torch.no_grad():
            finetuned_outputs = finetuned_model(**finetuned_inputs)
            finetuned_probs = torch.nn.functional.softmax(finetuned_outputs.logits, dim=-1)
            finetuned_pred_class = torch.argmax(finetuned_probs, dim=-1).item()
            finetuned_confidence = finetuned_probs[0][finetuned_pred_class].item()
        
        finetuned_result = "NBCD = 0" if finetuned_pred_class == 0 else "NBCD = 1"
        finetuned_prob_class0 = finetuned_probs[0][0].item()
        finetuned_prob_class1 = finetuned_probs[0][1].item()
        
        training_type_label = "二次微調" if selected_model_info.get('is_second_finetuning', False) else "第一次微調"
        
        finetuned_output = f"""
# 🟢 微調 Llama 預測結果

## 預測類別: **{finetuned_result}**

## 信心度: **{finetuned_confidence:.1%}**

## 機率分布:
- **Class 0 機率**: {finetuned_prob_class0:.2%}
- **Class 1 機率**: {finetuned_prob_class1:.2%}

---
### 模型資訊:
- **訓練類型**: {training_type_label}
- **模型名稱**: {selected_model_info['model_name']}
- **微調方法**: {selected_model_info['tuning_method']}
- **最佳化指標**: {selected_model_info['best_metric']}
- **訓練時間**: {selected_model_info['timestamp']}
- **模型路徑**: {model_path}

---
**注意**: 此預測僅供參考。
        """
        
        # 清空記憶體
        del finetuned_model
        del finetuned_tokenizer
        torch.cuda.empty_cache()
        
        return baseline_output, finetuned_output
        
    except Exception as 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_llama_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):
        training_type = model_info.get('training_type', '第一次微調')
        choice = f"路徑: {model_info['model_path']} | 類型: {training_type} | 方法: {model_info['tuning_method']} | 時間: {model_info['timestamp']}"
        model_choices.append(choice)
    
    return model_choices

def get_first_finetuning_models():
    """
    取得所有第一次微調的模型(用於二次微調選擇)
    """
    models_list_file = './saved_llama_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)
    
    # 只返回第一次微調的模型
    first_models = [m for m in models_list if not m.get('is_second_finetuning', False)]
    
    if len(first_models) == 0:
        return ["請先進行第一次微調"]
    
    model_choices = []
    for model_info in first_models:
        choice = f"{model_info['model_path']}"
        model_choices.append(choice)
    
    return model_choices

# ==================== Gradio 介面 (參考第四個文件的視覺化) ====================
with gr.Blocks(title="🦙 Llama NBCD 二次微調平台", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    # 🦙 Llama NBCD 二次微調完整平台
    
    ### 🌟 功能特色:
    - 🎯 第一次微調:從純 Llama 開始訓練
    - 🔄 第二次微調:基於第一次模型用新數據繼續訓練
    - 📊 自動比較有/無微調的表現差異
    - 🎨 可選擇最佳化指標(F1、Accuracy、Precision、Recall)
    - 🔮 訓練後可直接預測新樣本
    - 💾 自動儲存最佳模型
    - 🧹 自動記憶體管理
    
    ✅ **支持的微調方法**: LoRA, AdaLoRA, Adapter, BitFit, Prompt Tuning
    ⚠️ **暫不支持**: Prefix Tuning (版本兼容性問題,請使用 Prompt Tuning 替代)
    """)
    
    # Tab 1: 第一次微調
    with gr.Tab("1️⃣ 第一次微調"):
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📤 資料上傳")
                
                file_input = gr.File(
                    label="上傳 CSV 檔案",
                    file_types=[".csv"]
                )
                
                gr.Markdown("### 🤖 模型選擇")
                
                model_name_input = gr.Textbox(
                    value="meta-llama/Llama-3.2-1B",
                    label="Hugging Face 模型名稱",
                    info="例如: meta-llama/Llama-3.2-1B"
                )
                
                gr.Markdown("### 🔧 微調方法選擇")
                
                tuning_method = gr.Radio(
                    choices=["LoRA", "AdaLoRA", "Adapter", "BitFit", "Prompt Tuning"],
                    value="LoRA",
                    label="選擇微調方法",
                    info="不同的參數效率微調方法 (Prefix Tuning 暫不支持)"
                )
                
                gr.Markdown("### 🎯 最佳模型選擇")
                
                best_metric = gr.Dropdown(
                    choices=["f1", "accuracy", "precision", "recall", "sensitivity", "specificity"],
                    value="recall",
                    label="選擇最佳化指標",
                    info="模型會根據此指標選擇最佳檢查點"
                )
                
                gr.Markdown("### ⚙️ 資料平衡參數")
                
                target_samples_input = gr.Number(
                    value=700,
                    label="目標樣本數(每類別)"
                )
                
                use_weights_checkbox = gr.Checkbox(
                    value=True,
                    label="使用類別權重",
                    info="在損失函數中使用類別權重"
                )
                
                gr.Markdown("### ⚙️ 訓練參數")
                
                epochs_input = gr.Number(
                    value=3,
                    label="訓練輪數 (Epochs)"
                )
                
                batch_size_input = gr.Number(
                    value=4,
                    label="批次大小 (Batch Size)"
                )
                
                lr_input = gr.Number(
                    value=1e-4,
                    label="學習率 (Learning Rate)"
                )
                
                gr.Markdown("---")
                
                # LoRA 參數
                with gr.Column(visible=True) as lora_params:
                    gr.Markdown("### 🔷 LoRA 參數")
                    
                    lora_r_input = gr.Slider(
                        minimum=4,
                        maximum=64,
                        value=16,
                        step=4,
                        label="LoRA Rank (r)",
                        info="低秩分解的秩"
                    )
                    
                    lora_alpha_input = gr.Slider(
                        minimum=8,
                        maximum=128,
                        value=32,
                        step=8,
                        label="LoRA Alpha",
                        info="LoRA 縮放參數"
                    )
                    
                    lora_dropout_input = gr.Slider(
                        minimum=0.0,
                        maximum=0.5,
                        value=0.1,
                        step=0.05,
                        label="LoRA Dropout",
                        info="Dropout 率"
                    )
                    
                    lora_target_input = gr.Dropdown(
                        choices=["query,value", "query,key,value", "all"],
                        value="query,value",
                        label="目標模組",
                        info="用逗號分隔"
                    )
                
                # AdaLoRA 參數
                with gr.Column(visible=False) as adalora_params:
                    gr.Markdown("### 🔶 AdaLoRA 參數")
                    
                    adalora_init_r_input = gr.Slider(
                        minimum=4,
                        maximum=64,
                        value=12,
                        step=4,
                        label="初始 Rank",
                        info="訓練開始時的秩"
                    )
                    
                    adalora_target_r_input = gr.Slider(
                        minimum=4,
                        maximum=64,
                        value=8,
                        step=4,
                        label="目標 Rank",
                        info="訓練結束時的目標秩"
                    )
                    
                    adalora_alpha_input = gr.Slider(
                        minimum=8,
                        maximum=128,
                        value=32,
                        step=8,
                        label="LoRA Alpha",
                        info="縮放參數"
                    )
                    
                    adalora_tinit_input = gr.Number(
                        value=0,
                        label="Tinit",
                        info="開始剪枝的步數"
                    )
                    
                    adalora_tfinal_input = gr.Number(
                        value=0,
                        label="Tfinal",
                        info="結束剪枝的步數"
                    )
                    
                    adalora_delta_t_input = gr.Number(
                        value=1,
                        label="Delta T",
                        info="剪枝頻率"
                    )
                
                # Adapter 參數
                with gr.Column(visible=False) as adapter_params:
                    gr.Markdown("### 🔶 Adapter 參數")
                    
                    adapter_reduction_input = gr.Slider(
                        minimum=2,
                        maximum=64,
                        value=16,
                        step=2,
                        label="Reduction Factor",
                        info="降維因子,越大參數越少"
                    )
                
                # Prompt Tuning 參數
                with gr.Column(visible=False) as prompt_tuning_params:
                    gr.Markdown("### 🔷 Prompt Tuning 參數")
                    
                    prompt_tokens_input = gr.Slider(
                        minimum=1,
                        maximum=100,
                        value=20,
                        step=1,
                        label="Virtual Tokens 數量"
                    )
                
                # Prefix Tuning 參數
                with gr.Column(visible=False) as prefix_tuning_params:
                    gr.Markdown("### 🔶 Prefix Tuning 參數")
                    gr.Markdown("⚠️ **注意**: 目前版本可能有兼容性問題,建議使用 Prompt Tuning")
                    
                    prefix_tokens_input = gr.Slider(
                        minimum=1,
                        maximum=100,
                        value=30,
                        step=1,
                        label="Virtual Tokens 數量"
                    )
                
                train_button = gr.Button(
                    "🚀 開始第一次微調",
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column(scale=2):
                gr.Markdown("### 📊 第一次微調結果與比較")
                
                # 第一格:資料資訊
                data_info_output = gr.Markdown(
                    value="### 等待訓練...\n\n訓練完成後會顯示資料資訊和訓練配置",
                    label="資料資訊"
                )
                
                # 第二和第三格:並排顯示
                with gr.Row():
                    # 第二格:未微調 Llama
                    baseline_output = gr.Markdown(
                        value="### 未微調 Llama\n等待訓練完成...",
                        label="未微調 Llama"
                    )
                    
                    # 第三格:微調後 Llama
                    finetuned_output = gr.Markdown(
                        value="### 第一次微調 Llama\n等待訓練完成...",
                        label="第一次微調 Llama"
                    )
    
    # Tab 2: 二次微調
    with gr.Tab("2️⃣ 二次微調"):
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 🔄 選擇基礎模型")
                base_model_dropdown = gr.Dropdown(
                    label="選擇第一次微調的模型",
                    choices=["請先進行第一次微調"],
                    value="請先進行第一次微調"
                )
                refresh_base_models = gr.Button("🔄 重新整理模型列表", size="sm")
                
                gr.Markdown("### 📤 上傳新訓練數據")
                file_input_second = gr.File(label="上傳新的訓練數據 CSV", file_types=[".csv"])
                
                gr.Markdown("### ⚙️ 訓練參數")
                gr.Markdown("⚠️ 微調方法將自動繼承第一次微調的方法")
                best_metric_second = gr.Dropdown(
                    choices=["f1", "accuracy", "precision", "recall", "sensitivity", "specificity"],
                    value="f1",
                    label="選擇最佳化指標"
                )
                
                target_samples_second = gr.Number(
                    value=700,
                    label="目標樣本數(每類別)"
                )
                
                use_weights_second = gr.Checkbox(
                    value=True,
                    label="使用類別權重"
                )
                
                epochs_input_second = gr.Number(value=3, label="訓練輪數", info="建議比第一次少")
                batch_size_input_second = gr.Number(value=4, label="批次大小")
                lr_input_second = gr.Number(value=5e-5, label="學習率", info="建議比第一次小")
                
                train_button_second = gr.Button("🚀 開始二次微調", variant="primary", size="lg")
            
            with gr.Column(scale=2):
                gr.Markdown("### 📊 二次微調結果")
                data_info_output_second = gr.Markdown(value="等待訓練...")
                finetuned_output_second = gr.Markdown(value="### 二次微調\n等待訓練...")
    
    # Tab 3: 新數據測試
    with gr.Tab("3️⃣ 新數據測試"):
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📤 上傳測試數據")
                test_file_input = gr.File(label="上傳測試數據 CSV", file_types=[".csv"])
                
                gr.Markdown("### 🎯 選擇要比較的模型")
                gr.Markdown("可選擇 1-3 個模型進行比較")
                
                baseline_test_choice = gr.Radio(
                    choices=["評估純 Llama", "跳過"],
                    value="評估純 Llama",
                    label="純 Llama (Baseline)"
                )
                
                first_model_test_dropdown = gr.Dropdown(
                    label="第一次微調模型",
                    choices=["請選擇"],
                    value="請選擇"
                )
                
                second_model_test_dropdown = gr.Dropdown(
                    label="第二次微調模型",
                    choices=["請選擇"],
                    value="請選擇"
                )
                
                refresh_test_models = gr.Button("🔄 重新整理模型列表", size="sm")
                test_button = gr.Button("📊 開始測試", variant="primary", size="lg")
            
            with gr.Column(scale=2):
                gr.Markdown("### 📊 新數據測試結果 - 三模型比較")
                with gr.Row():
                    baseline_test_output = gr.Markdown(value="### 純 Llama\n等待測試...")
                    first_test_output = gr.Markdown(value="### 第一次微調\n等待測試...")
                    second_test_output = gr.Markdown(value="### 二次微調\n等待測試...")
    
    # Tab 4: 模型預測
    with gr.Tab("4️⃣ 模型預測"):
        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="請輸入要預測的文本...",
                    lines=10
                )
                
                predict_button = gr.Button(
                    "🔮 開始預測",
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column():
                gr.Markdown("### 預測結果比較")
                
                # 上框:未微調 Llama 預測結果
                baseline_prediction_output = gr.Markdown(
                    label="未微調 Llama",
                    value="等待預測..."
                )
                
                # 下框:微調 Llama 預測結果
                finetuned_prediction_output = gr.Markdown(
                    label="微調 Llama",
                    value="等待預測..."
                )
    
    # Tab 5: 使用說明
    with gr.Tab("📖 使用說明"):
        gr.Markdown("""
        ## 🔄 二次微調流程說明
        
        ### 步驟 1: 第一次微調
        1. 上傳訓練數據 A (CSV 格式: Text, label)
        2. 選擇微調方法 (LoRA / AdaLoRA / Adapter / BitFit / Prompt Tuning)
        3. 調整訓練參數
        4. 開始訓練
        5. 系統會自動比較純 Llama vs 第一次微調的表現
        
        ### 步驟 2: 二次微調
        1. 選擇已訓練的第一次微調模型
        2. 上傳新的訓練數據 B
        3. 調整訓練參數 (建議 epochs 更小, learning rate 更小)
        4. 開始訓練 (方法自動繼承第一次)
        5. 模型會基於第一次的權重繼續學習
        
        ### 步驟 3: 預測
        1. 選擇任一已訓練模型
        2. 輸入文本
        3. 查看預測結果
        
        ## 🎯 微調方法說明
        
        | 方法 | 參數量 | 記憶體 | 訓練速度 | 適用場景 |
        |------|--------|--------|----------|----------|
        | **LoRA** | 很少 (~1%) | 低 | 快 | 通用,效果好 |
        | **AdaLoRA** | 很少 (~1%) | 低 | 快 | 自適應,效果更優 |
        | **Adapter** | 少 (~2-5%) | 低 | 中 | 多任務學習 |
        | **BitFit** | 極少 (~0.1%) | 極低 | 極快 | 快速微調 |
        | **Prompt Tuning** | 極少 (可調) | 極低 | 快 | 小數據集 |
        
        ## 💡 二次微調建議
        
        ### 訓練參數調整:
        - **Epochs**: 第二次建議 3-5 輪 (第一次通常 8-10 輪)
        - **Learning Rate**: 第二次建議 5e-5 (第一次通常 1e-4)
        - **Warmup Steps**: 第二次建議減半
        
        ### 適用場景:
        1. **領域適應**: 第一次用通用醫療數據,第二次用特定醫院數據
        2. **增量學習**: 隨時間增加新病例數據
        3. **數據稀缺**: 先用大量相關數據預訓練,再用少量目標數據微調
        
        ## ⚠️ 注意事項
        
        - CSV 格式必須包含 `Text` 和 `label` 欄位
        - 第二次微調會自動使用第一次的微調方法
        - 建議第二次的學習率比第一次小,避免破壞已學習的知識
        - 訓練時間依資料量和硬體而定(10-30 分鐘)
        - 需要 Hugging Face Token 才能下載 Llama 模型
        - GPU 訓練效果最佳,CPU 會非常慢
        
        ## 📊 指標說明
        
        - **F1 Score**: 精確率和召回率的調和平均,平衡指標
        - **Accuracy**: 整體準確率
        - **Precision**: 預測為正類中的準確率
        - **Recall/Sensitivity**: 實際正類中被正確識別的比例
        - **Specificity**: 實際負類中被正確識別的比例
        
        ## 🔧 已修復的問題
        
        - ✅ **AdaLoRA**: 簡化配置參數,避免版本兼容性問題
        - ✅ **BitFit**: 正確處理 gradient 設置,包含分類頭訓練
        - ✅ **參數顯示**: AdaLoRA 現在會正確顯示專屬參數界面
        - ❌ **Prefix Tuning**: 因 PEFT 版本問題暫時移除,請用 Prompt Tuning 替代
        
        ## 🔐 設定 HF Token
        
        在環境變數中設定:
        ```
        export HF_TOKEN=your_token_here
        ```
        """)
    
    # ==================== 事件綁定 ====================
    
    # 根據選擇的微調方法顯示/隱藏相應參數
    def update_params_visibility(method):
        if method == "LoRA":
            return (
                gr.update(visible=True),   # lora_params
                gr.update(visible=False),  # adalora_params
                gr.update(visible=False),  # adapter_params
                gr.update(visible=False),  # prompt_tuning_params
                gr.update(visible=False)   # prefix_tuning_params
            )
        elif method == "AdaLoRA":
            return (
                gr.update(visible=False),  # lora_params
                gr.update(visible=True),   # adalora_params
                gr.update(visible=False),  # adapter_params
                gr.update(visible=False),  # prompt_tuning_params
                gr.update(visible=False)   # prefix_tuning_params
            )
        elif method == "Adapter":
            return (
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(visible=True),
                gr.update(visible=False),
                gr.update(visible=False)
            )
        elif method == "Prompt Tuning":
            return (
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(visible=True),
                gr.update(visible=False)
            )
        elif method == "Prefix Tuning":
            return (
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(visible=True)
            )
        else:  # BitFit
            return (
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(visible=False)
            )
    
    tuning_method.change(
        fn=update_params_visibility,
        inputs=[tuning_method],
        outputs=[lora_params, adalora_params, adapter_params, prompt_tuning_params, prefix_tuning_params]
    )
    
    # 設定第一次微調按鈕動作
    train_button.click(
        fn=train_first_wrapper,
        inputs=[
            file_input,
            model_name_input,
            target_samples_input,
            use_weights_checkbox,
            epochs_input,
            batch_size_input,
            lr_input,
            tuning_method,
            lora_r_input,
            lora_alpha_input,
            lora_dropout_input,
            lora_target_input,
            adalora_init_r_input,
            adalora_target_r_input,
            adalora_alpha_input,
            adalora_tinit_input,
            adalora_tfinal_input,
            adalora_delta_t_input,
            adapter_reduction_input,
            prompt_tokens_input,
            prefix_tokens_input,
            best_metric
        ],
        outputs=[data_info_output, baseline_output, finetuned_output]
    )
    
    # 重新整理基礎模型列表按鈕
    def refresh_base_models_list():
        choices = get_first_finetuning_models()
        return gr.update(choices=choices, value=choices[0])
    
    refresh_base_models.click(
        fn=refresh_base_models_list,
        outputs=[base_model_dropdown]
    )
    
    # 二次微調按鈕
    train_button_second.click(
        fn=train_second_wrapper,
        inputs=[
            base_model_dropdown,
            file_input_second,
            target_samples_second,
            use_weights_second,
            epochs_input_second,
            batch_size_input_second,
            lr_input_second,
            best_metric_second
        ],
        outputs=[data_info_output_second, finetuned_output_second]
    )
    
    # 重新整理測試模型列表
    def refresh_test_models_list():
        all_models = get_available_models()
        first_models = get_first_finetuning_models()
        
        # 篩選第二次微調模型
        with open('./saved_llama_models_list.json', 'r') as f:
            models_list = json.load(f)
        second_models = [m['model_path'] for m in models_list if m.get('is_second_finetuning', False)]
        
        if len(second_models) == 0:
            second_models = ["請選擇"]
        
        return (
            gr.update(choices=first_models if first_models[0] != "請先進行第一次微調" else ["請選擇"], value="請選擇"),
            gr.update(choices=second_models, value="請選擇")
        )
    
    refresh_test_models.click(
        fn=refresh_test_models_list,
        outputs=[first_model_test_dropdown, second_model_test_dropdown]
    )
    
    # 測試按鈕
    test_button.click(
        fn=test_new_data_wrapper,
        inputs=[test_file_input, baseline_test_choice, first_model_test_dropdown, second_model_test_dropdown],
        outputs=[baseline_test_output, first_test_output, second_test_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=[model_dropdown, text_input],
        outputs=[baseline_prediction_output, finetuned_prediction_output]
    )

if __name__ == "__main__":
    demo.launch()