"""
train.py - Qwen2.5-7B QLoRA SFT 训练脚本 (终极修复版)
适配: RTX 5070 Ti 16GB + Windows + PyTorch 2.11+cu128 + 最新版 trl
"""
import os
from pathlib import Path

# ⚠️ 核心修复方案：不依赖任何环境变量
# 直接拦截并替换 pathlib.Path 的 read_text 方法，强制所有未指定编码的文件读取使用 UTF-8
_original_read_text = Path.read_text
def _forced_utf8_read_text(self, *args, **kwargs):
    if 'encoding' not in kwargs:
        kwargs['encoding'] = 'utf-8'
    return _original_read_text(self, *args, **kwargs)

Path.read_text = _forced_utf8_read_text

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer, SFTConfig

# ==================== 配置区 ====================
MODEL_NAME = r"./models/qwen/Qwen2.5-7B-Instruct"
DATA_PATH = "dataset/train_clean.jsonl"
OUTPUT_DIR = "./qwen2.5-7b-sft-output"

# 【修改】将 max_seq_length 改为新版要求的 max_length
SFT_CONFIG = SFTConfig(
    output_dir=OUTPUT_DIR,
    num_train_epochs=3,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,
    learning_rate=2e-4,
    bf16=True,
    gradient_checkpointing=True,
    dataloader_num_workers=0,            # Windows 兼容设置，防止多进程数据加载报错
    logging_steps=10,
    save_strategy="epoch",
    optim="paged_adamw_8bit",            # 显存优化，防止 OOM
    report_to="none",

    # 【核心修复】新版 trl (>=v0.20.0) 已将 max_seq_length 更名为 max_length
    max_length=2048,                     # 最大序列长度
    packing=False,                       # 关闭打包（适用于自定义格式化）
    dataset_text_field="text",           # 占位字段（当使用 formatting_func 时此项仅作兼容保留）
)

LORA_CONFIG = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    task_type="CAUSAL_LM",
    bias="none",
)
# ================================================

def main():
    print("📦 加载模型与分词器...")

    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype="bfloat16",
        bnb_4bit_use_double_quant=True,
    )

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
        dtype="bfloat16"                 # 【优化】顺便将 torch_dtype 更新为新版推荐的 dtype
    )
    model = prepare_model_for_kbit_training(model)
    model = get_peft_model(model, LORA_CONFIG)
    model.print_trainable_parameters()

    print("📂 加载数据集...")
    dataset = load_dataset("json", data_files=DATA_PATH, split="train")

    def formatting_func(example):
        return f"### Instruction:\n{example['instruction']}\n\n### Response:\n{example['output']}{tokenizer.eos_token}"

    print("🚀 开始训练...")
    trainer = SFTTrainer(
        model=model,
        args=SFT_CONFIG,                 # 传入整合后的 SFTConfig
        train_dataset=dataset,
        processing_class=tokenizer,      # 新版分词器参数名
        formatting_func=formatting_func,
    )
    trainer.train()

    final_path = f"{OUTPUT_DIR}/final"
    trainer.save_model(final_path)
    tokenizer.save_pretrained(final_path)
    print(f"✅ 训练完成！LoRA权重已保存至: {final_path}")


if __name__ == "__main__":
    main()