# /// script
# dependencies = ["unsloth", "trl>=0.12.0", "datasets", "trackio"]
# ///
from unsloth import FastLanguageModel
from trl import SFTTrainer, SFTConfig
from datasets import load_dataset
model, tokenizer = FastLanguageModel.from_pretrained(
"Qwen/Qwen2.5-0.5B",
load_in_4bit=True,
max_seq_length=2048,
)
model = FastLanguageModel.get_peft_model(
model,
r=16,
lora_alpha=32,
lora_dropout=0,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj"],
)
dataset = load_dataset("trl-lib/Capybara", split="train")
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=dataset,
args=SFTConfig(
output_dir="./output",
push_to_hub=True,
hub_model_id="username/my-model",
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
num_train_epochs=1,
learning_rate=2e-4,
report_to="trackio",
),
)
trainer.train()
trainer.push_to_hub()