Update src/streamlit_app.py

src/streamlit_app.py

@@ -128,28 +128,8 @@ class TransformerDecoder(layers.Layer):
         self.layernorm_2 = layers.LayerNormalization(epsilon=1e-5)
         self.layernorm_3 = layers.LayerNormalization(epsilon=1e-5)
 
-    # def get_causal_attention_mask(self, inputs):
-    #     seq_len = tf.shape(inputs)[1]
-    #     causal_mask = tf.linalg.band_part(tf.ones((seq_len, seq_len), dtype=tf.bool), -1, 0)
-    #     return causal_mask[tf.newaxis, :, :]  # (1, seq_len, seq_len)
-
     def call(self, inputs, encoder_outputs, mask=None):
-        # Padding mask: (batch_size, 1, seq_len)
-        # if mask is not None:
-        #     padding_mask = tf.cast(mask[:, tf.newaxis, :], dtype=tf.bool)
-        # else:
-        #     padding_mask = None
-
-        # # Causal mask: (1, seq_len, seq_len)
-        # causal_mask = self.get_causal_attention_mask(inputs)
-
-        # # Combine masks for self-attention
-        # if padding_mask is not None:
-        #     combined_mask = tf.logical_and(padding_mask, causal_mask)
-        # else:
-        #     combined_mask = causal_mask
-
-        # Self-attention with combined mask
+        
         attention_output_1 = self.attention_1(
             query=inputs,
             value=inputs,
@@ -197,33 +177,38 @@ embed_dim = 512
 dense_dim = 2048
 num_heads = 8
 num_blocks = 7
-    
+
 encoder_inputs = tf.keras.Input(shape=(None,), dtype="int32", name="encoder_inputs")
 decoder_inputs = tf.keras.Input(shape=(None,), dtype="int32", name="decoder_inputs")
-    
-# Paddincfg masks
+
+# Masks
 encoder_mask = tf.keras.layers.Lambda(lambda x: tf.cast(tf.not_equal(x, 0), tf.bool))(encoder_inputs)
-cross_attention_mask = tf.keras.layers.Lambda(lambda x: tf.cast(x[:, tf.newaxis, tf.newaxis, :], tf.bool))(encoder_mask) 
-    
-# Embeddings
+cross_attention_mask = tf.keras.layers.Lambda(lambda x: tf.cast(x[:, tf.newaxis, tf.newaxis, :], tf.bool))(encoder_mask)
+
+    # Embeddings
 encoder_embed = PositionalEmbedding(sequence_length, vocab_size, embed_dim)(encoder_inputs)
-decoder_embed = PositionalEmbedding(256, 257, embed_dim,mask_zero=False)(decoder_inputs)
-    
-# Encoder blocks
+decoder_embed = PositionalEmbedding(256, 257, embed_dim, mask_zero=False)(decoder_inputs)
+
+    # Pre-instantiate blocks
+encoder_blocks = [TransformerEncoder(embed_dim, dense_dim, num_heads) for _ in range(num_blocks)]
+decoder_blocks = [TransformerDecoder(embed_dim, dense_dim, num_heads) for _ in range(num_blocks)]
+
+    # Encoder
 x = encoder_embed
-for _ in range(num_blocks):
-    x = TransformerEncoder(embed_dim, dense_dim, num_heads)(x, mask=encoder_mask)
+for block in encoder_blocks:
+    x = block(x, mask=encoder_mask)
 encoder_outputs = x
-    
-# Decoder blocks
+
+    # Decoder
 x = decoder_embed
-for _ in range(num_blocks):
-    x = TransformerDecoder(embed_dim, dense_dim, num_heads)(x, encoder_outputs, mask=cross_attention_mask)
+for block in decoder_blocks:
+    x = block(x, encoder_outputs, mask=cross_attention_mask)
 
-# Final layers
+    # Output layers
 x = layers.LayerNormalization(epsilon=1e-5)(x)
 x = layers.Dropout(0.1)(x)
 decoder_outputs = layers.Dense(256)(x)
+
 transformer = keras.Model([encoder_inputs, decoder_inputs], decoder_outputs)
 
 start_token = 256