From 10591c24ef4cfb083d1017834f3512d8d1b2a70c Mon Sep 17 00:00:00 2001
From: Nathan Hinton <nathan@hintonclan.org>
Date: Wed, 1 Jul 2026 19:21:52 -0600
Subject: [PATCH] bigger model better?

---
 train.py | 38 ++++++++++++++++++++++++++++++++++----
 1 file changed, 34 insertions(+), 4 deletions(-)

diff --git a/train.py b/train.py
index ffdba98..adf69f6 100644
--- a/train.py
+++ b/train.py
@@ -159,10 +159,25 @@ class CharacterTextModel(tf.keras.Model):
         self.ln3 = tf.keras.layers.LayerNormalization()
         self.dropout3 = tf.keras.layers.Dropout(dropout_rate)
 
-        # Layer 4 (Added for deeper text/context comprehension)
+        # Layer 4
         self.lstm4 = tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True)
+        self.ln4 = tf.keras.layers.LayerNormalization()
         self.dropout4 = tf.keras.layers.Dropout(dropout_rate)
 
+        # Layer 5
+        self.lstm5 = tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True)
+        self.ln5 = tf.keras.layers.LayerNormalization()
+        self.dropout5 = tf.keras.layers.Dropout(dropout_rate)
+
+        # Layer 6
+        self.lstm6 = tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True)
+        self.ln6 = tf.keras.layers.LayerNormalization()
+        self.dropout6 = tf.keras.layers.Dropout(dropout_rate)
+
+        # Layer _final (Added for deeper text/context comprehension)
+        self.lstm_final = tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True)
+        self.dropout_final = tf.keras.layers.Dropout(dropout_rate)
+
         # Final output dense layer
         self.dense = tf.keras.layers.Dense(vocab_size)
 
@@ -171,9 +186,9 @@ class CharacterTextModel(tf.keras.Model):
 
         # Unpack states cleanly for 3 layers
         if states is None:
-            state_1, state_2, state_3, state_4 = None, None, None, None
+            state_1, state_2, state_3, state_4, state_5, state_6, state_final = None, None, None, None, None, None, None
         else:
-            state_1, state_2, state_3, state_4 = states
+            state_1, state_2, state_3, state_4, state_5, state_6, state_final = states
 
         # Pass through Layer 1
         x, h1, c1 = self.lstm1(x, initial_state=state_1, training=training)
@@ -192,13 +207,28 @@ class CharacterTextModel(tf.keras.Model):
 
         # Pass through Layer 4
         x, h4, c4 = self.lstm4(x, initial_state=state_4, training=training)
+        x = self.ln4(x, training=training)
         x = self.dropout4(x, training=training)
 
+        # Pass through Layer 5
+        x, h5, c5 = self.lstm5(x, initial_state=state_5, training=training)
+        x = self.ln5(x, training=training)
+        x = self.dropout5(x, training=training)
+
+        # Pass through Layer 6
+        x, h6, c6 = self.lstm6(x, initial_state=state_6, training=training)
+        x = self.ln6(x, training=training)
+        x = self.dropout6(x, training=training)
+
+        # Pass through Layer _final
+        x, h_final, c_final = self.lstm_final(x, initial_state=state_final, training=training)
+        x = self.dropout_final(x, training=training)
+
         # Output logits
         x = self.dense(x, training=training)
 
         if return_state:
-            return x, [(h1, c1), (h2, c2), (h3, c3), (h4, c4)]
+            return x, [(h1, c1), (h2, c2), (h3, c3), (h4, c4), (h5, c5), (h6, c6), (h_final, c_final)]
         return x
 
 # Automatically scales to whatever dimensions were chosen or loaded!