from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model
import torch

MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"  # 可换成Qwen、ChatGLM等
OUTPUT_DIR = "./lora-weights"

# 1. 加载模型和分词器
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    load_in_4bit=True,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# 2. 配置 LoRA
lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj","v_proj"], 
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)

# 3. 加载数据
dataset = load_dataset("json", data_files="dataset.json")
def tokenize(batch):
    text = batch["instruction"] + batch["input"] + batch["output"]
    return tokenizer(text, truncation=True, max_length=512)
tokenized_dataset = dataset.map(tokenize, batched=True)

# 4. 训练参数
training_args = TrainingArguments(
    per_device_train_batch_size=1,
    gradient_accumulation_steps=8,
    learning_rate=2e-4,
    num_train_epochs=3,
    fp16=True,
    logging_steps=10,
    output_dir=OUTPUT_DIR,
    save_strategy="epoch",
)

# 5. Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
)
trainer.train()

# 6. 保存权重
model.save_pretrained(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)
print(f"训练完成，LoRA 权重已保存在 {OUTPUT_DIR}")