adding wandb

train.py

@@ -3,6 +3,7 @@ import math
 from transformers import GPT2Tokenizer
 from datasets import load_dataset
 import numpy as np
+import wandb  # Import W&B library
 
 from model import minGRULM
 from util import generate_text
@@ -10,21 +11,16 @@ 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
-
 reset_state_every = 16
 validate_every = 100  # Perform validation every 100 training steps
-
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Total context size is {batch_size * seq_length} tokens")
 
@@ -33,10 +29,25 @@ tokenizer.pad_token = tokenizer.eos_token
 vocab_size = tokenizer.vocab_size
 print(f"Tokenzier has {vocab_size} unique tokens")
 
+# 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),
+    settings=wandb.Settings(api_key="860f8753998c6e6dc356914de07e8855aa2f9642")
+})
+
 # ============================
 # Load and Preprocess Dataset
 # ============================
-
 dataset = load_dataset(dataset_path)
 
 def process_function(examples):
@@ -49,23 +60,22 @@ 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
 # ============================
-
-# Split the training set into 90% train and 10% validation
 split_dataset = tokenized_datasets['train'].train_test_split(test_size=1 / validate_every)
 train_dataset = split_dataset['train']
 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
 # ============================
-
 model = minGRULM(
     vocab_size = vocab_size,
     d_model = 384,
@@ -75,6 +85,7 @@ 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
 
 optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
 
@@ -83,7 +94,6 @@ h_states = None
 # ============================
 # Training Loop with Validation
 # ============================
-
 step = 0
 for epoch in range(num_epochs):
     print(f"Starting Epoch {epoch + 1}/{num_epochs}")
@@ -92,10 +102,9 @@ for epoch in range(num_epochs):
         input_ids = torch.tensor(batch['input_ids']).to(device)
 
         # Reset hidden states if needed
-        h_states = h_states if (step % reset_state_every != 0) else None
+        h_states = h_states if (step % reset_state_every!= 0) else None
         str_states = (
-            ''.join(['{:.3f}, '.format(h_states[0][0][0][j].item()) for j in range(10)])
-            if h_states is not None else 'None'
+            ''.join(['{:.3f}, '.format(h_states[0][0][0][j].item()) for j in range(10)] if h_states is not None else 'None'
         )
 
         optimizer.zero_grad()
@@ -104,6 +113,13 @@ for epoch in range(num_epochs):
         optimizer.step()
 
         step += 1
+        # Log step information
+        wandb.log({
+            "step": step,
+            "epoch": epoch + 1,
+            "loss": loss.item(),
+            "hidden_state": str_states
+        })
         print(f"Epoch: {epoch + 1}/{num_epochs}, Step: {step}, Loss: {loss.item():.4f}, Hidden State: {str_states}")
 
         # Perform validation at specified intervals
@@ -122,12 +138,9 @@ for epoch in range(num_epochs):
                     validation_loss += val_loss.item()
                     valid_steps += 1
 
-                    # Optionally, limit the number of batches for faster validation
-                    # Uncomment the following lines to validate on only the first 100 batches
-                    # if valid_steps >= 100:
-                    #     break
-
             avg_validation_loss = validation_loss / valid_steps if valid_steps > 0 else float('inf')
+            # Log validation loss
+            wandb.log({"validation_loss": avg_validation_loss, "step": step})
             print(f"----- Validation after Step {step}: Average Loss = {avg_validation_loss:.4f} -----")
             model.train()  # Switch back to training mode
 
@@ -142,3 +155,5 @@ for epoch in range(num_epochs):
                 prompt = sample_ids.unsqueeze(0)  # Shape: [1, input_len]
                 generated_text = generate_text(model, tokenizer, prompt, num_predict)
                 print(f"Generated Text:\n{generated_text}\n")
+                # Optionally, log generated text (e.g., as HTML to preserve formatting)
+                # wandb.log({"generated_text": wandb.Html(f"<pre>{generated_text}</pre>")}, step=step)