From eeb4315aa5f89e1ba89f1cb6d1f7efe55f370c0e Mon Sep 17 00:00:00 2001 From: Nathan Hinton Date: Mon, 6 Jul 2026 17:41:08 -0600 Subject: [PATCH] New train version --- train.py | 30 +++--- train_v2.py | 282 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 298 insertions(+), 14 deletions(-) create mode 100644 train_v2.py diff --git a/train.py b/train.py index a5aceca..3fcba11 100644 --- a/train.py +++ b/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") diff --git a/train_v2.py b/train_v2.py new file mode 100644 index 0000000..c15adb2 --- /dev/null +++ b/train_v2.py @@ -0,0 +1,282 @@ +import os +import json +import glob +import argparse +import sys +import re + +# ============================================================================== +# AUTO-CONFIG: Find and link hidden XLA compiler libraries inside the venv +# ============================================================================== +def configure_xla_paths(): + venv_base = sys.prefix + for root, _, files in os.walk(venv_base): + if "libdevice.10.bc" in files: + cuda_dir = root.split("/nvvm")[0] + os.environ["XLA_FLAGS"] = f"--xla_gpu_cuda_data_dir={cuda_dir}" + print(f"[XLA Config] Successfully linked compiler data to: {cuda_dir}") + return True + print("[XLA Config] Warning: libdevice.10.bc not found. Training may crash.") + return False + +configure_xla_paths() + +import tensorflow as tf +import numpy as np + +# ========================================== +# 1. MODEL ARCHITECTURE +# ========================================== +class CharacterTextModel(tf.keras.Model): + 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) + + 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)] + 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: + states = [None] * len(self.lstm_layers) + + new_states = [] + for i, (lstm, dropout) in enumerate(zip(self.lstm_layers, self.dropout_layers)): + x, h, c = lstm(x, initial_state=states[i], training=training) + x = dropout(x, training=training) + new_states.append([h, c]) + + # FIXED: Moved out of the loop so multi-layer architectures function properly + x = self.dense(x, training=training) + + if return_state: + return x, new_states + return x + + + +# ========================================== +# 2. HELPER UTILITIES +# ========================================== +def load_or_create_vocab(text_file, vocab_file): + if os.path.exists(vocab_file): + print(f"[Vocab] Loading existing vocabulary from '{vocab_file}'...") + with open(vocab_file, 'r', encoding='utf-8') as f: + return json.load(f) + + if not os.path.exists(text_file): + print(f"[Error] Text file '{text_file}' not found. Cannot build vocabulary.") + quit(1) + + print(f"[Vocab] Building vocabulary from {text_file}...") + with open(text_file, 'r', encoding='utf-8') as f: + text = f.read() + temp_vectorizer = tf.keras.layers.TextVectorization(split="character", standardize="lower") + temp_vectorizer.adapt(tf.data.Dataset.from_tensor_slices([text])) + vocab = temp_vectorizer.get_vocabulary() + with open(vocab_file, 'w', encoding='utf-8') as f: + json.dump(vocab, f, ensure_ascii=False) + return vocab + +def produce_sample(model, vectorize_layer, vocab, seed, num_generate=300, temperature=0.7): + input_chars = vectorize_layer(tf.constant([seed])) + input_ids = input_chars[0][:len(seed)].numpy().tolist() + + generated_ids = [] + states = None + + predictions, states = model(tf.expand_dims(input_ids, 0), states=states, return_state=True, training=False) + 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) + + for _ in range(num_generate - 1): + predictions, states = model(tf.expand_dims([predicted_id], 0), states=states, return_state=True, training=False) + 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) + + return "".join([vocab[idx] for idx in generated_ids]) + +class GenerationCallback(tf.keras.callbacks.Callback): + def __init__(self, vectorize_layer, vocab, seed_text): + super().__init__() + self.vectorize_layer = vectorize_layer + self.vocab = vocab + self.seed_text = seed_text + + def on_epoch_end(self, epoch, logs=None): + print(f"\n--- End of Epoch {epoch+1} ---") + if epoch % 10 == 0: + for temp in [0.3, 0.6, 0.9, 1.2]: + sample = produce_sample(self.model, self.vectorize_layer, self.vocab, seed=self.seed_text, num_generate=350, temperature=temp) + print(f"Seed (Temp {temp}): \"{self.seed_text}\" -> {sample}\n") + + + + + + + +# ========================================== +# 3. MAIN EXECUTION PIPELINE +# ========================================== +def main(): + parser = argparse.ArgumentParser(description="Dynamically build, train, or infer from custom character LSTM models.") + parser.add_argument("--modelfile", type=str, default=None, help="Path to a specific '.weights.h5' checkpoint file.") + parser.add_argument("--prompt", type=str, default=None, help="Prompt string to run inference generation.") + parser.add_argument("--train", action="store_true", help="Flag to trigger or resume training on a model.") + args = parser.parse_args() + + # Default settings for clean, scratch runs + DEFAULT_EMBEDDING_DIM = 512 + DEFAULT_RNN_UNITS = 512 + DEFAULT_NUM_LAYERS = 1 + DEFAULT_DROPOUT_RATE = 0.3 + + TEXT_FILE = "emily_post.txt" + VOCAB_FILE = "vocab.txt" + SEQ_LENGTH = 512 + BATCH_SIZE = 128 + EPOCHS_PER_RUN = 200 # Easily run in chunks of 20 epochs + BUFFER_SIZE = 10000 + SEED_TEXT = "In the case of unlawful enterance" + + vocab = load_or_create_vocab(TEXT_FILE, VOCAB_FILE) + vocab_size = len(vocab) + + vectorize_layer = tf.keras.layers.TextVectorization( + split="character", standardize="lower", vocabulary=vocab, output_mode="int" + ) + + # Architectural parameters (will overwrite if loading an old model profile) + emb_dim, rnn_units, num_layers, dropout_rate = DEFAULT_EMBEDDING_DIM, DEFAULT_RNN_UNITS, DEFAULT_NUM_LAYERS, DEFAULT_DROPOUT_RATE + checkpoint_dir = f"./checkpoints_emb{emb_dim}_rnn{rnn_units}_layers{num_layers}" + + # Determine execution behavior based on user arguments + is_resuming = args.modelfile is not None + initial_epoch = 0 + + if is_resuming: + # Check for a matching configuration JSON file alongside the weights file + base_path, _ = os.path.splitext(args.modelfile) + # We strip off the ".weights" part if present to look for the matching json + if base_path.endswith('.weights'): + base_path = base_path[:-8] + + config_path = f"{base_path}_config.json" + + if os.path.exists(config_path): + print(f"[Config] Found matching footprint configuration at: {config_path}") + with open(config_path, 'r') as f: + saved_config = json.load(f) + emb_dim = saved_config["embedding_dim"] + rnn_units = saved_config["rnn_units"] + num_layers = saved_config.get("num_layers", DEFAULT_NUM_LAYERS) + dropout_rate = saved_config.get("dropout_rate", DEFAULT_DROPOUT_RATE) + saved_lr = saved_config.get("learning_rate", None) + checkpoint_dir = saved_config.get("checkpoint_dir", os.path.dirname(args.modelfile)) + print(f"-> Architecture Adjusted to match saved model: Embedding={emb_dim}, RNN={rnn_units}, Layers={num_layers}") + else: + print(f"[Warning] No unique configuration file found at '{config_path}'. Falling back to script defaults.") + filename = os.path.basename(args.modelfile) + epoch_match = re.search(r'ckpt_(\d+)', filename) + if epoch_match: + initial_epoch = int(epoch_match.group(1)) + print(f"[Tracking] Detected training resume point. Resuming starting at Epoch: {initial_epoch}") + + else: + print("[Mode] No model file passed. Initializing a brand-new training run profile.") + os.makedirs(checkpoint_dir, exist_ok=True) + + # Initialize model dynamically using the determined metrics + model = CharacterTextModel(vocab_size=vocab_size, embedding_dim=emb_dim, rnn_units=rnn_units, num_layers=num_layers, dropout_rate=dropout_rate) + model.build(input_shape=(BATCH_SIZE, SEQ_LENGTH)) + + loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) + model.compile(optimizer='adam', loss=loss_fn) + + if is_resuming: + print(f"[Weights] Restoring model layer states from: {args.modelfile}") + model.load_weights(args.modelfile) + if saved_lr is not None: + model.optimizer.learning_rate.assign(saved_lr) + print(f"-> Optimizer Learning Rate restored to: {saved_lr}") + + # Process data pipeline if training is requested OR if no parameters were specified at all + should_train = args.train or (args.modelfile is None and args.prompt is None) + + if should_train: + if not os.path.exists(TEXT_FILE): + print(f"[Error] Raw text source '{TEXT_FILE}' required for training. Exiting.") + quit(1) + + print(f"[Data] Preparing dataset batches from {TEXT_FILE}...") + with open(TEXT_FILE, 'r', encoding='utf-8') as f: + text = f.read() + all_ids = vectorize_layer(tf.constant([text]))[0] + ids_dataset = tf.data.Dataset.from_tensor_slices(all_ids) + sequences = ids_dataset.batch(SEQ_LENGTH + 1, drop_remainder=True) + + def split_input_target(sequence): + return sequence[:-1], sequence[1:] + + dataset = sequences.map(split_input_target).shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True).prefetch(tf.data.AUTOTUNE) + + # Callbacks + # Note: Weights and JSON foot-prints share matching index sequences + checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}") + + class SaveConfigCallback(tf.keras.callbacks.Callback): + def on_epoch_end(self, epoch, logs=None): + # Save out metadata alongside weights files every epoch + epoch_config_path = f"{checkpoint_prefix.format(epoch=epoch+1)}_config.json" + current_footprint = { + "embedding_dim": emb_dim, + "rnn_units": rnn_units, + "num_layers": num_layers, + "dropout_rate": dropout_rate, + "checkpoint_dir": checkpoint_dir, + "learning_rate": float(self.model.optimizer.learning_rate.numpy()) + } + with open(epoch_config_path, 'w') as f: + json.dump(current_footprint, f, indent=4) + + checkpoint_callback = tf.keras.callbacks.ModelCheckpoint( + filepath=checkpoint_prefix + ".weights.h5", + save_weights_only=True + ) + + early_stopping = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=4, restore_best_weights=True) + lr_scheduler = tf.keras.callbacks.ReduceLROnPlateau(monitor='loss', factor=0.5, patience=2, min_lr=1e-6, verbose=1) + gen_callback = GenerationCallback(vectorize_layer, vocab, SEED_TEXT) + + total_epochs = initial_epoch + EPOCHS_PER_RUN + print(f"[Train] Commencing training loop inside checkpoint tree: {checkpoint_dir}") + print(f"[Train] Running Epochs {initial_epoch + 1} through {total_epochs}...") + + model.fit( + dataset, + epochs=total_epochs, + initial_epoch = initial_epoch, + callbacks=[gen_callback, checkpoint_callback, SaveConfigCallback(), early_stopping, lr_scheduler] + ) + + # Process custom Inference Generations + if args.prompt is not None: + print(f"\n--- Custom Inference Mode (Seed: \"{args.prompt}\") ---") + result = produce_sample(model, vectorize_layer, vocab, seed=args.prompt, num_generate=300, temperature=0.6) + print(f"Result:\n{result}\n") + elif not should_train and args.modelfile is not None: + # User specified a file but gave no training or prompt orders; run a quick test generation + print(f"\n--- Default Quick Test (Seed: \"{SEED_TEXT}\") ---") + result = produce_sample(model, vectorize_layer, vocab, seed=SEED_TEXT, num_generate=300, temperature=0.6) + print(f"Result:\n{result}\n") + +if __name__ == "__main__": + main()