# Import required libraries
import torch
from transformers import MT5ForConditionalGeneration, MT5Tokenizer

# Initialize model and tokenizer
model_name = "google/mt5-small"
tokenizer = MT5Tokenizer.from_pretrained(model_name)
model = MT5ForConditionalGeneration.from_pretrained(model_name)

# Tokenize the datasets
train_inputs = tokenize_data(df_4_train, max_length=512, max_target_length=100)
val_inputs = tokenize_data(df_4_val, max_length=512, max_target_length=100)
test_inputs = tokenize_data(df_4_test, max_length=512, max_target_length=100)

# Set up training arguments
training_args = Seq2SeqTrainingArguments(
    output_dir="./results",
    eval_strategy="epoch",
    learning_rate=1e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=5,
    weight_decay=0.01,
    save_total_limit=2,
    predict_with_generate=True,
    save_safetensors=False
)