algo and input changes

train.py

@@ -25,7 +25,7 @@ import numpy as np
 # ==========================================
 # 1. ARCHITECTURE CONFIGURATION
 # ==========================================
-EMBEDDING_DIM = 1024
+EMBEDDING_DIM = 512
 RNN_UNITS = 512
 
 # Pipeline & Training Hyperparameters
@@ -37,7 +37,7 @@ VOCAB_FILE = "vocab.txt"
 CHECKPOINT_DIR = f"./checkpoints_emb{EMBEDDING_DIM}_rnn{RNN_UNITS}"
 CONFIG_FILE = os.path.join(CHECKPOINT_DIR, "config.json")
 
-SEQ_LENGTH = 300
+SEQ_LENGTH = 128
 BATCH_SIZE = 128
 EPOCHS = 100
 BUFFER_SIZE = 10000
@@ -111,47 +111,27 @@ if os.path.exists(TEXT_FILE):
 # ==========================================
 # 4. MODEL ARCHITECTURE
 # ==========================================
-#class CharacterTextModel(tf.keras.Model):
-#    def __init__(self, vocab_size, embedding_dim, rnn_units):
-#        super().__init__()
-#        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
-#        self.lstm1 = tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True)
-#        self.lstm2 = tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True)
-#        self.dense = tf.keras.layers.Dense(vocab_size)
-#
-#    def call(self, inputs, states=None, return_state=False, training=False):
-#        x = self.embedding(inputs, training=training)
-#        if states is None:
-#            state_1, state_2 = None, None
-#        else:
-#            state_1, state_2 = states
-#        x, h1, c1 = self.lstm1(x, initial_state=state_1, training=training)
-#        x, h2, c2 = self.lstm2(x, initial_state=state_2, training=training)
-#        x = self.dense(x, training=training)
-#        if return_state:
-#            return x, [(h1, c1), (h2, c2)]
-#        return x
 
 class CharacterTextModel(tf.keras.Model):
     def __init__(self, vocab_size, embedding_dim, rnn_units, num_layers=2, dropout_rate=0.2):
         super().__init__()
         self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
 
-        # Create paired lists of LSTMs and Dropouts based on your desired depth
-        self.lstm_layers = []
+        # Create paired lists of GRUs and Dropouts based on your desired depth
+        self.gru_layers = []
         self.dropout_layers = []
 
-        # If num_layers=3, this loop runs 2 times, leaving the 3rd layer as the final layer
+        # Loop runs (num_layers - 1) times, leaving the final layer separate
         for _ in range(num_layers - 1):
-            self.lstm_layers.append(
-                tf.keras.layers.LSTM(rnn_units, return_sequences=True)
+            self.gru_layers.append(
+                tf.keras.layers.GRU(rnn_units, return_sequences=True)
             )
             self.dropout_layers.append(
                 tf.keras.layers.Dropout(dropout_rate)
             )
 
-        # Keep return_state ONLY on the final layer if needed for text generation
-        self.lstm_final = tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True)
+        # Switched to GRU. Note: GRU only has 1 state output (h), not 2 (h, c)
+        self.gru_final = tf.keras.layers.GRU(rnn_units, return_sequences=True, return_state=True)
         self.dropout_final = tf.keras.layers.Dropout(dropout_rate)
         self.dense = tf.keras.layers.Dense(vocab_size)
 
@@ -159,17 +139,17 @@ class CharacterTextModel(tf.keras.Model):
         x = self.embedding(inputs, training=training)
 
         # Loop through both lists simultaneously using zip()
-        for lstm, dropout in zip(self.lstm_layers, self.dropout_layers):
-            x = lstm(x, training=training)
+        for gru, dropout in zip(self.gru_layers, self.dropout_layers):
+            x = gru(x, training=training)
             x = dropout(x, training=training)
 
-        # Final LSTM Layer
-        x, h, c = self.lstm_final(x, initial_state=states, training=training)
+        # Final GRU Layer returns the output 'x' and a single hidden state 'h'
+        x, h = self.gru_final(x, initial_state=states, training=training)
         x = self.dropout_final(x, training=training)
         x = self.dense(x, training=training)
 
         if return_state:
-            return x, (h, c)
+            return x, h  # Returns just 'h' instead of '(h, c)'
         return x
 
 # Automatically scales to whatever dimensions were chosen or loaded!