import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
    TrainingArguments,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer  # Supervised Fine-Tuning Trainer
from datasets import load_dataset

# ========================
# 1. 配置模型和分词器
# ========================
model_name = "meta-llama/Llama-2-7b-hf"  # 或你有权访问的本地路径

# 4-bit 量化配置（可选但推荐）
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token  # LLaMA 没有 pad_token

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)

# 准备模型用于 k-bit 训练
model = prepare_model_for_kbit_training(model)

# ========================
# 2. 配置 LoRA
# ========================
peft_config = LoraConfig(
    r=8,                         # 秩（rank）
    lora_alpha=16,               # 缩放因子
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",       # 因果语言模型（decoder-only）
    target_modules=["q_proj", "v_proj"]  # 常见目标模块（LLaMA/Mistral）
)

model = get_peft_model(model, peft_config)
model.print_trainable_parameters()  # 查看可训练参数量

# ========================
# 3. 加载数据集
# ========================
# 假设你的数据是 JSONL 格式，包含 "text" 字段（或自定义 prompt）
dataset = load_dataset("json", data_files="your_data.jsonl", split="train")

# 或使用 Hugging Face 上的数据集
# dataset = load_dataset("your_username/your_dataset", split="train")

# ========================
# 4. 训练配置
# ========================
training_args = TrainingArguments(
    output_dir="./lora-finetuned",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    num_train_epochs=3,
    logging_steps=10,
    save_strategy="epoch",
    evaluation_strategy="no",
    optim="paged_adamw_32bit",   # 支持分页优化器，减少 OOM
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    fp16=False,                  # 使用 bf16 或关闭
    bf16=True,
    max_grad_norm=0.3,
    group_by_length=True,
    report_to="none",            # 可改为 "wandb" 或 "tensorboard"
)

# ========================
# 5. 使用 SFTTrainer
# ========================
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
    dataset_text_field="text",   # 数据集中包含文本的字段名
    max_seq_length=512,
    packing=False,               # 若为 True，会自动拼接样本以提高效率
)

# 开始训练
trainer.train()

# 保存适配器
trainer.save_model()  # 仅保存 LoRA 权重