Updated training data and model definitions

train.py

@@ -25,11 +25,12 @@ import numpy as np
 # ==========================================
 # 1. ARCHITECTURE CONFIGURATION
 # ==========================================
-EMBEDDING_DIM = 64
+EMBEDDING_DIM = 512
-RNN_UNITS = 256
+RNN_UNITS = 512
 
 # Pipeline & Training Hyperparameters
-TEXT_FILE = "combined_training_data.txt"
+#TEXT_FILE = "combined_training_data.txt"
+TEXT_FILE = "emily_post.txt"
 VOCAB_FILE = "vocab.txt"
 
 # DYNAMICALLY MANGLED DIRECTORY:
@@ -37,11 +38,11 @@ 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 = 128
+SEQ_LENGTH = 64
 BATCH_SIZE = 128
-EPOCHS = 100
+EPOCHS = 21
 BUFFER_SIZE = 10000
-SEED_TEXT = "it was a dark and stormy night"
+SEED_TEXT = "In the case of unlawful enterance"
 
 print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
 os.makedirs(CHECKPOINT_DIR, exist_ok=True)
@@ -113,45 +114,37 @@ if os.path.exists(TEXT_FILE):
 # ==========================================
 
 class CharacterTextModel(tf.keras.Model):
-    def __init__(self, vocab_size, embedding_dim, rnn_units, num_layers=2, dropout_rate=0.4):
+    def __init__(self, vocab_size, embedding_dim, rnn_units, num_layers=3, dropout_rate=0.4):
         super().__init__()
         self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
 
-        # Create paired lists of GRUs and Dropouts based on your desired depth
+        # Explicitly use return_state=True for ALL GRU layers
-        self.gru_layers = []
+        self.gru_layers = [
-        self.dropout_layers = []
+            tf.keras.layers.GRU(rnn_units, return_sequences=True, return_state=True)
-
+            for _ in range(num_layers)
-        # Loop runs (num_layers - 1) times, leaving the final layer separate
+        ]
-        for _ in range(num_layers - 1):
+        self.dropout_layers = [tf.keras.layers.Dropout(dropout_rate) for _ in range(num_layers)]
-            self.gru_layers.append(
-                tf.keras.layers.GRU(rnn_units, return_sequences=True)
-            )
-            self.dropout_layers.append(
-                tf.keras.layers.Dropout(dropout_rate)
-            )
-
-        # 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)
 
     def call(self, inputs, states=None, return_state=False, training=False):
         x = self.embedding(inputs, training=training)
 
-        # Loop through both lists simultaneously using zip()
+        # If no states are provided, initialize a list of None
-        for gru, dropout in zip(self.gru_layers, self.dropout_layers):
+        if states is None:
-            x = gru(x, training=training)
+            states = [None] * len(self.gru_layers)
-            x = dropout(x, training=training)
+
+        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)
+            x = dropout(x, training=training)
+            new_states.append(h)
 
-        # 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  # Returns just 'h' instead of '(h, c)'
+            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))
@@ -177,17 +170,32 @@ model.compile(optimizer='adam', loss=loss)
 # 5. SAMPLING & GENERATION LOGIC
 # ==========================================
 def produce_sample(model, seed, num_generate=300, temperature=0.7):
+    # Vectorize the initial seed text
     input_chars = vectorize_layer(tf.constant([seed]))
     input_ids = input_chars[0][:len(seed)].numpy().tolist()
-    generated_ids, states = [], None
 
-    for _ in range(num_generate):
+    generated_ids = []
-        current_tokens = input_ids[-SEQ_LENGTH:]
+    states = None
-        predictions, states = model(tf.expand_dims(current_tokens, 0), states=states, return_state=True, training=False)
+
+    # Step 1: "Warm up" the model with the seed text to build the initial states
+    # We pass the entire seed here
+    predictions, states = model(tf.expand_dims(input_ids, 0), states=states, return_state=True, training=False)
+
+    # Get the very last prediction from the seed sequence
+    predictions = predictions[0, -1, :] / temperature
+    predicted_id = tf.random.categorical(tf.expand_dims(predictions, 0), num_samples=1)[0, 0].numpy()
+    generated_ids.append(predicted_id)
+
+    # Step 2: Generation loop using ONLY the single newest token and updating states
+    for _ in range(num_generate - 1):
+        # Pass ONLY the last predicted token, plus the existing states
+        predictions, states = model(tf.expand_dims([predicted_id], 0), states=states, return_state=True, training=False)
+
+        # Scale by temperature and sample
         predictions = predictions[0, -1, :] / temperature
         predicted_id = tf.random.categorical(tf.expand_dims(predictions, 0), num_samples=1)[0, 0].numpy()
+
         generated_ids.append(predicted_id)
-        input_ids.append(predicted_id)
 
     return "".join([vocab[idx] for idx in generated_ids])
 
@@ -203,16 +211,25 @@ checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_pre
 
 # Stop training when validation loss stops improving for 3 epochs
 early_stopping = tf.keras.callbacks.EarlyStopping(
-    monitor='val_loss',
+    monitor='loss',
     patience=3,
-    restore_best_weights=True # Automatically rolls back to epoch 15 weights!
+    restore_best_weights=True # Automatically rolls back best weights
+)
+
+# Create the learning rate scheduler callback
+lr_scheduler = tf.keras.callbacks.ReduceLROnPlateau(
+    monitor='loss',        # Can change to 'val_loss' if using a validation split
+    factor=0.5,            # Multiply the learning rate by 0.5 when triggered (cuts it in half)
+    patience=2,            # Number of epochs with no improvement before dropping LR
+    min_lr=1e-6,           # Don't let the learning rate drop lower than this
+    verbose=1              # Prints a message when the learning rate changes
 )
 
 # ==========================================
 # 6. START RUN
 # ==========================================
 if dataset is not None:
-    model.fit(dataset, epochs=EPOCHS, callbacks=[GenerationCallback(), checkpoint_callback])
+    model.fit(dataset, epochs=EPOCHS, callbacks=[GenerationCallback(), checkpoint_callback, early_stopping, lr_scheduler])
 else:
     print(f"\n--- Inference Mode ({EMBEDDING_DIM}dim, {RNN_UNITS}units) ---")
     print(f"Result: {produce_sample(model, seed=SEED_TEXT, num_generate=300, temperature=0.6)}")