Updated training data and model definitions
This commit is contained in:
+24510
File diff suppressed because it is too large
Load Diff
@@ -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)
|
||||
input_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)
|
||||
|
||||
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)}")
|
||||
|
||||
Reference in New Issue
Block a user