test2.py

import os
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from peft import LoraConfig, PeftModel
from trl import SFTTrainer

# --- 1. 設定基礎模型 ---
# 在此處選擇您要微調的基礎模型 ID
# model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
model_name = "AI-Model-Scope/Llama3-TAIDE-LX-8B-Chat-v1" 

# --- 2. 設定資料集路徑 ---
# 請確認這個路徑指向您上一階段產出的 .jsonl 檔案
dataset_name = "judge_training_data/judge_finetuning_dataset_v2.jsonl"

# --- 3. 設定新模型的名稱 ---
# 微調完成後，您的新模型將會儲存在這個路徑
new_model_name = "MyJudgeModel-Llama3-TAIDE-v1"

# =================================================================================
# 以下是 QLoRA 微調的核心設定，通常不需要修改
# =================================================================================

# --- 4. QLoRA 參數設定 ---
# 這些設定能讓模型以 4-bit 的精度載入，大幅降低 VRAM 需求
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=False,
)

# --- 5. LoRA 參數設定 ---
# 這是 PEFT (參數高效微調) 的設定
peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.1,
    r=64,
    bias="none",
    task_type="CAUSAL_LM",
)

# --- 6. 載入基礎模型與 Tokenizer ---
print(f"正在從 {model_name} 載入模型...")
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto" # 自動將模型分配到可用的 GPU
)
model.config.use_cache = False # 訓練時建議關閉

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
# Llama 3 需要手動設定 Pad Token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

# --- 7. 載入並準備資料集 ---
print(f"正在從 {dataset_name} 載入資料集...")
# Hugging Face datasets 函式庫可以直接讀取 .jsonl
dataset = load_dataset("json", data_files=dataset_name, split="train")

# --- 8. 設定訓練參數 (Training Arguments) ---
training_arguments = TrainingArguments(
    output_dir=f"./results/{new_model_name}", # 訓練過程中的檢查點會儲存在這裡
    num_train_epochs=1,                     # 訓練輪數 (Epochs)，建議從 1 開始
    per_device_train_batch_size=4,          # 每個 GPU 的批次大小
    gradient_accumulation_steps=1,          # 梯度累積步數
    optim="paged_adamw_32bit",              # 優化器
    save_steps=50,                          # 每隔多少步儲存一次檢查點
    logging_steps=10,                       # 每隔多少步輸出一筆 log
    learning_rate=2e-4,                     # 學習率
    weight_decay=0.001,
    fp16=False,
    bf16=True,                             # 如果您的 GPU 支援 (Ampere 架構或更新)，請設為 True
    max_grad_norm=0.3,
    max_steps=-1,                           # 如果設為正數，會覆蓋 num_train_epochs
    warmup_ratio=0.03,
    group_by_length=True,
    lr_scheduler_type="constant",
)

# --- 9. 建立 SFTTrainer 並開始訓練！ ---
# SFTTrainer 是 trl 函式庫中一個高度封裝的訓練器，非常方便
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_config,
    dataset_text_field="instruction",  # **重要**：這裡我們不直接用欄位，而是用一個格式化函數
    max_seq_length=2048,               # 最大序列長度，可根據您的 VRAM 調整
    tokenizer=tokenizer,
    args=training_arguments,
    packing=False,
    # **重要**：定義一個函數來將我們的資料格式化成 "指令 -> 回應" 的形式
    formatting_func=lambda example: f"### 指令:\n{example['instruction']}\n\n### 回應:\n{example['output']}"
)

print("訓練即將開始...")
trainer.train()

# --- 10. 儲存最終的模型 ---
print("訓練完成，正在儲存最終模型...")
trainer.model.save_pretrained(new_model_name)
tokenizer.save_pretrained(new_model_name)

print(f"恭喜！您的裁判模型已成功訓練並儲存於 '{new_model_name}'")