# 核心依赖
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# HuggingFace 生态
pip install transformers datasets accelerate

# QLoRA 必需组件
pip install bitsandbytes  # 4-bit量化工具
pip install peft           # LoRA 实现
pip install trl            # SFTTrainer 所在的包

# 验证安装
python -c "import torch; print('CUDA可用:', torch.cuda.is_available())"
python -c "import bitsandbytes; print('bitsandbytes版本:', bitsandbytes.__version__)"

# train_qlora.py
# 完整的 QLoRA 微调脚本
# 运行方式：python train_qlora.py

import json
import torch
from datasets import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
    TrainingArguments,
)
from peft import (
    LoraConfig,
    TaskType,
    get_peft_model,
    prepare_model_for_kbit_training,
)
from trl import SFTTrainer, SFTConfig

# ============================================================
# 第一步：配置参数（根据你的GPU调整这里）
# ============================================================

MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct"  # 基础模型
DATA_PATH = "training_data_final.json"    # 上一章准备的数据集
OUTPUT_DIR = "./qlora-output"             # 训练结果保存目录

# GPU显存不够？参考下方"不同GPU的推荐配置"调整
BATCH_SIZE = 2                 # 每个GPU每步处理的样本数
GRAD_ACCUMULATION = 8          # 梯度累积步数（等效batch_size = 2×8 = 16）
MAX_SEQ_LENGTH = 1024          # 最大序列长度（越长显存越多）
NUM_EPOCHS = 3                 # 训练轮数
LEARNING_RATE = 2e-4           # 学习率

print("="*60)
print("QLoRA 微调开始")
print(f"模型: {MODEL_NAME}")
print(f"数据: {DATA_PATH}")
print(f"设备: {'CUDA' if torch.cuda.is_available() else 'CPU（警告：训练会非常慢）'}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"显存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
print("="*60)

# ============================================================
# 第二步：加载和格式化数据集
# ============================================================

def format_prompt(example: dict) -> str:
    """
    将数据集中的一条样本转换为模型的输入格式。

    Qwen2.5 使用 ChatML 格式，我们需要把 instruction/input/output
    转换成对话格式，让模型学会按照这种格式回答问题。
    """
    instruction = example.get('instruction', '')
    input_text = example.get('input', '').strip()
    output = example.get('output', '')

    # 如果有 input（待处理的原始材料），拼接到 instruction 后面
    if input_text:
        user_content = f"{instruction}\n\n{input_text}"
    else:
        user_content = instruction

    # ChatML 格式（Qwen2.5 的标准对话格式）
    prompt = f"<|im_start|>system\n你是一个有帮助的AI助手。<|im_end|>\n"
    prompt += f"<|im_start|>user\n{user_content}<|im_end|>\n"
    prompt += f"<|im_start|>assistant\n{output}<|im_end|>"

    return prompt


print("\n[1/5] 加载数据集...")
with open(DATA_PATH, 'r', encoding='utf-8') as f:
    raw_data = json.load(f)

print(f"原始数据量: {len(raw_data)} 条")

# 格式化数据
formatted_data = [{"text": format_prompt(item)} for item in raw_data]

# 转为 HuggingFace Dataset 格式
dataset = Dataset.from_list(formatted_data)

# 划分训练集和验证集（90% 训练，10% 验证）
split = dataset.train_test_split(test_size=0.1, seed=42)
train_dataset = split['train']
eval_dataset = split['test']

print(f"训练集: {len(train_dataset)} 条")
print(f"验证集: {len(eval_dataset)} 条")

# ============================================================
# 第三步：加载量化模型
# ============================================================

print("\n[2/5] 加载量化模型（这需要几分钟，请耐心等待）...")

# 4-bit 量化配置
# NF4（NormalFloat4）是 QLoRA 论文提出的量化类型，
# 专门为神经网络权重的正态分布设计，比普通 int4 效果更好
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,                        # 开启4-bit量化
    bnb_4bit_quant_type="nf4",               # 使用NF4量化（QLoRA推荐）
    bnb_4bit_compute_dtype=torch.float16,    # 计算时用float16（精度和速度的平衡）
    bnb_4bit_use_double_quant=True,          # 双重量化：对量化系数再量化，节省约0.4GB
)

# 加载模型（量化模型会自动分配到GPU）
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    device_map="auto",         # 自动分配到可用设备
    trust_remote_code=True,    # Qwen 模型需要这个参数
)

# 加载 tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    padding_side="right",      # 右填充（训练时推荐）
)

# 确保有 padding token（有些模型没有，需要手动设置）
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

