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/")