from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model
from datasets import load_dataset
model_name = "mistral-7b"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# LoRA 설정
lora_config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
)
model = get_peft_model(model, lora_config)
dataset = load_dataset("your_dataset")
training_args = TrainingArguments(
output_dir="./lora-out",
per_device_train_batch_size=2,
gradient_accumulation_steps=2,
learning_rate=2e-4,
num_train_epochs=3,
fp16=True,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=dataset["train"]
)
trainer.train()
model.save_pretrained("./lora-out")