Implement custom trainer callback

custom_trainer_callback.py

@@ -0,0 +1,40 @@
+"""
+Custom trainer callback to store config and evaluation results
+"""
+import os
+import json
+from copy import deepcopy
+from transformers import TrainerCallback
+
+
+class CustomTrainerCallback(TrainerCallback):
+    """Keep track of configs and evaluation results"""
+
+    def __init__(self, args):
+        self.config = args
+        self.checkpoint_metrics = {}
+        self.last_metrics = {'eval_accuracy': 0}
+
+
+    def on_evaluate(self, args, state, control, metrics, **kwargs):
+        """Get evaluation accuracy after each evaluation step"""
+        metric_name = 'eval_accuracy'
+        if metric_name in metrics:
+            self.last_metrics[metric_name] = metrics[metric_name]
+
+
+    def on_save(self, args, state, control, **kwargs):
+        """Store evaluation results after each save step"""
+        self.checkpoint_metrics[state.global_step] = deepcopy(self.last_metrics)
+
+
+    def on_train_end(self, _train_args, state, control, **kwargs):
+        """Store results and configs after training is done"""
+        chk_dirs = [d for d in os.listdir(self.config['name']) if d.startswith('checkpoint-')]
+        checkpoint_steps = [int(c.replace('checkpoint-', '').strip()) for c in chk_dirs]
+        self.config['checkpoints'] = {c: self.checkpoint_metrics[c] for c in checkpoint_steps}
+        config_path = os.path.join(self.config['name'], "config.json")
+        with open(config_path, "w", encoding='utf-8') as outfile:
+            json.dump(self.config, outfile, indent=4)
+        best_accuracy = max([v['eval_accuracy'] for v in self.checkpoint_metrics.values()])
+        print(f'==== BEST EVAL ACCURACY : {best_accuracy} ====')

finetune_t5.py

@@ -16,6 +16,7 @@
 )
 from Levenshtein import ratio as sim_ratio
 from load_datasets import DatasetManager
+from custom_trainer_callback import CustomTrainerCallback
 
 
 # Args
@@ -162,8 +163,9 @@ def compute_metrics(eval_preds):
 )
 
 # Callbacks
+callback = CustomTrainerCallback(vars(args))
 early_stopping = EarlyStoppingCallback(early_stopping_patience=args.early_stopping_patience)
-callbacks = [early_stopping]
+callbacks = [callback, early_stopping]
 
 # Training args
 training_args = Seq2SeqTrainingArguments(