From e4300773e15f3e8a4d39fd0679b5bc3992309ee1 Mon Sep 17 00:00:00 2001 From: Nathan Hinton Date: Wed, 1 Jul 2026 21:52:04 -0600 Subject: [PATCH] rebasing to main --- train.py | 110 +++++++++++++++++-------------------------------------- 1 file changed, 33 insertions(+), 77 deletions(-) diff --git a/train.py b/train.py index 806b19f..7d48e22 100644 --- a/train.py +++ b/train.py @@ -19,12 +19,13 @@ def configure_xla_paths(): configure_xla_paths() import tensorflow as tf +#tf.debugging.set_log_device_placement(True) import numpy as np # ========================================== # 1. ARCHITECTURE CONFIGURATION # ========================================== -EMBEDDING_DIM = 256 +EMBEDDING_DIM = 1024 RNN_UNITS = 512 # Pipeline & Training Hyperparameters @@ -36,8 +37,8 @@ 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 = 150 -BATCH_SIZE = 64 +SEQ_LENGTH = 300 +BATCH_SIZE = 128 EPOCHS = 100 BUFFER_SIZE = 10000 SEED_TEXT = "it was a dark and stormy night" @@ -132,108 +133,63 @@ if os.path.exists(TEXT_FILE): # return x class CharacterTextModel(tf.keras.Model): - def __init__(self, vocab_size, embedding_dim, rnn_units, dropout_rate=0.2): + 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) - # Layer 1 - self.lstm1 = tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True) - self.ln1 = tf.keras.layers.LayerNormalization() - self.dropout1 = tf.keras.layers.Dropout(dropout_rate) + # Create paired lists of LSTMs and Dropouts based on your desired depth + self.lstm_layers = [] + self.dropout_layers = [] - # Layer 2 - self.lstm2 = tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True) - self.ln2 = tf.keras.layers.LayerNormalization() - self.dropout2 = tf.keras.layers.Dropout(dropout_rate) + # If num_layers=3, this loop runs 2 times, leaving the 3rd layer as the final layer + for _ in range(num_layers - 1): + self.lstm_layers.append( + tf.keras.layers.LSTM(rnn_units, return_sequences=True) + ) + self.dropout_layers.append( + tf.keras.layers.Dropout(dropout_rate) + ) - # Layer 3 - self.lstm3 = tf.keras.layers.LSTM(rnn_units, return_sequences=True, return_state=True) - self.ln3 = tf.keras.layers.LayerNormalization() - self.dropout3 = tf.keras.layers.Dropout(dropout_rate) - - # 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) + # 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) self.dropout_final = tf.keras.layers.Dropout(dropout_rate) - - # Final output dense layer 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) - # Unpack states cleanly for 3 layers - if states is 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, state_5, state_6, state_final = states + # Loop through both lists simultaneously using zip() + for lstm, dropout in zip(self.lstm_layers, self.dropout_layers): + x = lstm(x, training=training) + x = dropout(x, training=training) - # Pass through Layer 1 - x, h1, c1 = self.lstm1(x, initial_state=state_1, training=training) - x = self.ln1(x, training=training) - x = self.dropout1(x, training=training) - - # Pass through Layer 2 - x, h2, c2 = self.lstm2(x, initial_state=state_2, training=training) - x = self.ln2(x, training=training) - x = self.dropout2(x, training=training) - - # Pass through Layer 3 - x, h3, c3 = self.lstm3(x, initial_state=state_3, training=training) - x = self.ln3(x, training=training) - x = self.dropout3(x, training=training) - - # 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) + # Final LSTM Layer + x, h, c = self.lstm_final(x, initial_state=states, 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), (h5, c5), (h6, c6), (h_final, c_final)] + return x, (h, c) 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)) -latest_checkpoint = tf.train.latest_checkpoint(CHECKPOINT_DIR) -if latest_checkpoint: +import glob + +# Find all files matching the pattern +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) print(f"Restoring model layers from: {latest_checkpoint}") model.load_weights(latest_checkpoint) else: print("Starting a clean initialization.") + loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) model.compile(optimizer='adam', loss=loss)