adding validation

train.py

@@ -1,5 +1,3 @@
-#@title Utility functions for sampling
-
 import torch
 import math
 from transformers import GPT2Tokenizer
@@ -9,11 +7,9 @@ import numpy as np
 from model import minGRULM
 from util import generate_text
 
-
-
-
-
-
+# ============================
+# Configuration Parameters
+# ============================
 
 dataset_path  = 'flpelerin/tinystories-100k'
 
@@ -28,19 +24,19 @@ num_predict   = 250
 
 reset_state_every = 16
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"total context size is {batch_size * seq_length} tokens");
-
-
+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")
 
 tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
 tokenizer.pad_token = tokenizer.eos_token
 vocab_size = tokenizer.vocab_size
-print(f"tokenizer has {vocab_size} unique tokens")
-
-
+print(f"Tokenzier has {vocab_size} unique tokens")
 
+# ============================
+# Load and Preprocess Dataset
+# ============================
 
 dataset = load_dataset(dataset_path)
 
@@ -48,17 +44,23 @@ def process_function(examples):
     return tokenizer(examples['text'], padding='longest', truncation=True)
 
 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")
+print(f"Dataset has {tokenized_datasets['train'].num_rows} rows of {batch_size} times {seq_length} tokens")
 
+# ============================
+# 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=0.1)
+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)}")
 
-#model = minGRULM(
-#    vocab_size = vocab_size,
-#    d_model = 768,
-#    d_inner = 1536,
-#    n_layers = 12
-#)
+# ============================
+# Initialize the Model
+# ============================
 
 model = minGRULM(
     vocab_size = vocab_size,
@@ -67,29 +69,30 @@ model = minGRULM(
     n_layers = 6
 )
 
-
 model.to(device)
-print(f"model has {sum(p.numel() for p in model.parameters()):,} parameters")
-
-
-
+print(f"Model has {sum(p.numel() for p in model.parameters()):,} parameters")
 
 optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
 
 h_states = None
 
+# ============================
+# Training Loop with Validation
+# ============================
+
 step = 0
 for epoch in range(num_epochs):
-    for i in range(0, len(tokenized_datasets['train']), batch_size):
-        batch = tokenized_datasets['train'][i:i + batch_size]
+    print(f"Starting Epoch {epoch + 1}/{num_epochs}")
+    for i in range(0, len(train_dataset), batch_size):
+        batch = train_dataset[i:i + batch_size]
         input_ids = torch.tensor(batch['input_ids']).to(device)
 
-        #if (i / batch_size) % reset_state_every == 0:
-        #    print(f"resetting state, {(i / batch_size)} % {reset_state_every} == 0")
-        #    h_states = None
-            
-        h_states = h_states if (i / batch_size) % reset_state_every != 0 else None
-        str_states = ''.join(['{:.3f}, '.format(h_states[0][0][0][i].item()) for i in range(10)]) if h_states is not None else 'None'
+        # Reset hidden states if needed
+        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'
+        )
 
         optimizer.zero_grad()
         _, h_states, loss = model.forward(input_ids, h_states)
@@ -97,17 +100,41 @@ for epoch in range(num_epochs):
         optimizer.step()
 
         step += 1
-        print(f"Epoch: {epoch} / {num_epochs}, Step: {step}, Loss: {loss.item():.4f}, Hidden State: {str_states}")
+        print(f"Epoch: {epoch + 1}/{num_epochs}, Step: {step}, Loss: {loss.item():.4f}, Hidden State: {str_states}")
 
-        if step % infer_step == 0:
+        # Perform validation at specified intervals
+        if step % validate_every == 0:
             model.eval()
+            validation_loss = 0.0
+            valid_steps = 0
+
+            with torch.no_grad():
+                for vi in range(0, len(valid_dataset), batch_size):
+                    val_batch = valid_dataset[vi:vi + batch_size]
+                    val_input_ids = torch.tensor(val_batch['input_ids']).to(device)
 
-            ids = input_ids[0][:input_len]
-            text = tokenizer.decode(ids)
-            print(f"input: {text}")
+                    # Forward pass
+                    _, _, val_loss = model.forward(val_input_ids, None)
+                    validation_loss += val_loss.item()
+                    valid_steps += 1
 
-            prompt = ids[None, ...]
-            text = generate_text(model, tokenizer, prompt, num_predict)
-            print(f"output: {text}")
+                    # 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
 
-            model.train()
+            avg_validation_loss = validation_loss / valid_steps if valid_steps > 0 else float('inf')
+            print(f"----- Validation after Step {step}: Average Loss = {avg_validation_loss:.4f} -----")
+            model.train()  # Switch back to training mode
+
+        # Perform inference at specified steps
+        if step % infer_step == 0:
+            with torch.no_grad():
+                # Select a single input from the current batch for inference
+                sample_ids = input_ids[0][:input_len]
+                input_text = tokenizer.decode(sample_ids, skip_special_tokens=True)
+                print(f"Input for Inference: {input_text}")
+
+                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")