Custom callback support

Author: Meorge

aitextgen/aitextgen.py

@@ -581,6 +581,7 @@ def train(
         progress_bar_refresh_rate: int = 20,
         freeze_layers: bool = False,
         num_layers_freeze: int = None,
+        custom_callbacks: dict = {},
         **kwargs,
     ) -> None:
 
@@ -667,6 +668,7 @@ def train(
                 save_gdrive,
                 avg_loss_smoothing,
                 is_gpu_used,
+                custom_callbacks
             )
         )
 

aitextgen/train_callback.py

@@ -23,6 +23,7 @@ def __init__(
         save_gdrive,
         avg_loss_smoothing,
         is_gpu_used,
+        custom_callbacks
     ):
         self.training_bar = None
         self.model = model
@@ -39,6 +40,7 @@ def __init__(
         self.steps = 0
         self.current_loss = None
         self.prev_avg_loss = None
+        self.custom_callbacks = custom_callbacks
 
     @property
     def save_every_check(self):
@@ -54,17 +56,22 @@ def on_train_begin(self, args, state, control, **kwargs):
                 file=sys.stdout,
             )
 
+        self.custom_callbacks.get('on_train_begin', lambda: None)()
+
     def on_train_end(self, args, state, control, **kwargs):
         if state.is_local_process_zero:
             self.training_bar.close()
             self.training_bar = None
 
+        self.custom_callbacks.get('on_train_end', lambda: None)()
+
     def on_evaluate(self, args, state, control, metrics, **kwargs):
+        print(f'''on_evaluate called: train_loss={metrics.get('train_loss')}''')
         if state.is_local_process_zero:
             self.current_loss = float(metrics.get("train_loss"))
 
     def on_step_end(self, args, state, control, **kwargs):
-
+        print(f'on_step_end called\nargs={args}\nstate={state}\ncontrol={control}\nkwargs={kwargs}\n=====')
         if state.is_local_process_zero:
             self.steps += 1
             avg_loss = 0
@@ -100,6 +107,8 @@ def on_step_end(self, args, state, control, **kwargs):
                 if self.current_loss:
                     self.training_bar.set_description(desc)
 
+            self.custom_callbacks.get('on_step_end', lambda steps, max, curr, avg, trainer: None)(self.steps, state.max_steps, self.current_loss, avg_loss, self.trainer)
+
             if self.save_every > 0 and self.steps % self.save_every == 0:
                 self.save_pytorch_model()
 
@@ -138,6 +147,8 @@ def generate_sample_text(self):
 
         self.training_bar.write("=" * 10)
 
+        self.custom_callbacks.get('on_sample_text_generated', lambda texts: None)(gen_texts)
+
     def save_pytorch_model(self):
         # only runs on state.is_local_process_zero
         self.training_bar.write(
@@ -154,6 +165,8 @@ def save_pytorch_model(self):
                     os.path.join("/content/drive/My Drive/", self.run_id, pt_file),
                 )
 
+        self.custom_callbacks.get('on_model_saved', lambda curr, out: None)(self.steps, self.output_dir)
+
     def average_loss(self, current_loss, prev_avg_loss, smoothing):
         if prev_avg_loss is None:
             return current_loss