Support for Callback Functions in Trainer Class

[Lucas-Fernandes-Martins]

Hi,

I have a quick question concerning using callback functions inside a training loop using the Trainer classhttps://github.com/google-research/kauldron/blob/main/kauldron/train/trainer_lib.py.

I ask because in the Transformers library, you have a TrainerCallback class https://huggingface.co/docs/transformers/v4.50.0/en/main_classes/callback#transformers.TrainerCallback whose objects you can pass as arguments to a training loop.

I'd like to ask if there's indeed something similar in kauldron, and I misread the documentation (in which case I apologise), or if this is something to be considered as a potential future feature.

Thank you very much for your time.

Simple example using the Transformers library for context:

from transformers import Trainer, TrainingArguments, TrainerCallback

# Define a custom callback
class MyCallback(TrainerCallback):
    def on_epoch_end(self, args, state, control, **kwargs):
        print(f"Epoch {state.epoch} has ended.")

# model = ...
# train_dataset = ...
# eval_dataset = ...
training_args = TrainingArguments(output_dir="./results", num_train_epochs=3)

# Initialize the Trainer and pass the custom callback
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    callbacks=[MyCallback()],  # Add the custom callback here
)

trainer.train()