Fixed more bugs, solved an issue with 16FP used for the GRU.

ThoughtVectors.py

@@ -1,13 +1,12 @@
-# thought_vectors.py
+import atexit
 import random
 import math
 import torch
 import torch.nn as nn
 import torch.optim as optim
 import sentencepiece as spm
-from torch.utils.data import Dataset, DataLoader
 from torch.amp import autocast, GradScaler
-from typing import List, Union, Generator, Sequence, Optional
+from typing import List, Union, Generator, Sequence
 import os
 import csv
 import tempfile
@@ -156,9 +155,12 @@ class ThoughtEncoder(nn.Module):
                 prev_context = torch.zeros_like(text_context)
             
             combined = torch.cat([text_context, prev_context], dim=1)
+            
             next_thought = self.fc_thought(combined)
             
-            thought_hidden = self.thought_rnn(next_thought.unsqueeze(1), thought_hidden)[1]
+            with autocast('cuda', enabled=False):  # FP32
+                thought_hidden = self.thought_rnn(next_thought.unsqueeze(1).float(), thought_hidden.float())[1]
+            
             termination_logit = self.fc_terminate(thought_hidden.squeeze(0))
             termination_score = torch.sigmoid(termination_logit).squeeze(-1)
             
@@ -372,13 +374,20 @@ class ThoughtVectors:
         self.dropout = None
         self.max_len = None
         self.termination_threshold = None
