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