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

Updated training data and model definitions

Browse Source
This commit is contained in:
Nathan Hinton
2026-07-06 09:46:11 -06:00
parent 5009183c67
commit 4f90c64db8
2 changed files with 24567 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+57 -40
View File
@@ -25,11 +25,12 @@ import numpy as np
# ==========================================
# 1. ARCHITECTURE CONFIGURATION
# ==========================================
EMBEDDING_DIM = 64
RNN_UNITS = 256
EMBEDDING_DIM = 512
RNN_UNITS = 512
# Pipeline & Training Hyperparameters
TEXT_FILE = "combined_training_data.txt"
#TEXT_FILE = "combined_training_data.txt"
TEXT_FILE = "emily_post.txt"
VOCAB_FILE = "vocab.txt"
# DYNAMICALLY MANGLED DIRECTORY:
@@ -37,11 +38,11 @@ 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 = 128
SEQ_LENGTH = 64
BATCH_SIZE = 128
EPOCHS = 100
EPOCHS = 21
BUFFER_SIZE = 10000
SEED_TEXT = "it was a dark and stormy night"
SEED_TEXT = "In the case of unlawful enterance"
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
os.makedirs(CHECKPOINT_DIR, exist_ok=True)
@@ -113,45 +114,37 @@ if os.path.exists(TEXT_FILE):
# ==========================================
class CharacterTextModel(tf.keras.Model):
def __init__(self, vocab_size, embedding_dim, rnn_units, num_layers=2, dropout_rate=0.4):
def __init__(self, vocab_size, embedding_dim, rnn_units, num_layers=3, dropout_rate=0.4):
super().__init__()
self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
# Create paired lists of GRUs and Dropouts based on your desired depth
self.gru_layers = []
self.dropout_layers = []
# Loop runs (num_layers - 1) times, leaving the final layer separate
for _ in range(num_layers - 1):
self.gru_layers.append(
tf.keras.layers.GRU(rnn_units, return_sequences=True)
)
self.dropout_layers.append(
tf.keras.layers.Dropout(dropout_rate)
)
# Switched to GRU. Note: GRU only has 1 state output (h), not 2 (h, c)
self.gru_final = tf.keras.layers.GRU(rnn_units, return_sequences=True, return_state=True)
self.dropout_final = tf.keras.layers.Dropout(dropout_rate)
# Explicitly use return_state=True for ALL GRU layers
self.gru_layers = [
tf.keras.layers.GRU(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)
# Loop through both lists simultaneously using zip()
for gru, dropout in zip(self.gru_layers, self.dropout_layers):
x = gru(x, training=training)
x = dropout(x, training=training)
# If no states are provided, initialize a list of None
if states is None:
states = [None] * len(self.gru_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)
x = dropout(x, training=training)
new_states.append(h)
# Final GRU Layer returns the output 'x' and a single hidden state 'h'
x, h = self.gru_final(x, initial_state=states, training=training)
x = self.dropout_final(x, training=training)
x = self.dense(x, training=training)
if return_state:
return x, h # Returns just 'h' instead of '(h, c)'
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))
@@ -177,17 +170,32 @@ model.compile(optimizer='adam', loss=loss)
# 5. SAMPLING & GENERATION LOGIC
# ==========================================
def produce_sample(model, seed, num_generate=300, temperature=0.7):
# Vectorize the initial seed text
input_chars = vectorize_layer(tf.constant([seed]))
input_ids = input_chars[0][:len(seed)].numpy().tolist()
generated_ids, states = [], None
for _ in range(num_generate):
current_tokens = input_ids[-SEQ_LENGTH:]
predictions, states = model(tf.expand_dims(current_tokens, 0), states=states, return_state=True, training=False)
generated_ids = []
states = None
# Step 1: "Warm up" the model with the seed text to build the initial states
# We pass the entire seed here
predictions, states = model(tf.expand_dims(input_ids, 0), states=states, return_state=True, training=False)
# Get the very last prediction from the seed sequence
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)
# Step 2: Generation loop using ONLY the single newest token and updating states
for _ in range(num_generate - 1):
# Pass ONLY the last predicted token, plus the existing states
predictions, states = model(tf.expand_dims([predicted_id], 0), states=states, return_state=True, training=False)
# Scale by temperature and sample
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)
input_ids.append(predicted_id)
return "".join([vocab[idx] for idx in generated_ids])
@@ -203,16 +211,25 @@ checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_pre
# Stop training when validation loss stops improving for 3 epochs
early_stopping = tf.keras.callbacks.EarlyStopping(
monitor='val_loss',
monitor='loss',
patience=3,
restore_best_weights=True # Automatically rolls back to epoch 15 weights!
restore_best_weights=True # Automatically rolls back best weights
)
# Create the learning rate scheduler callback
lr_scheduler = tf.keras.callbacks.ReduceLROnPlateau(
monitor='loss', # Can change to 'val_loss' if using a validation split
factor=0.5, # Multiply the learning rate by 0.5 when triggered (cuts it in half)
patience=2, # Number of epochs with no improvement before dropping LR
min_lr=1e-6, # Don't let the learning rate drop lower than this
verbose=1 # Prints a message when the learning rate changes
)
# ==========================================
# 6. START RUN
# ==========================================
if dataset is not None:
model.fit(dataset, epochs=EPOCHS, callbacks=[GenerationCallback(), checkpoint_callback])
model.fit(dataset, epochs=EPOCHS, callbacks=[GenerationCallback(), checkpoint_callback, early_stopping, lr_scheduler])
else:
print(f"\n--- Inference Mode ({EMBEDDING_DIM}dim, {RNN_UNITS}units) ---")
print(f"Result: {produce_sample(model, seed=SEED_TEXT, num_generate=300, temperature=0.6)}")