# Define training arguments for Seq2Seq model
training_args = Seq2SeqTrainingArguments(
    output_dir="dhongi",  # Directory to save the model checkpoints
    eval_strategy="epoch",  # Perform evaluation at the end of each epoch
    learning_rate=2e-5,  # Set a low learning rate for stable training
    per_device_train_batch_size=16,  # Set batch size for training per device (GPU/CPU)
    per_device_eval_batch_size=16,  # Set batch size for evaluation per device
    weight_decay=0.01,  # Apply weight decay for regularization to avoid overfitting
    save_total_limit=3,  # Keep only the last 3 model checkpoints to save storage
    num_train_epochs=2,  # Number of epochs for training
    predict_with_generate=True,  # Use generate() method to make predictions (important for seq2seq models)
    fp16=True,  # Use mixed precision training for faster training and reduced memory usage on GPUs
    # bf16=True,  # Uncomment to use bfloat16 precision for XPU hardware (like Intel's Xe)
    push_to_hub=False,  # Do not push the trained model to the Hugging Face Hub after training
)

# Initialize the trainer with the model, arguments, and datasets
trainer = Seq2SeqTrainer(
    model=model,  # The model to train
    args=training_args,  # Pass the training arguments defined above
    train_dataset=train_data,  # The dataset to use for training
    eval_dataset=val_data,  # The dataset to use for evaluation
    tokenizer=tokenizer,  # The tokenizer to process inputs and outputs
    data_collator=data_collator,  # The data collator used to batch the data
    compute_metrics=compute_metrics,  # Function to compute metrics during evaluation
)

# Start training the model with the defined parameters
trainer.train()  # Begin the training process