Update model.py

model.py

@@ -2,16 +2,21 @@ import pickle
 import tensorflow as tf
 import pandas as pd
 import numpy as np
+import keras
+from keras import layers
 from keras.applications import efficientnet
-from keras.layers import TextVectorization
-IMAGE_SIZE = (299, 299)
+
+
 # CONTANTS
-MAX_LENGTH = 40
+MAX_LENGTH = 80
 # VOCABULARY_SIZE = 10000
 BATCH_SIZE = 32
 BUFFER_SIZE = 1000
 EMBEDDING_DIM = 512
 UNITS = 512
+IMAGE_SIZE = (299, 299)
+FF_DIM = 512
+VOCAB_SIZE = 19774
 
 
 # LOADING DATA
@@ -33,96 +38,108 @@ idx2word = tf.keras.layers.StringLookup(
 
 # MODEL
 def CNN_Encoder():
-
-    inception_v3 = efficientnet.EfficientNetB0(
+    base_model = efficientnet.EfficientNetB0(
         input_shape=(*IMAGE_SIZE, 3),
         include_top=False,
         weights="imagenet",
     )
-
-    output = inception_v3.output
-    output = tf.keras.layers.Reshape(
-        (-1, output.shape[-1]))(output)
-
-    cnn_model = tf.keras.models.Model(inception_v3.input, output)
+    # We freeze our feature extractor
+    base_model.trainable = False
+    base_model_out = base_model.output
+    base_model_out = layers.Reshape((-1, base_model_out.shape[-1]))(base_model_out)
+    cnn_model = keras.models.Model(base_model.input, base_model_out)
     return cnn_model
-       
 
 
+#the extracted image features are passed to transformer based encoder that generates the new representation of inputs
 class TransformerEncoderLayer(tf.keras.layers.Layer):
-
-    def __init__(self, embed_dim, num_heads):
-        super().__init__()
-        self.layer_norm_1 = tf.keras.layers.LayerNormalization()
-        self.layer_norm_2 = tf.keras.layers.LayerNormalization()
-        self.attention = tf.keras.layers.MultiHeadAttention(
-            num_heads=num_heads, key_dim=embed_dim)
-        self.dense = tf.keras.layers.Dense(embed_dim, activation="relu")
-    
-
-    def call(self, x, training):
-        x = self.layer_norm_1(x)
-        x = self.dense(x)
-
-        attn_output = self.attention(
-            query=x,
-            value=x,
-            key=x,
+    def __init__(self, embed_dim, dense_dim, num_heads, **kwargs):
+        super().__init__(**kwargs)
+        self.embed_dim = embed_dim
+        self.dense_dim = dense_dim
+        self.num_heads = num_heads
+        #captures dependencies between different parts of the input sequence.
+        self.attention_1 = layers.MultiHeadAttention(
+            num_heads=num_heads, key_dim=embed_dim, dropout=0.0
+        )
+        self.layernorm_1 = layers.LayerNormalization()
+        self.layernorm_2 = layers.LayerNormalization()
+        self.dense_1 = layers.Dense(embed_dim, activation="relu")
+
+    def call(self, inputs, training, mask=None):
+        inputs = self.layernorm_1(inputs)
+        inputs = self.dense_1(inputs)
+
+        attention_output_1 = self.attention_1(
+            query=inputs,
+            value=inputs,
+            key=inputs,
             attention_mask=None,
-            training=training
+            training=training,
         )
-
-        x = self.layer_norm_2(x + attn_output)
-        return x
-
+        out_1 = self.layernorm_2(inputs + attention_output_1)
+        return out_1
 
 class Embeddings(tf.keras.layers.Layer):
+    def __init__(self, sequence_length, vocab_size, embed_dim, **kwargs):
+        super().__init__(**kwargs)
+        self.token_embeddings = layers.Embedding(
+            input_dim=vocab_size, output_dim=embed_dim
+        )
+        self.position_embeddings = layers.Embedding(
+            input_dim=sequence_length, output_dim=embed_dim
+        )
+        self.sequence_length = sequence_length
+        self.vocab_size = vocab_size
+        self.embed_dim = embed_dim
+        self.embed_scale = tf.math.sqrt(tf.cast(embed_dim, tf.float32))#4== 2.0
 
-    def __init__(self, vocab_size, embed_dim, max_len):
-        super().__init__()
-        self.token_embeddings = tf.keras.layers.Embedding(
-            vocab_size, embed_dim)
-        self.position_embeddings = tf.keras.layers.Embedding(
-            max_len, embed_dim, input_shape=(None, max_len))
-    
-
-    def call(self, input_ids):
-        length = tf.shape(input_ids)[-1]
-        position_ids = tf.range(start=0, limit=length, delta=1)
-        position_ids = tf.expand_dims(position_ids, axis=0)
+    def call(self, inputs):
+        length = tf.shape(inputs)[-1] #gets the length of input sequence or token
+        positions = tf.range(start=0, limit=length, delta=1)# Create a tensor of positions from 0 to length-1(if length is 5 then positions = [0,1,2,3,4])
+        embedded_tokens = self.token_embeddings(inputs)# Embed tokens using the token embedding layer(self.token_embeddings layer converts these discrete token indices into continuous, high-dimensional vectors (embeddings). Each token in the caption is represented by a dense vector that captures its semantic meaning.)
+        embedded_tokens = embedded_tokens * self.embed_scale  # Scale the embedded tokens by the square root of the embedding dimension
+        embedded_positions = self.position_embeddings(positions) # Embed positions using the position embedding layer
+        return embedded_tokens + embedded_positions 
 
-        token_embeddings = self.token_embeddings(input_ids)
-        position_embeddings = self.position_embeddings(position_ids)
+    def compute_mask(self, inputs, mask=None):
+        return tf.math.not_equal(inputs, 0)
 
-        return token_embeddings + position_embeddings
 
 
 class TransformerDecoderLayer(tf.keras.layers.Layer):
-
-    def __init__(self, embed_dim, units, num_heads):
-        super().__init__()
-        self.embedding = Embeddings(
-            tokenizer.vocabulary_size(), embed_dim, MAX_LENGTH)
-
-        self.attention_1 = tf.keras.layers.MultiHeadAttention(
+    
+    def __init__(self, embed_dim, ff_dim, num_heads, **kwargs):
+        super().__init__(**kwargs)
+        self.embed_dim = embed_dim
+        self.ff_dim = ff_dim
+        self.num_heads = num_heads
+        #Defines two multi-head self-attention mechanisms (attention_1 and attention_2) with dropout.
+        self.attention_1 = layers.MultiHeadAttention(
             num_heads=num_heads, key_dim=embed_dim, dropout=0.1
         )
-        self.attention_2 = tf.keras.layers.MultiHeadAttention(
+        self.attention_2 = layers.MultiHeadAttention(
             num_heads=num_heads, key_dim=embed_dim, dropout=0.1
         )
+        #two feed-forward network layers with ReLU activation.
+        self.ffn_layer_1 = layers.Dense(ff_dim, activation="relu")
+        self.ffn_layer_2 = layers.Dense(embed_dim)
+#normalization to the outputs of attention mechanisms 
+        self.layernorm_1 = layers.LayerNormalization()
+        self.layernorm_2 = layers.LayerNormalization()
+        self.layernorm_3 = layers.LayerNormalization()
+#Utilizes a PositionalEmbedding layer to incorporate positional information into the inputs.
+        self.embedding = Embeddings(
+            embed_dim=EMBEDDING_DIM,
+            sequence_length=MAX_LENGTH,
+            vocab_size=VOCAB_SIZE,
+        )
+    #softmax activation for generating predictions.
+        self.out = layers.Dense(VOCAB_SIZE, activation="softmax")
 
-        self.layernorm_1 = tf.keras.layers.LayerNormalization()
-        self.layernorm_2 = tf.keras.layers.LayerNormalization()
-        self.layernorm_3 = tf.keras.layers.LayerNormalization()
-
-        self.ffn_layer_1 = tf.keras.layers.Dense(units, activation="relu")
-        self.ffn_layer_2 = tf.keras.layers.Dense(embed_dim)
-
-        self.out = tf.keras.layers.Dense(tokenizer.vocabulary_size(), activation="softmax")
-
-        self.dropout_1 = tf.keras.layers.Dropout(0.3)
-        self.dropout_2 = tf.keras.layers.Dropout(0.5)
-    
+        self.dropout_1 = layers.Dropout(0.3)
+        self.dropout_2 = layers.Dropout(0.5)
+        self.supports_masking = True # can be useful for handling variable-length sequences.
 
     def call(self, input_ids, encoder_output, training, mask=None):
         embeddings = self.embedding(input_ids)
@@ -165,6 +182,23 @@ class TransformerDecoderLayer(tf.keras.layers.Layer):
         preds = self.out(ffn_out)
         return preds
 
+    def get_causal_attention_mask(self, inputs):
+        input_shape = tf.shape(inputs)
+        batch_size, sequence_length = input_shape[0], input_shape[1]
+        i = tf.range(sequence_length)[:, tf.newaxis]
+        j = tf.range(sequence_length)
+        mask = tf.cast(i >= j, dtype="int32")
+        mask = tf.reshape(mask, (1, input_shape[1], input_shape[1]))
+        mult = tf.concat(
+            [
+                tf.expand_dims(batch_size, -1),
+                tf.constant([1, 1], dtype=tf.int32),
+            ],
+            axis=0,
+        )
+        return tf.tile(mask, mult)
+        
+
 
     def get_causal_attention_mask(self, inputs):
         input_shape = tf.shape(inputs)
@@ -178,7 +212,7 @@ class TransformerDecoderLayer(tf.keras.layers.Layer):
             axis=0
         )
         return tf.tile(mask, mult)
-
+    
 
 class ImageCaptioningModel(tf.keras.Model):
 
@@ -303,8 +337,8 @@ def generate_caption(img, caption_model, add_noise=False):
 
 
 def get_caption_model():
-    encoder = TransformerEncoderLayer(EMBEDDING_DIM, 1)
-    decoder = TransformerDecoderLayer(EMBEDDING_DIM, UNITS, 8)
+    encoder = TransformerEncoderLayer(embed_dim=EMBEDDING_DIM, dense_dim=FF_DIM, num_heads=1)
+    decoder = TransformerDecoderLayer(embed_dim=EMBEDDING_DIM, ff_dim=FF_DIM, num_heads=2)
 
     cnn_model = CNN_Encoder()
 
@@ -316,7 +350,7 @@ def get_caption_model():
         return batch
 
     caption_model.call = call_fn
-    sample_x, sample_y = tf.random.normal((1, 299, 299, 3)), tf.zeros((1, 40))
+    sample_x, sample_y = tf.random.normal((1, 299, 299, 3)), tf.zeros((1, 80))
 
     caption_model((sample_x, sample_y))
 
@@ -325,8 +359,8 @@ def get_caption_model():
     caption_model.decoder(sample_y, sample_enc_out, training=False)
 
     try:
-        caption_model.load_weights('saved_models/image_captioning_coco_weights.h5')
+        caption_model.load_weights('saved_models/my_checkpoint.ckpt')
     except FileNotFoundError:
-        caption_model.load_weights('Image-Captioning/saved_models/image_captioning_coco_weights.h5')
+        caption_model.load_weights('Image-Captioning/saved_models/my_checkpoint.ckpt')
 
-    return caption_model
+    return caption_model