# 安装: pip install transformers peft accelerate bitsandbytes datasets

from transformers import (
    AutoModelForCausalLM, AutoTokenizer,
    TrainingArguments, Trainer, DataCollatorForLanguageModeling
)
from peft import LoraConfig, get_peft_model, TaskType
from datasets import Dataset
import torch

# 1. 加载基座模型（用 4-bit 量化节省显存）
model_name = "Qwen/Qwen2-7B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    load_in_4bit=True  # 4-bit 量化，显存降低 4 倍
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# 2. 配置 LoRA
lora_config = LoraConfig(
    r=16,                   # 秩，越大效果越好但参数越多
    lora_alpha=32,          # 缩放因子
    target_modules=[        # 在哪些层插入 LoRA
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ],
    lora_dropout=0.1,       # Dropout 防止过拟合
    bias="none",
    task_type=TaskType.CAUSAL_LM
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# 输出: trainable params: 41,943,040 || all params: 7,057,072,128 || trainable%: 0.59%
# 原始 7B 参数中只训练 4200 万（0.59%）！

# 3. 准备训练数据
training_data = [
    {"text": "用户：帮我写一段 Python 代码，实现快速排序。\n助手：好的，以下是快速排序的 Python 实现：\n```python\ndef quicksort(arr):\n    if len(arr) <= 1:\n        return arr\n    pivot = arr[len(arr)//2]\n    left = [x for x in arr if x < pivot]\n    middle = [x for x in arr if x == pivot]\n    right = [x for x in arr if x > pivot]\n    return quicksort(left) + middle + quicksort(right)\n```"},
    # ... 更多训练数据
]

dataset = Dataset.from_list(training_data)

def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=512,
        padding="max_length"
    )

tokenized_dataset = dataset.map(tokenize_function)

# 4. 训练
training_args = TrainingArguments(
    output_dir="./qwen2-lora-code",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,  # 等效 batch size = 16
    num_train_epochs=3,
    learning_rate=2e-4,
    fp16=True,
    logging_steps=10,
    save_strategy="epoch",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
trainer.train()

# 5. 保存 LoRA 权重（只有几 MB！）
model.save_pretrained("./qwen2-lora-code-final")

# 6. 推理时加载 LoRA
from peft import PeftModel

base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    load_in_4bit=True
)
model = PeftModel.from_pretrained(base_model, "./qwen2-lora-code-final")

inputs = tokenizer("用 Python 实现冒泡排序", return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens=256)
print(tokenizer.decode(outputs[0]))