+        
+        # Track files to clean up
+        self._temp_files = set()
+        # Register cleanup function
+        atexit.register(self._cleanup)
     
     def train(self,
               group_data: Union[str, List[List[str]]],
               test_data: Union[str, List[List[str]]] = None,
               model: str = None,
-              num_epochs: int = 30,
-              batch_size: int = 8,
+              num_epochs: int = 2048,
+              batch_size: int = 128,
+              batches_per_val: int = 16,
+              val_batches: int = 8,
               accum_steps: int = 4,
               learning_rate: float = 2e-4,
               weight_decay: float = 1e-5,
@@ -405,6 +414,8 @@ class ThoughtVectors:
             model (str, optional): Path to pre-existing model to load. If None, creates a new model. Defaults to None.
             num_epochs (int, optional): Number of training epochs. Defaults to 30.
             batch_size (int, optional): Batch size. Defaults to 8.
+            batches_per_val (int, optional): Validate model every x batches. Defaults to 16
+            val_batches (int, optional): Number of batches to use in validation, speeds up validation. Defaults to 8
             accum_steps (int, optional): Number of gradient accumulation steps. Defaults to 4.
             learning_rate (float, optional): Learning rate for Adam optimizer. Defaults to 2e-4.
             weight_decay (float, optional): Weight decay for regularization. Defaults to 1e-5.
@@ -422,7 +433,7 @@ class ThoughtVectors:
             dropout (float, optional): Dropout probability. Defaults to 0.1.
             max_len (int, optional): Maximum sequence length for positional encoding. Defaults to 1024.
             termination_threshold (float, optional): Threshold for stopping thought generation. Defaults to 0.75.
-            patience (int, optional): Epochs to wait for improvement before early stopping. Defaults to 5.
+            patience (int, optional): BATCHES to wait for improvement before early stopping. Defaults to 5.
         """
         self.d_model = d_model
         self.max_thoughts = max_thoughts
@@ -442,7 +453,7 @@ class ThoughtVectors:
         if model and os.path.exists(model):
             print("loading model...")
             loaded_instance = self.load(model)
-            self.__dict__.update(loaded_instance.__dict__)  # actually update self
+            self.__dict__.update(loaded_instance.__dict__)
             print(f"Loaded pre-existing model from {model}")
         else:
             print("building models")
@@ -452,11 +463,11 @@ class ThoughtVectors:
                         temp_file.write(sentence.strip() + "\n")
                 temp_dataset_path = temp_file.name
             
-            temp_dir = tempfile.mkdtemp()
-            temp_spm_prefix = os.path.join(temp_dir, "spm")
+            # Use a persistent directory instead of temp_dir to avoid premature cleanup
+            os.makedirs("temp_spm_dir", exist_ok=True)
+            temp_spm_prefix = os.path.join("temp_spm_dir", "spm")
             try:
                 print("training sentence piece (could take a while on massive datasets)")
-                print("Also you will get some log spamming from it cause it won't shut up")
                 spm.SentencePieceTrainer.Train(
                     input=temp_dataset_path,
                     model_prefix=temp_spm_prefix,
@@ -467,13 +478,25 @@ class ThoughtVectors:
                     input_sentence_size=8_388_608,
                     train_extremely_large_corpus=True
                 )
+                # Ensure files are copied to the final location
                 shutil.copy(f"{temp_spm_prefix}.model", f"{spm_model_prefix}.model")
                 shutil.copy(f"{temp_spm_prefix}.vocab", f"{spm_model_prefix}.vocab")
                 if not self.sp.Load(f"{spm_model_prefix}.model"):
                     raise RuntimeError("Failed to load SentencePiece model.")
             finally:
                 os.remove(temp_dataset_path)
-                shutil.rmtree(temp_dir, ignore_errors=True)
+                # Only remove temp_spm_dir after successful save, handled in save method
+            
+            self.vocab_size = self.sp.GetPieceSize()
+            self.encoder = ThoughtEncoder(
+                vocab_size=self.vocab_size, d_model=d_model, max_thoughts=max_thoughts,
+                nhead=encoder_nhead, num_layers=encoder_layers, dropout=dropout,
+                termination_threshold=termination_threshold, max_len=self.max_len
+            ).to(self.device)
+            self.decoder = ThoughtDecoder(
+                vocab_size=self.vocab_size, d_model=d_model, num_layers=decoder_layers,
+                nhead=decoder_nhead, dropout=dropout, max_len=self.max_len
+            ).to(self.device)
             
             self.vocab_size = self.sp.GetPieceSize()
             self.encoder = ThoughtEncoder(
@@ -494,6 +517,7 @@ class ThoughtVectors:
         
         best_val_loss = float('inf')
         patience_counter = 0
+        best_batch = 0  # Track best batch number
         
         print("Beginning training")
         try:
@@ -501,8 +525,8 @@ class ThoughtVectors:
                 total_train_loss = 0.0
                 num_train_batches = (len(train_lazy_data) + batch_size - 1) // batch_size
                 
-                # training
                 for i in range(0, len(train_lazy_data), batch_size):
+                    batch_idx = i // batch_size + 1
                     raw_batch = train_lazy_data[i:i + batch_size]
                     tokenized_batch = [
                         [[self.sp.bos_id()] + self.sp.EncodeAsIds(s) + [self.sp.eos_id()] for s in group]
@@ -519,101 +543,114 @@ class ThoughtVectors:
                                 thought_vectors = self.encoder(group_tensor, force_single_vector)
                                 output_logits = self.decoder(thought_vectors, group_tensor[:, :-1])
                                 
-                                # Masked loss
                                 mask = (group_tensor[:, 1:] != self.sp.pad_id()).float()
                                 loss = criterion(output_logits.reshape(-1, self.vocab_size),
                                                  group_tensor[:, 1:].reshape(-1))
                                 loss = (loss * mask.reshape(-1)).sum() / mask.sum().clamp(min=1.0)
+                                if torch.isnan(loss):
+                                    print(f"NaN detected: mask_sum={mask.sum().item()}, loss_pre_mask={loss.item()}")
+                                    print(f"group_tensor={group_tensor}\n\n")
+                                    print(f"thought vectors={thought_vectors}\n\n")
+                                    print("Raising keyboard interrupt to allow preservation.")
+                                    raise KeyboardInterrupt
                                 if not force_single_vector:
                                     loss += length_penalty * thought_vectors.shape[1]
                                 batch_loss += loss / accum_steps
                     
                     scaler.scale(batch_loss).backward()
                     total_train_loss += batch_loss.item() * accum_steps
-                    if (i // batch_size + 1) % accum_steps == 0:
+                    if (batch_idx % accum_steps) == 0:
+                        scaler.unscale_(optimizer)
+                        torch.nn.utils.clip_grad_norm_(
+                            list(self.encoder.parameters()) + list(self.decoder.parameters()), max_norm=0.5)
                         scaler.step(optimizer)
                         scaler.update()
                         optimizer.zero_grad()
                         torch.cuda.empty_cache()
                     
-                    print(f"Batch {i // batch_size + 1}/{num_train_batches} - Loss: {batch_loss.item():.4f}")
-                    # Debugging
+                    # Per group of batches validation
+                    if (batch_idx + 1) % batches_per_val == 0 or batch_idx == num_train_batches:
+                        if test_lazy_data is None:
+                            avg_val_loss = batch_loss.item()  # Use train loss if no val
+                        else:
+                            self.encoder.eval()
+                            self.decoder.eval()
+                            total_val_loss = 0.0
+                            num_val_batches = min(val_batches,
+                                                  (len(test_lazy_data) + batch_size - 1) // batch_size)  # Quick val subset
+                            
+                            with torch.no_grad():
+                                val_indices = list(range(0, len(test_lazy_data), batch_size))[:num_val_batches]
+                                for i in val_indices:
+                                    raw_batch = test_lazy_data[i:i + batch_size]
+                                    tokenized_batch = [
+                                        [[self.sp.bos_id()] + self.sp.EncodeAsIds(s) + [self.sp.eos_id()] for s in group]
+                                        for group in raw_batch
+                                    ]
+                                    group_batch = group_collate_fn(tokenized_batch)
+                                    
+                                    batch_val_loss = 0.0
+                                    for group_tensor in group_batch:
+                                        if group_tensor.shape[1] <= self.max_len and group_tensor.numel() > 0:
+                                            group_tensor = group_tensor.to(self.device)
+                                            thought_vectors = self.encoder(group_tensor)
+                                            output_logits = self.decoder(thought_vectors, group_tensor[:, :-1])
+                                            val_loss = criterion(output_logits.reshape(-1, self.vocab_size),
+                                                                 group_tensor[:, 1:].reshape(-1))
+                                            val_loss += length_penalty * thought_vectors.shape[1]
+                                            batch_val_loss += val_loss.item()
+                                    total_val_loss += batch_val_loss
+                                    print(f" Input: {self.sp.DecodeIds(group_tensor[0].tolist())}")
+                                    print(f"Output: {self.sp.DecodeIds(output_logits.argmax(-1)[0].tolist())}\n")
+                            
+                            avg_val_loss = total_val_loss / num_val_batches if num_val_batches > 0 else float('inf')
+                            print(f"Val Loss: {avg_val_loss:.4f}\n\n")
+                            
+                            # Per-batch patience
+                        if avg_val_loss < best_val_loss:
+                            best_val_loss = avg_val_loss
+                            best_batch = batch_idx + epoch * num_train_batches  # Total batch count
+                            if patience_counter > 0:
+                                patience_counter -= 1
+                            self.save(save_path, spm_model_prefix)
+                        else:
+                            patience_counter += 1
+                            if patience_counter >= patience:
+                                print(
+                                    f"Early stopping triggered at batch {batch_idx} (total {best_batch + patience}), best val loss: {best_val_loss:.4f}")
+                                self.load(save_path)  # Roll back to best
+                                self.save(save_path, spm_model_prefix)  # Save and clean up
+                                return  # Exit training
+                    
+                    print(f"Batch {batch_idx}/{num_train_batches} - Loss: {batch_loss.item():.4f}")
                     print(f" Input: {self.sp.DecodeIds(group_tensor[0].tolist())}")
-                    print(f"Output: {self.sp.DecodeIds(output_logits.argmax(-1)[0].tolist())}\n\n")
-                
-                avg_train_loss = total_train_loss / num_train_batches if num_train_batches > 0 else float('inf')
-                
-                # Validation
-                if test_lazy_data is not None:
-                    self.encoder.eval()
-                    self.decoder.eval()
-                    total_val_loss = 0.0
-                    num_val_batches = (len(test_lazy_data) + batch_size - 1) // batch_size
-                    print(f"Starting validation with {num_val_batches} batches")
+                    print(f"Output: {self.sp.DecodeIds(output_logits.argmax(-1)[0].tolist())}\n")
                     
-                    with torch.no_grad():
-                        for i in range(0, len(test_lazy_data), batch_size):
-                            raw_batch = test_lazy_data[i:i + batch_size]
-                            tokenized_batch = [
-                                [[self.sp.bos_id()] + self.sp.EncodeAsIds(s) + [self.sp.eos_id()] for s in group]
-                                for group in raw_batch
-                            ]
-                            group_batch = group_collate_fn(tokenized_batch)
-                            
-                            batch_val_loss = 0.0
-                            for group_tensor in group_batch:
-                                if group_tensor.shape[1] <= self.max_len and group_tensor.numel() > 0:
-                                    group_tensor = group_tensor.to(self.device)
-                                    thought_vectors = self.encoder(group_tensor)
-                                    output_logits = self.decoder(thought_vectors, group_tensor[:, :-1])
-                                    val_loss = criterion(output_logits.reshape(-1, self.vocab_size),
-                                                         group_tensor[:, 1:].reshape(-1))
-                                    val_loss += length_penalty * thought_vectors.shape[1]
-                                    batch_val_loss += val_loss.item()
-                            total_val_loss += batch_val_loss
-                            print(f"VAL batch loss: {batch_val_loss:.4f}")
-                            print(f" Input: {self.sp.DecodeIds(group_tensor[0].tolist())}")
-                            print(f"Output: {self.sp.DecodeIds(output_logits.argmax(-1)[0].tolist())}\n\n")
-                    avg_val_loss = total_val_loss / num_val_batches if num_val_batches > 0 else float('inf')
-                else:
-                    avg_val_loss = avg_train_loss
-                
-                print(f"Epoch {epoch + 1}/{num_epochs} - Train Loss: {avg_train_loss:.4f} - Val Loss: {avg_val_loss:.4f}\n\n")
-                
-                if avg_val_loss < best_val_loss:
-                    best_val_loss = avg_val_loss
-                    patience_counter = 0
-                    self.save(save_path, spm_model_prefix, clean=False)
-                else:
-                    self.load(save_path)
-                    self.save(save_path, spm_model_prefix, clean=True)
-                    patience_counter += 1
-                    if patience_counter >= patience:
-                        print(f"Early stopping triggered after {epoch + 1} epochs.")
-                        break
+                    self.encoder.train()
+                    self.decoder.train()
                 
+                avg_train_loss = total_train_loss / num_train_batches if num_train_batches > 0 else float('inf')
+                print(f"Epoch {epoch + 1}/{num_epochs} - Train Loss: {avg_train_loss:.4f}\n")
                 scheduler.step()
-        # Allow saving model on keyboard interrupt
+            
         except KeyboardInterrupt:
             while True:
-                saving = input("Would you like to save the best (b), current (c), or no (n) model?")
-                saving = saving.strip().lower
-                if saving == "best" or saving == "b":
-                    self.load(save_path)
-                    self.save(save_path, spm_model_prefix, clean=True)
+                saving = input("Would you like to save the best (b), current (c), or no (n) model: ")
+                saving = saving.strip().lower()
+                if saving in ("best", "b"):
                     break
-                elif saving == "current" or saving == "c" or saving == "curr":
-                    self.save(save_path, spm_model_prefix, clean=True)
+                elif saving in ("current", "c", "curr"):
+                    self.save(save_path, spm_model_prefix)
                     break
-                elif saving == "n" or saving == "no" or saving == "":
-                    self.save(save_path, spm_model_prefix, clean=True)
-                    os.remove(save_path)
+                elif saving in ("n", "no", ""):
+                    if os.path.exists(save_path):
+                        os.remove(save_path)
                     break
                 else:
-                    print("invalid input, type 'b' 'c' or 'n'")
-                
-        print(f"Best model saved to {save_path}")
+                    print("invalid input, type 'b', 'c', or 'n'")
     
+        print(f"Model saved to {save_path}")
+        
     def encode(self, text: Union[str, List[str]], force_single_vector: bool = False) -> torch.Tensor:
         """Encodes text into thought vectors.
 
@@ -721,14 +758,7 @@ class ThoughtVectors:
         
         return results
     
-    def save(self, path: str, spm_model_prefix: str, clean: bool = True) -> None:
-        """Saves the model and SentencePiece data to a tar file.
-
-        Args:
-            path (str): Path to save the tar file.
-            spm_model_prefix (str): Prefix for SentencePiece model files.
-            clean (bool, optional): If True, removes temporary files after saving. Defaults to True.
-        """
+    def save(self, path: str, spm_model_prefix: str) -> None:
         if self.encoder is None or self.decoder is None:
             raise RuntimeError("No model to save.")
         
@@ -749,16 +779,11 @@ class ThoughtVectors:
                 'dropout': self.dropout,
                 'max_len': self.max_len,
                 'termination_threshold': self.termination_threshold,
-                'spm_model_path': f"{spm_model_prefix}.model"  # Added to match loading logic
+                'spm_model_path': f"{spm_model_prefix}.model"
             }, "model.pth")
             tar.add("model.pth")
             tar.add(f"{spm_model_prefix}.model")
             tar.add(f"{spm_model_prefix}.vocab")
