# %% Cell 6: Training Arguments
from transformers import TrainingArguments

total_steps = len(dataset) // (BATCH_SIZE * GRAD_ACCUM) * EPOCHS
print(f"Estimated steps: {total_steps:,}")
print(f"Estimated time : ~{total_steps / 8 / 3600:.1f} hours (at ~8 it/s on A100)\n")

args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    overwrite_output_dir=True,
    num_train_epochs=EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    gradient_accumulation_steps=GRAD_ACCUM,
    learning_rate=LR,
    lr_scheduler_type="cosine",
    warmup_steps=500,
    optim="adamw_8bit",
    weight_decay=0.01,
    max_grad_norm=0.3,
    bf16=torch.cuda.is_bf16_supported(),
    fp16=not torch.cuda.is_bf16_supported(),
    logging_steps=50,
    save_strategy="steps",
    save_steps=2000,
    save_total_limit=2,
    group_by_length=True,
    remove_unused_columns=False,
    report_to="none",
    seed=3407,
)

print("✅ Training arguments ready.")
print(f"  Effective batch: {BATCH_SIZE * GRAD_ACCUM}")
print(f"  LR: {LR}")
print(f"  Context: {MAX_SEQ_LENGTH}")
print(f"  Epochs: {EPOCHS}")
print("\n→ Run Cell 7 to initialize the trainer.")