src/trainer.py

from transformers import Trainer, TrainingArguments
import os

def train_model(model, datasets):
    training_args = TrainingArguments(
        output_dir="models/model_output",
        evaluation_strategy="epoch",
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        num_train_epochs=3,
        weight_decay=0.01
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=datasets["train"],
        eval_dataset=datasets["test"]
    )

    trainer.train()
    model.save_pretrained("models/model_output")