bionickatana/emily_post_llm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

New train version

Browse Source
This commit is contained in:
Nathan Hinton
2026-07-06 17:41:08 -06:00
parent 4f90c64db8
commit eeb4315aa5
2 changed files with 298 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+16 -14
View File
@@ -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")
train_v2.py
+282
View File
@@ -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()