print(f"模型加载完成，参数量: {model.num_parameters() / 1e9:.1f}B")

# ============================================================
# 第四步：配置 LoRA
# ============================================================

print("\n[3/5] 配置 LoRA...")

# 为 4-bit 量化训练做准备（必须在添加LoRA之前调用）
# 这一步会：
# 1. 把所有非LoRA参数设为不可训练
# 2. 把 LayerNorm 转为 float32（提升训练稳定性）
model = prepare_model_for_kbit_training(model)

# LoRA 配置
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,    # 任务类型：自回归语言模型
    r=16,                             # LoRA 秩（rank）：控制新增参数量
                                      # r=8 更省显存，r=32 效果可能更好
    lora_alpha=32,                    # 缩放系数（通常设为 r 的2倍）
    lora_dropout=0.05,               # Dropout 防止过拟合
    target_modules=[                  # 给哪些层添加 LoRA
        "q_proj", "k_proj",           # 注意力的 Query 和 Key
        "v_proj", "o_proj",           # 注意力的 Value 和 Output
        "gate_proj", "up_proj", "down_proj"  # FFN 层
    ],
    bias="none",                      # 不训练 bias（节省参数）
)

# 给模型添加 LoRA 层
model = get_peft_model(model, lora_config)

# 打印可训练参数量（LoRA 只训练极少数参数）
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
all_params = sum(p.numel() for p in model.parameters())
print(f"可训练参数: {trainable_params:,} ({trainable_params/all_params*100:.2f}%)")
print(f"全部参数: {all_params:,}")
# 正常情况下，可训练参数占比约为 0.1%-1%

# ============================================================
# 第五步：配置训练参数并启动训练
# ============================================================

print("\n[4/5] 配置训练参数...")

# SFTConfig 继承自 TrainingArguments，专为指令微调设计
training_args = SFTConfig(
    output_dir=OUTPUT_DIR,                    # 保存检查点的目录
    num_train_epochs=NUM_EPOCHS,              # 训练轮数
    per_device_train_batch_size=BATCH_SIZE,   # 每个GPU每步处理的样本数
    per_device_eval_batch_size=BATCH_SIZE,    # 评估时的batch大小
    gradient_accumulation_steps=GRAD_ACCUMULATION,  # 梯度累积（等效增大batch size）
    learning_rate=LEARNING_RATE,              # 学习率
    lr_scheduler_type="cosine",              # 余弦学习率衰减（比线性更平滑）
    warmup_ratio=0.05,                        # 预热比例（训练初期缓慢提升学习率）

    # 显存优化
    fp16=True,                                # 使用 float16 计算（节省显存）
    gradient_checkpointing=True,             # 梯度检查点（训练时不保存所有中间计算结果，需要时重新计算，用时间换显存）
    optim="paged_adamw_32bit",               # QLoRA 推荐的分页优化器（一种节省显存的 Adam 变体，把优化器状态分页存储到CPU内存）

    # 保存和评估策略
    evaluation_strategy="steps",            # 每N步评估一次（而不是每epoch）
    eval_steps=100,                          # 每100步评估一次
    save_strategy="steps",
    save_steps=100,                          # 每100步保存一次检查点
    save_total_limit=3,                      # 最多保留3个检查点（节省磁盘空间）
    load_best_model_at_end=True,             # 训练结束后加载最佳模型

    # 日志
    logging_steps=10,                        # 每10步打印一次训练损失
    report_to="none",                        # 不上报到 wandb/tensorboard（避免需要登录）

    # SFT 特有参数
    max_seq_length=MAX_SEQ_LENGTH,           # 最大序列长度
    dataset_text_field="text",               # 数据集中存储文本的字段名
    packing=False,                           # 不打包（打包可提高效率但可能影响效果）
)

print("\n[5/5] 启动训练...")

# 初始化 SFTTrainer
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
)

# 开始训练！
# 训练日志会实时显示损失值（loss），正常应从约2.0逐渐下降到0.5以下
train_result = trainer.train()

# 打印训练结果摘要
print("\n" + "="*60)
print("训练完成！")
print(f"训练步数: {train_result.global_step}")
print(f"训练损失: {train_result.training_loss:.4f}")
print("="*60)

# 保存最终模型（只保存 LoRA 权重，很小，通常只有几十MB）
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)
print(f"模型已保存到: {OUTPUT_DIR}")

# 原理：
# 正常训练：保存所有中间激活值（占用大量显存）
# 梯度检查点：不保存中间激活值，反向传播时重新计算
# 代价：训练速度下降约 20-30%，但显存节省约 50%（实际效果因模型和配置而异，供参考）

