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
+        # Explicitly use return_state=True for ALL LSTM layers
-        self.gru_layers = [
+        self.lstm_layers = [
-            tf.keras.layers.GRU(rnn_units, return_sequences=True, return_state=True)
+            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)):
+        for i, (lstm, dropout) in enumerate(zip(self.lstm_layers, self.dropout_layers)):
-            # Pass the specific state for this layer, and collect the new one
+            # LSTM returns: output, hidden state (h), and cell state (c)
-            x, h = gru(x, initial_state=states[i], training=training)
+            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)
 
         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")

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()