update

train.py

@@ -12,43 +12,43 @@ from util import generate_text
 # Configuration Parameters
 # ============================
 dataset_path  = 'flpelerin/tinystories-10k'
+
 num_epochs    = 1
 batch_size    = 4
 seq_length    = 256
 learning_rate = 1e-4
-infer_step    = 50
+
 input_len     = 50
 num_predict   = 250
+
+infer_every       = 50
 reset_state_every = 16
-validate_every = 100  # Perform validation every 100 training steps
+validate_every    = 100  # Perform validation every 100 training steps
+
+
+
+# ============================
+# Initialize the Device
+# ============================
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Total context size is {batch_size * seq_length} tokens")
 
+
+
+# ============================
+# Initialize the Tokenizer
+# ============================
+
+
 tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
 tokenizer.pad_token = tokenizer.eos_token
 vocab_size = tokenizer.vocab_size
 print(f"Tokenzier has {vocab_size} unique tokens")
 
 
-wandb.login(key="860f8753998c6e6dc356914de07e8855aa2f9642")
 
-# Initialize W&B Logging
-wandb.init(
-    project="minGRU-Training", 
-    config={
-        "dataset_path": dataset_path,
-        "num_epochs": num_epochs,
-        "batch_size": batch_size,
-        "seq_length": seq_length,
-        "learning_rate": learning_rate,
-        "infer_step": infer_step,
-        "input_len": input_len,
-        "num_predict": num_predict,
-        "reset_state_every": reset_state_every,
-        "validate_every": validate_every,
-        "device": str(device)
-    }
-)
+
 
 # ============================
 # Load and Preprocess Dataset
@@ -65,7 +65,9 @@ def process_function(examples):
 
 tokenized_datasets = dataset.map(process_function, batched=True)
 print(f"Dataset has {tokenized_datasets['train'].num_rows} rows of {batch_size} times {seq_length} tokens")
-wandb.log({"dataset_rows": tokenized_datasets['train'].num_rows, "dataset_token_count": batch_size * seq_length})  # Log dataset stats
+
+
+
 
 # ============================
 # Split Dataset into Train and Validation
@@ -76,7 +78,8 @@ valid_dataset = split_dataset['test']
 
 print(f"Training set size: {len(train_dataset)}")
 print(f"Validation set size: {len(valid_dataset)}")
-wandb.log({"train_set_size": len(train_dataset), "valid_set_size": len(valid_dataset)})  # Log set sizes
+
+
 
 # ============================
 # Initialize the Model
@@ -89,17 +92,59 @@ model = minGRULM(
 )
 
 model.to(device)
-print(f"Model has {sum(p.numel() for p in model.parameters()):,} parameters")
-wandb.log({"model_parameters": sum(p.numel() for p in model.parameters())})  # Log model parameter count
+parameters_count = sum(p.numel() for p in model.parameters())
+print(f"Model has {parameters_count:,} parameters")
+
+
+
+# ============================
+# Initialize the Weights and Biases Run
+# ============================
+wandb.login(key="860f8753998c6e6dc356914de07e8855aa2f9642")
+wandb.init(
+    project="minGRU-Training", 
+    config={
+        "dataset_path": dataset_path,
+        
+        "num_epochs": num_epochs,
+        "batch_size": batch_size,
+        "seq_length": seq_length,
+        "learning_rate": learning_rate,
+
+        "input_len": input_len,
+        "num_predict": num_predict,
+
+        "infer_every": infer_every,
+        "reset_state_every": reset_state_every,
+        "validate_every": validate_every,
+
+        "dataset_rows": tokenized_datasets['train'].num_rows,
+        "dataset_token_count": batch_size * seq_length,
+        
+        "train_set_size": len(train_dataset),
+        "valid_set_size": len(valid_dataset),
+
+        "model_parameters": parameters_count,
+
+        "vocab_size": vocab_size,
+        "d_model": 384,
+        "d_inner": 768,
+        "n_layers": 6,
+
+        "device": str(device)
+    }
+)
 
-optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
 
-h_states = None
 
 # ============================
 # Training Loop with Validation
 # ============================
+
+optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+h_states = None
 step = 0
+
 for epoch in range(num_epochs):
     print(f"Starting Epoch {epoch + 1}/{num_epochs}")
     for i in range(0, len(train_dataset), batch_size):
@@ -150,7 +195,7 @@ for epoch in range(num_epochs):
             model.train()  # Switch back to training mode
 
         # Perform inference at specified steps
-        if step % infer_step == 0:
+        if step % infer_every == 0:
             with torch.no_grad():
                 # Select a single input from the current batch for inference
                 sample_ids = input_ids[0][:input_len]