# 在 TrainingArguments 中已经配置了：
# gradient_checkpointing=True

# 效果对比（7B模型，batch=2）：
# 不开启：约20GB显存
# 开启后：约12GB显存（节省40%）

{'loss': 2.3421, 'learning_rate': 0.0002, 'epoch': 0.05}
{'loss': 1.8732, 'learning_rate': 0.000198, 'epoch': 0.1}
{'loss': 1.2541, 'learning_rate': 0.000185, 'epoch': 0.3}
...
{'loss': 0.4823, 'learning_rate': 0.0001, 'epoch': 1.5}

# verify_finetune.py
# 对比基础模型和微调后模型的输出

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel

# ---- 加载微调后的模型 ----

MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct"
LORA_PATH = "./qlora-output"   # 训练结果目录

# 4-bit量化配置（推理时也用量化，节省显存）
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
)

print("加载基础模型...")
base_model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
)

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

# 加载 LoRA 权重（叠加在基础模型上）
print("加载 LoRA 权重...")
finetuned_model = PeftModel.from_pretrained(base_model, LORA_PATH)
finetuned_model.eval()  # 切换到推理模式


def generate_response(model, tokenizer, user_input: str, max_new_tokens: int = 300) -> str:
    """
    用给定的模型生成回答。

    Args:
        model: 要测试的模型（基础模型或微调后的模型）
        tokenizer: 对应的 tokenizer
        user_input: 用户输入的问题
        max_new_tokens: 最多生成多少个新 token

    Returns:
        模型生成的文本
    """
    # 构建 ChatML 格式的输入（与训练时保持一致！）
    messages = [
        {"role": "system", "content": "你是一个有帮助的AI助手。"},
        {"role": "user", "content": user_input}
    ]

    # 用 tokenizer 的 chat_template 格式化输入
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True  # 在末尾加上 <|im_start|>assistant\n
    )

    # 转为模型输入
    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    # 生成回答
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.7,        # 控制随机性（0=完全确定，1=较随机）
            do_sample=True,         # 采样模式（而不是贪婪搜索）
            repetition_penalty=1.1, # 轻微惩罚重复（避免模型陷入循环）
            pad_token_id=tokenizer.eos_token_id,
        )

    # 只取新生成的部分（去掉输入的 prompt）
    new_tokens = outputs[0][inputs['input_ids'].shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True)


# ---- 测试问题列表 ----
# 从你的训练数据主题中选几个有代表性的问题
test_questions = [
    "Python 中的装饰器是什么？请用一个简单的例子说明。",
    "如何处理 Python 中的 FileNotFoundError 异常？",
    "列表推导式和普通 for 循环有什么区别？",
]

# ---- 对比基础模型和微调后模型 ----
print("\n" + "="*60)
print("对比测试：基础模型 vs 微调后模型")
print("="*60)

for i, question in enumerate(test_questions, 1):
    print(f"\n[问题 {i}]: {question}")
    print("-" * 40)

    # 基础模型的回答
    base_answer = generate_response(base_model, tokenizer, question)
    print(f"[基础模型]:\n{base_answer}")

    print("-" * 40)

    # 微调后模型的回答
    finetuned_answer = generate_response(finetuned_model, tokenizer, question)
    print(f"[微调后模型]:\n{finetuned_answer}")

    print("="*60)

print("\n提示：如果微调后的回答在风格、格式或专业度上更接近你的训练数据，说明微调成功。")
print("下一步：阅读《微调效果评估与部署》，了解如何系统地量化评估效果。")

# merge_lora.py
# 将 LoRA 权重合并到基础模型中，生成完整的模型文件

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct"
LORA_PATH = "./qlora-output"
MERGED_PATH = "./merged-model"   # 合并后的模型保存位置

print("加载基础模型（合并时用 float16，不量化）...")
# 注意：合并时不使用量化，保持 float16 精度
base_model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,  # 使用 float16，不是4-bit
    device_map="auto",
    trust_remote_code=True,
)

print("加载 LoRA 权重...")
model = PeftModel.from_pretrained(base_model, LORA_PATH)

print("合并权重（这需要几分钟）...")
merged_model = model.merge_and_unload()  # 将 LoRA 合并进基础模型

print(f"保存合并后的模型到 {MERGED_PATH}...")
merged_model.save_pretrained(MERGED_PATH, safe_serialization=True)

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
tokenizer.save_pretrained(MERGED_PATH)

print("完成！合并后的模型可以直接用于推理或转换为 GGUF 格式部署到 Ollama。")