Custom evaluate for Trainer Class

From the official documentation it say that “You can also subclass and override this method to inject custom behavior.” for the evaluate function, but I have a question about how make this, it should be like this?

class NewTrainer(Trainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model.eval()
    def create_optimizer_and_scheduler(self, num_training_steps: int):
        self.optimizer = AdamW(model.parameters(), lr=self.args.learning_rate, correct_bias=False)
        self.total_steps = self.args.per_device_train_batch_size * self.args.num_train_epochs
        self.lr_scheduler = get_linear_schedule_with_warmup(self.optimizer, num_warmup_steps = 0, num_training_steps = self.total_steps)
    def evaluate(self, test_batches): 
            model.eval()
            tmp_m, tmp_m_sentence = test_model(model, test_batches)

where test_model is a function that evaluate the model in the test datasets, in this function I split the text in sub chuncks and make the attention mechanism for the classification.

1 Like