add build call

modeling4.py

@@ -1,4 +1,3 @@
-# @title Model Architecture
 import tensorflow as tf
 from tensorflow.keras import layers, activations, initializers
 
@@ -28,13 +27,19 @@ class MiniSunModel(tf.keras.Model):
         self.token_embedding = layers.Embedding(config.vocab_size, config.hidden_size)
         self.position_embedding = layers.Embedding(config.max_position_embeddings, config.hidden_size)
 
-        # Transformer decoder blocks
-        self.decoder_blocks = [self._build_decoder_block() for _ in range(config.num_hidden_layers)]
+        # Initialize an empty list for decoder blocks
+        self.decoder_blocks = []
 
         # Final normalization and head
         self.layer_norm = layers.LayerNormalization(epsilon=1e-6)
         self.lm_head = layers.Dense(config.vocab_size, kernel_initializer=initializers.he_normal())
 
+    def build(self, input_shape):
+        # Create transformer decoder blocks based on the model configuration
+        self.decoder_blocks = [self._build_decoder_block() for _ in range(self.config.num_hidden_layers)]
+        # Call the superclass's build method
+        super(MiniSunModel, self).build(input_shape)
+
     def _build_decoder_block(self):
         # Decoder block consisting of multi-head attention and feed-forward layers
         return [
@@ -91,7 +96,6 @@ class MiniSunModel(tf.keras.Model):
     
     def compute_loss(self, labels, logits):
         """Computes the loss between labels and logits."""
-        # Ensure labels and logits are not None
         if labels is None or logits is None:
             raise ValueError("Labels and logits cannot be None.")
         return tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)
@@ -117,7 +121,6 @@ class MiniSunModel(tf.keras.Model):
 
         return {m.name: m.result() for m in self.metrics}
 
-
 def create_model(config):
     model = MiniSunModel(config)
 
@@ -137,7 +140,6 @@ def cosine_annealing_with_warmup(step, config):
     if step < warmup_steps:
         return config.learning_rate * (step / warmup_steps)
     else:
-        # Calculate the cosine decay
         cos_step = step - warmup_steps
         total_cos_steps = config.total_steps - warmup_steps
         return 0.5 * config.learning_rate * (1 + tf.cos(tf.constant(np.pi) * cos_step / total_cos_steps))
@@ -146,21 +148,16 @@ def cosine_annealing_with_restarts(step, config, restart_period, cycle_num):
     """Learning rate schedule with warm-up and cosine annealing with restarts."""
     warmup_steps = int(config.total_steps * config.warmup_steps)
     
-    # Determine the current cycle based on step and restart_period
     current_cycle = step // restart_period
-
-    # Calculate the effective step within the current cycle
     effective_step = step % restart_period
 
     if effective_step < warmup_steps:
         return config.learning_rate * (effective_step / warmup_steps)
     else:
-        # Calculate the cosine decay within the current cycle
         cos_step = effective_step - warmup_steps
         total_cos_steps = restart_period - warmup_steps
         return 0.5 * config.learning_rate * (1 + tf.cos(tf.constant(np.pi) * cos_step / total_cos_steps))
 
-
 # Configuration
 config = MiniSunConfig()
 
@@ -169,4 +166,4 @@ model = create_model(config)
 
 # Create a LearningRateScheduler callback
 lr_scheduler = tf.keras.callbacks.LearningRateScheduler(lambda step: cosine_annealing_with_warmup(step, config))
-#lr_scheduler_with_restarts = tf.keras.callbacks.LearningRateScheduler(lambda step: cosine_annealing_with_restarts(step, config, restart_period=1000, cycle_num=1))
+# lr_scheduler_with_restarts = tf.keras.callbacks.LearningRateScheduler(lambda step: cosine_annealing_with_restarts(step, config, restart_period=1000, cycle_num=1))