308 lines
13 KiB
Python
308 lines
13 KiB
Python
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
|
|
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
|
|
|
|
# ==========================================
|
|
# 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):
|
|
# 1. Safely tokenize the seed character-by-character using the vectorizer
|
|
# We split the string into individual characters so the vectorizer processes them properly
|
|
seed_chars = tf.strings.unicode_split(seed, input_encoding='UTF-8')
|
|
|
|
input_ids = tf.squeeze(vectorize_layer(seed_chars))
|
|
if input_ids.ndim == 0: # Safeguard if the seed is a single character
|
|
input_ids = tf.expand_dims(input_ids, 0)
|
|
|
|
#input_ids = vectorize_layer(seed_chars).numpy().tolist()
|
|
|
|
generated_ids = []
|
|
states = None
|
|
|
|
# 2. Process the entire seed sequence first to warm up the LSTM states
|
|
predictions, states = model(tf.expand_dims(input_ids, 0), states=states, return_state=True, training=False)
|
|
|
|
# Grab the logits for the very last character of the seed
|
|
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)
|
|
|
|
# 3. Autoregressive loop for the remaining characters
|
|
for _ in range(num_generate - 1):
|
|
# We pass only the single, last-predicted character ID
|
|
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)
|
|
|
|
# 4. Join the vocabulary tokens, ignoring any special [PAD] or [UNK] tokens if they sneak in
|
|
generated_text = "".join([vocab[idx] for idx in generated_ids if idx < len(vocab)])
|
|
|
|
# Return the seed + the newly generated text
|
|
return seed + generated_text
|
|
|
|
|
|
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 = 256
|
|
DEFAULT_RNN_UNITS = 1024
|
|
DEFAULT_NUM_LAYERS = 3
|
|
DEFAULT_DROPOUT_RATE = 0.4
|
|
|
|
#TEXT_FILE = "emily_post.txt"
|
|
TEXT_FILE = "combined_training_data.txt"
|
|
VOCAB_FILE = "vocab.txt"
|
|
SEQ_LENGTH = 512
|
|
BATCH_SIZE = 128
|
|
EPOCHS_PER_RUN = 22
|
|
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)
|
|
|
|
dummy_input = tf.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=tf.int32)
|
|
_ = model(dummy_input, training=False)
|
|
|
|
#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}...")
|
|
print(f"[DEBUG] Verification -> initial_epoch: {initial_epoch}, total_epochs: {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()
|