-        
-        if clean:
-            os.remove("model.pth")
-            os.remove(f"{spm_model_prefix}.model")
-            os.remove(f"{spm_model_prefix}.vocab")
     
     @classmethod
     def load(cls, path: str) -> 'ThoughtVectors':
@@ -772,17 +797,12 @@ class ThoughtVectors:
         """
         translator = cls()
         translator.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # load based on available device, not model train device
-        try:
-            with tarfile.open(path, "r") as tar:
-                tar.extractall()
-            checkpoint = torch.load("model.pth", map_location=torch.device(translator.device))
-            translator.sp = spm.SentencePieceProcessor()
-            if not translator.sp.Load("spm.model"):
-                raise RuntimeError("Failed to load SentencePiece model.")
-        finally:
-            for f in ["model.pth", "spm.model", "spm.vocab"]:
-                if os.path.exists(f):
-                    os.remove(f)
+        with tarfile.open(path, "r") as tar:
+            tar.extractall()
+        checkpoint = torch.load("model.pth", map_location=torch.device(translator.device))
+        translator.sp = spm.SentencePieceProcessor()
+        if not translator.sp.Load("spm.model"):
+            raise RuntimeError("Failed to load SentencePiece model.")
         
         translator.vocab_size = checkpoint['vocab_size']
         translator.d_model = checkpoint['d_model']
@@ -810,44 +830,40 @@ class ThoughtVectors:
         translator.decoder.load_state_dict(checkpoint['decoder_state_dict'])
         
         return translator
-
-
-if __name__ == "__main__":
-    def process_train(input_file):
-        data = []
-        with open(input_file, 'r', encoding='utf-8') as f_in:
-            reader = csv.reader(f_in)
-            for row in reader:
-                sentence1 = row[0].strip()
-                data.append([sentence1])
-
-        return data
     
-    def process_test(input_file):
-        data = []
-        with open(input_file, 'r', encoding='utf-8') as f_in:
-            reader = csv.reader(f_in)
-            for row in reader:
-                sentence1 = row[0].strip()
-                data.append([sentence1])
-        return data
-        
-    data = process_train("SNLItrain.csv")
-    val = process_test("SNLIval.csv")
-
+    def _add_temp_file(self, filepath: str):
+        """Add a file to the set of temporary files to clean up on close."""
+        self._temp_files.add(os.path.abspath(filepath))
+    
+    def _cleanup(self):
+        """Remove all tracked temporary files."""
+        for filepath in self._temp_files:
+            if os.path.exists(filepath):
+                try:
+                    os.remove(filepath)
+                except OSError as e:
+                    print(f"Failed to clean up {filepath}: {e}")
+        # Clear the set after cleanup
+        self._temp_files.clear()
+
+
+# example usage
+if __name__ == "__main__":
     tv = ThoughtVectors()
     tv.train(
-        group_data=data,
-        test_data=val,
-        model="thought_vectors_prototype-0.1.tar",
-        num_epochs=1000,
-        batch_size=128,
+        group_data="train.csv",
+        test_data="val.csv",
+        model="thought_vectors_prototype.tar",
+        num_epochs=2048,
+        batch_size=256,
+        batches_per_val=32,
+        val_batches=16,
         accum_steps=1,
-        learning_rate=5e-4,
-        weight_decay=2e-5,
+        learning_rate=2e-4,
+        weight_decay=1e-5,
         length_penalty=0.001,
         single_vector_prob=0.1,
-        save_path="thought_vectors_prototype-0.2.tar",
+        save_path="thought_vectors_prototype-0.2.0.tar",
         spm_model_prefix="spm",
         vocab_size=8192,
         d_model=512,
@@ -859,12 +875,11 @@ if __name__ == "__main__":
         dropout=0.1,
         max_len=256,
         termination_threshold=0.8,
-        patience=5
+        patience=10
     )
-    loaded_tv = ThoughtVectors.load("thought_vectors_prototype-0.2.tar")
-    thought_vectors = loaded_tv.encode("AI is smart")
-    generated_text_greedy = loaded_tv.decode(thought_vectors, temperature=0.7, beam_width=0)
-    generated_text_beam = loaded_tv.decode(thought_vectors, beam_width=5)
+    thought_vectors = tv.encode("AI is smart")
+    generated_text_greedy = tv.decode(thought_vectors, temperature=0.7, beam_width=0)
+    generated_text_beam = tv.decode(thought_vectors, beam_width=5)
     print(f"Greedy decoding: {generated_text_greedy}")
     print(f"Beam search decoding: {generated_text_beam}")