trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = ds_train['train'],
    args = SFTConfig(
        dataset_text_field = "text",
        per_device_train_batch_size = 1,
        gradient_accumulation_steps = 1, 
        warmup_steps = 5,
        num_train_epochs = 1,
        learning_rate = 2e-5, 
        logging_steps = 5,
        optim = "adamw_8bit",
        weight_decay = 0.01,
        lr_scheduler_type = "linear",
        seed = 3407,
        report_to = "wandb", 
        output_dir = "outputs",
        save_strategy = "steps",
        save_steps = 50,
        save_total_limit = 2,
    ),
)