New train version

train.py

@@ -38,9 +38,9 @@ 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 = 64
+SEQ_LENGTH = 512
 BATCH_SIZE = 128
-EPOCHS = 21
+EPOCHS = 210
 BUFFER_SIZE = 10000
 SEED_TEXT = "In the case of unlawful enterance"
 
@@ -114,13 +114,13 @@ if os.path.exists(TEXT_FILE):
 # ==========================================
 
 class CharacterTextModel(tf.keras.Model):
-    def __init__(self, vocab_size, embedding_dim, rnn_units, num_layers=3, dropout_rate=0.4):
+    def __init__(self, vocab_size, embedding_dim, rnn_units, num_layers=1, dropout_rate=0.3):
         super().__init__()
         self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
 
-        # Explicitly use return_state=True for ALL GRU layers
-        self.gru_layers = [
-            tf.keras.layers.GRU(rnn_units, return_sequences=True, return_state=True)
+        # Explicitly use return_state=True for ALL LSTM layers
+        self.lstm_layers = [
+            tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True)
             for _ in range(num_layers)
         ]
         self.dropout_layers = [tf.keras.layers.Dropout(dropout_rate) for _ in range(num_layers)]
@@ -131,20 +131,22 @@ class CharacterTextModel(tf.keras.Model):
 
         # If no states are provided, initialize a list of None
         if states is None:
-            states = [None] * len(self.gru_layers)
+            states = [None] * len(self.lstm_layers)
 
         new_states = []
-        for i, (gru, dropout) in enumerate(zip(self.gru_layers, self.dropout_layers)):
-            # Pass the specific state for this layer, and collect the new one
-            x, h = gru(x, initial_state=states[i], training=training)
+        for i, (lstm, dropout) in enumerate(zip(self.lstm_layers, self.dropout_layers)):
+            # LSTM returns: output, hidden state (h), and cell state (c)
+            x, h, c = lstm(x, initial_state=states[i], training=training)
             x = dropout(x, training=training)
-            new_states.append(h)
+            # Store both states as a list/tuple for this layer
+            new_states.append([h, c])
 
-        x = self.dense(x, training=training)
+            x = self.dense(x, training=training)
 
         if return_state:
             return x, new_states
         return x
+
 # Automatically scales to whatever dimensions were chosen or loaded!
 model = CharacterTextModel(vocab_size=vocab_size, embedding_dim=EMBEDDING_DIM, rnn_units=RNN_UNITS)
 model.build(input_shape=(BATCH_SIZE, SEQ_LENGTH))
@@ -156,7 +158,7 @@ checkpoint_files = glob.glob(os.path.join(CHECKPOINT_DIR, "ckpt_*.weights.h5"))
 
 if checkpoint_files:
     # Sort files naturally or by modification time to get the latest one
-    latest_checkpoint = max(checkpoint_files, key=os.path.getmtime)
+    latest_checkpoint = min(checkpoint_files, key=os.path.getmtime)
     print(f"Restoring model layers from: {latest_checkpoint}")
     model.load_weights(latest_checkpoint)
 else:
@@ -203,7 +205,7 @@ class GenerationCallback(tf.keras.callbacks.Callback):
     def on_epoch_end(self, epoch, logs=None):
         print(f"\n--- End of Epoch {epoch+1} ---")
         if epoch % 10 == 0:
-            for temp in [0.1, 0.4, 0.7]:
+            for temp in [0.3, 0.6, 0.9, 1.2]:
                 print(f"Seed: \"{SEED_TEXT}\" -> {produce_sample(self.model, seed=SEED_TEXT, num_generate=200, temperature=temp)}\n")
 
 checkpoint_prefix = os.path.join(CHECKPOINT_DIR, "ckpt_{epoch}.weights.h5")