File size: 886 Bytes
52c8e55 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig
from common import BASE
ds = load_dataset("json", data_files="data.jsonl")["train"]
trainer = SFTTrainer(
model=BASE,
train_dataset=ds,
peft_config=LoraConfig(
r=32, lora_alpha=64, lora_dropout=0.05,
target_modules="all-linear", task_type="CAUSAL_LM",
),
args=SFTConfig(
output_dir="out",
num_train_epochs=3,
per_device_train_batch_size=2,
gradient_accumulation_steps=4,
learning_rate=2e-4,
bf16=False,
max_length=1024,
logging_steps=10,
save_strategy="no",
dataloader_num_workers=0, # avoids Windows DataLoader errors
report_to="none",
),
)
trainer.train()
trainer.save_model("lora-out")
print("Adapter saved to lora-out/") |