Update train.py

train.py

@@ -1,4 +1,3 @@
-
 import tensorflow as tf
 from transformers import TFXLMRobertaModel, AutoTokenizer, TFAutoModel
 from datasets import load_dataset, concatenate_datasets
@@ -6,8 +5,6 @@ from datetime import datetime
 import logging
 from pyprojroot.here import here
 
-
-
 class mean_pooling_layer(tf.keras.layers.Layer):
     def __init__(self):
         super(mean_pooling_layer, self).__init__()
@@ -39,6 +36,7 @@ def create_model():
     output_id = base_student_model.roberta(input_ids_id, attention_mask=attention_mask_id).last_hidden_state[:,0,:]
 
     student_model = tf.keras.Model(inputs=[input_ids_en, attention_mask_en, input_ids_id, attention_mask_id], outputs=[output_en, output_id])
+    student_model.load_weights("disk/model/2023-05-25_07-52-43/multiqa-Mmini-L6-H384.h5")
     return student_model
 
 class sentence_translation_metric(tf.keras.callbacks.Callback):
@@ -74,21 +72,19 @@ class sentence_translation_metric(tf.keras.callbacks.Callback):
         logs["val_avg_acc"] = avg_acc
 
 
-class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
-  def __init__(self, d_model, warmup_steps=100000):
+class ConstantScheduler(tf.keras.optimizers.schedules.LearningRateSchedule):
+  def __init__(self, max_lr, warmup_steps=5000):
     super().__init__()
-
-    self.d_model = d_model
-    self.d_model = tf.cast(self.d_model, tf.float32)
-
+    self.max_lr = tf.cast(max_lr, tf.float32)
     self.warmup_steps = warmup_steps
 
   def __call__(self, step):
-    step = tf.cast(step, dtype=tf.float32)
-    arg1 = tf.math.rsqrt(step)
-    arg2 = step * (self.warmup_steps ** -1.5)
+    step = tf.cast(step, tf.float32)
+    condition  = tf.cond(step < self.warmup_steps, lambda: step / self.warmup_steps, lambda: 1.0)
+    return self.max_lr * condition 
+  
 
-    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)
+          
 
 
 if __name__ == "__main__":
@@ -101,8 +97,8 @@ if __name__ == "__main__":
         dataset_1 = concatenate_datasets([dataset_1, dataset[split]])
     
 
-    batch_size = 384
-    dataset = dataset_1.train_test_split(test_size=0.01, shuffle=True)
+    batch_size = 512
+    dataset = dataset_1.train_test_split(test_size=0.005, shuffle=True)
     train_dataset = dataset["train"]
     val_dataset = dataset["test"]
     print(val_dataset.shape)
@@ -127,7 +123,8 @@ if __name__ == "__main__":
     val_dataset = val_dataset.batch(batch_size, drop_remainder=True).cache()
 
 
-    learning_rate = CustomSchedule(384)
+    learning_rate = ConstantScheduler(1e-3, warmup_steps=10000)
+
 
     optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98,
                                         epsilon=1e-9)
@@ -137,7 +134,7 @@ if __name__ == "__main__":
     loss = tf.keras.losses.MeanSquaredError() 
     
     date_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
-    output_path = here(f"disk/model/{date_time}/multiqa-mpnet-dot-v1.h5")
+    output_path = here(f"disk/model/{date_time}/multiqa-Mmini-L6-H384.h5")
 
     model_checkpoint = tf.keras.callbacks.ModelCheckpoint(
                     filepath = output_path,
@@ -146,16 +143,16 @@ if __name__ == "__main__":
                     mode = 'auto',
                     verbose = 1,
                     save_best_only = True,
-                    initial_value_threshold = 0.1 
+                    initial_value_threshold = 0.5,
                     )
 
     early_stopping = tf.keras.callbacks.EarlyStopping(
                     monitor = "val_avg_acc",
                     mode = 'auto',
                     restore_best_weights=False,
-                    patience = 2,
+                    patience = 4,
                     verbose=1,
-                    start_from_epoch = 25,
+                    start_from_epoch = 5,
     )
 
 
@@ -169,13 +166,7 @@ if __name__ == "__main__":
                     append = False
     )
 
-    reduce_rl = tf.keras.callbacks.ReduceLROnPlateau(
-        monitor = "",
-        factor = 0.1,
-        patience = 2,
-        min_lr = 1e-6,
-        verbose = 1
-    )
+    
 
 
     callbacks = [sentence_translation_metric(), model_checkpoint, csv_logger,early_stopping]
@@ -192,7 +183,8 @@ if __name__ == "__main__":
 
     student_model.fit(train_dataset, epochs=20, validation_data=val_dataset, callbacks=callbacks)
 
-    last_epoch_save = here(f"disk/model/last_epoch/{date_time}/multiqa-mpnet-dot-v1.h5")
+
+    last_epoch_save = here(f"disk/model/last_epoch/{date_time}.h5")
     student_model.save_weights(last_epoch_save)