from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="./model",
    overwrite_output_dir=True,
    push_to_hub=True,
    hub_model_id="Ransaka/sinhala-bert-yt-comments",
    num_train_epochs=2,
    per_device_train_batch_size=32,
    save_steps=5_000,
    logging_steps = 1000,
    save_total_limit=2,
    use_mps_device = True, # disable this if you're running non-mac env
    hub_private_repo = False, # please set true if you want to save model privetly
    save_safetensors= True,
    learning_rate = 1e-4,
    report_to='wandb'
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=tokenized_dataset_train
)