Update app.py

app.py

@@ -1,4 +1,6 @@
 import tensorflow as tf 
+from tensorflow.keras.optimizers import Adam
+from tensorflow.keras.optimizers.schedules import LearningRateSchedule
 import tensorflow_text as text
 from tensorflow.train import Checkpoint
 import pandas as pd
@@ -34,8 +36,29 @@ transformer = Transformer(vocab_size_enc = VOCAB_SIZE,
                           FFN_units = FFN_UNITS,
                           nb_proj = NB_PROJ,
                           dropout_rate = DROPOUT_RATE)
-                          
-ckpt = Checkpoint()   
+
+class CustomSchedule(LearningRateSchedule):
+    def __init__(self, d_model, warmup_steps = 4000):
+        super(CustomSchedule, self).__init__()
+        
+        self.d_model = tf.cast(d_model, tf.float32)
+        self.warmup_steps = warmup_steps
+    
+    def __call__(self, step):
+        arg1 = tf.math.rsqrt(step)
+        arg2 = step * (self.warmup_steps**-1.5)
+        
+        return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)
+
+leaning_rate = CustomSchedule(D_MODEL)
+
+optimizer = Adam(leaning_rate,
+                 beta_1=0.9,
+                 beta_2=0.98,
+                 epsilon=1e-9)
+                                                    
+ckpt = tf.train.Checkpoint(transformer = transformer,
+                           optimizer = optimizer)  
 ckpt.restore("ckpt-10")  
 print("Checkpoint Restaurado")