# /// script
# dependencies = ["trl", "peft", "transformers", "datasets", "accelerate", "torch"]
# ///

print("Starting...")

from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig

print("Loading dataset...")
ds = load_dataset("trl-lib/Capybara", split="train[:100]")
print(f"Loaded {len(ds)} examples")

config = SFTConfig(
    output_dir="test-output",
    num_train_epochs=1,
    per_device_train_batch_size=1,
    logging_steps=5,
    report_to="none",
    push_to_hub=True,
    hub_model_id="luiscosio/qwen25-minimal-test",
)

peft_config = LoraConfig(r=8, lora_alpha=16, task_type="CAUSAL_LM", target_modules=["q_proj", "v_proj"])

print("Creating trainer...")
trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B",
    train_dataset=ds,
    args=config,
    peft_config=peft_config,
)

print("Training...")
trainer.train()
trainer.push_to_hub()
print("Done!")