Update app.py

app.py

@@ -81,7 +81,7 @@ class PositionalEmbedding(layers.Layer):
     return output
 
 
-class AttentionalPooling(layers.Layer):
+class AttentionalPooling(tf.keras.layers.Layer):
     def __init__(self, embed_dim, num_heads=6):
         super().__init__()
         self.embed_dim = embed_dim
@@ -100,7 +100,7 @@ class AttentionalPooling(layers.Layer):
         return self.norm(attn_output)
 
 
-class TransformerBlock(layers.Layer):
+class TransformerBlock(tf.keras.layers.Layer):
   def __init__(self, embed_dim, dense_dim, num_heads, dropout_rate=0.1, ln_epsilon=1e-6, is_multimodal=False, **kwargs):
     super().__init__(**kwargs)
     self.embed_dim = embed_dim
@@ -128,7 +128,7 @@ class TransformerBlock(layers.Layer):
 
 
     # Feed-Forward Network
-    self.dense_proj = keras.Sequential([
+    self.dense_proj = tf.keras.Sequential([
         layers.Dense(self.dense_dim, activation="gelu"),
         layers.Dropout(self.dropout_rate),
         layers.Dense(self.embed_dim)
@@ -279,7 +279,7 @@ for layer in vit_tiny_model.layers:
     layer.trainable = True
 
 
-class CoCaEncoder(keras.Model):
+class CoCaEncoder(tf.keras.Model):
     def __init__(self,
             vit, **kwargs):
 
@@ -317,7 +317,7 @@ class CoCaEncoder(keras.Model):
 
 
 
-class CoCaDecoder(keras.Model):
+class CoCaDecoder(tf.keras.Model):
     def __init__(self,
             cls_token_id,
             num_heads,
@@ -368,7 +368,7 @@ class CoCaDecoder(keras.Model):
 
 
 # день 6
-class CoCaModel(keras.Model):
+class CoCaModel(tf.keras.Model):
     def __init__(self,
         vit,
         cls_token_id,
@@ -491,7 +491,7 @@ dummy_features = tf.zeros((1, 3, img_size, img_size), dtype=tf.float32)
 dummy_captions = tf.zeros((1, sentence_length-1), dtype=tf.int64)
 _ = coca_model((dummy_features, dummy_captions))
 
-optimizer = keras.optimizers.Adam(learning_rate=1e-4)
+optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)
 coca_model.compile(optimizer)
 
 save_dir = "models/"
@@ -540,7 +540,7 @@ class BahdanauAttention(layers.Layer):
 
 
 
-class ImageCaptioningModel(keras.Model):
+class ImageCaptioningModel(tf.keras.Model):
     def __init__(self, vocab_size, max_caption_len, embedding_dim=512, lstm_units=512, dropout_rate=0.5, **kwargs):
         super().__init__(**kwargs)