Update ThoughtVectors.py

Added improvements, including more error handling, proper loss calculation/application, and more

ThoughtVectors.py

@@ -1,4 +1,3 @@
-# thought_vectors.py
 import random
 import math
 import torch
@@ -6,7 +5,7 @@ import torch.nn as nn
 import torch.optim as optim
 import sentencepiece as spm
 from torch.utils.data import Dataset, DataLoader
-from torch.cuda.amp import autocast, GradScaler
+from torch.amp import autocast, GradScaler
 from typing import List, Union, Generator, Sequence, Optional
 import os
 import csv
@@ -48,8 +47,7 @@ class PositionalEncoding(nn.Module):
                                  (-math.log(10000.0) / d_model))
             
             pe[:, 0::2] = torch.sin(position * div_term)
-            pe[:, 1::2] = torch.cos(position * div_term[:d_model // 2 + 1] if d_model % 2 else div_term)
-            
+            pe[:, 1::2] = torch.cos(position * div_term[:self.d_model // 2 + (self.d_model % 2)])
             pe = pe.unsqueeze(0)
             self.register_buffer('pe', pe)
         else:
@@ -107,7 +105,8 @@ class ThoughtEncoder(nn.Module):
                  nhead: int = 8, num_layers: int = 2, dropout: float = 0.1, pad_id: int = 0,
                  termination_threshold: float = 0.75, max_len: int = 5000):
         super().__init__()
-        assert nhead % 2 == 0, "nhead must be even for thought_attention"
+        if nhead % 2 != 0:
+            raise ValueError(f"nhead must be even for thought_attention, got {nhead}")
         self.d_model = d_model
         self.max_thoughts = max_thoughts
         self.pad_id = pad_id
@@ -164,7 +163,7 @@ class ThoughtEncoder(nn.Module):
             
             thought_vectors_list.append(next_thought)
             finished = finished | (termination_score > self.termination_threshold)
-            if finished.all():
+            if finished.all() and not force_single_vector:
                 break
         
         return torch.stack(thought_vectors_list, dim=1)
@@ -202,6 +201,8 @@ class ThoughtDecoder(nn.Module):
         Returns:
             torch.Tensor: Mask tensor, shape (sz, sz).
         """
+        if self.positional_encoding.max_len > 0 and sz > self.positional_encoding.max_len:
+            raise ValueError(f"Sequence length {sz} exceeds max_len {self.positional_encoding.max_len}")
         mask = (torch.triu(torch.ones(sz, sz, device=device)) == 1).transpose(0, 1)
         mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
         return mask
@@ -302,43 +303,40 @@ class LazyList(Sequence):
             else:
                 self._length = 4
         return self._length
-
-    def __getitem__(self, index: int) -> List[str]:
-        """Return a single row with semi-random skipping; index is ignored.
-
-        Args:
-            index (int): Ignored parameter (kept for Sequence compatibility).
-
-        Returns:
-            List[str]: A row from the data, with ~25% chance of skipping each row.
-        """
+    
+    def __getitem__(self, index: Union[int, slice]) -> Union[List[str], List[List[str]]]:
+        if isinstance(index, slice):
+            start, stop, step = index.indices(len(self))
+            count = (stop - start) // (step or 1)
+            return [self._get_single_item() for _ in range(max(0, count))]
+        return self._get_single_item()
+    
+    def _get_single_item(self) -> List[str]:
         if self.data_list is not None:
-            # For in-memory list, pick a random row (faster than iterating)
             return random.choice(self.data_list)
-
         elif self.filepath and os.path.exists(self.filepath):
             self._open_file()
             try:
                 while True:
                     row = next(self._reader)
-                    if row and random.random() <= self.sample_prob:  # 75% chance to return
+                    if row and random.random() <= self.sample_prob:
                         return row
             except StopIteration:
-                # Reset file pointer to start if we reach the end
                 self._file_handle.seek(0)
                 self._reader = csv.reader(self._file_handle)
-                row = next(self._reader)  # Return first row if available
-                if row:
-                    return row
-
-        # Dummy data fallback
-        dummy_data = [
+                row = next(self._reader)
+                return row if row else random.choice(self._dummy_data())
+        print("WARN: file error, file may be empty, missing, or inaccessible. Dummy data returned")  # change to raise in release
+        return random.choice(self._dummy_data())  # remove in release
+    
+    # remove in release
+    def _dummy_data(self) -> List[List[str]]:
+        return [
             ["Hello world", "AI is cool"],
             ["This is a test", "Another sentence"],
             ["Python is fun", "Coding rocks"],
             ["Short sentence", "Quick test"],
         ]
-        return random.choice(dummy_data)
 
     def __del__(self):
         """Close the file handle when the object is destroyed."""
@@ -441,6 +439,7 @@ class ThoughtVectors:
         test_lazy_data = LazyList(test_data) if test_data is not None else None
         
         if model and os.path.exists(model):
+            print("loading model...")
             checkpoint = torch.load(model, map_location=self.device)
             self.sp.Load(checkpoint['spm_model_path'])
             self.vocab_size = self.sp.GetPieceSize()
@@ -457,6 +456,7 @@ class ThoughtVectors:
             self.decoder.load_state_dict(checkpoint['decoder_state_dict'])
             print(f"Loaded pre-existing model from {model}")
         else:
+            print("building models")
             with tempfile.NamedTemporaryFile(mode="w", encoding="utf8", delete=False) as temp_file:
                 for group in train_lazy_data:
                     for sentence in group:
@@ -466,6 +466,8 @@ class ThoughtVectors:
             temp_dir = tempfile.mkdtemp()
             temp_spm_prefix = os.path.join(temp_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,
@@ -499,56 +501,54 @@ class ThoughtVectors:
                                lr=learning_rate, weight_decay=weight_decay)
         scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)
         criterion = nn.CrossEntropyLoss(ignore_index=self.sp.pad_id())
-        scaler = GradScaler()
+        scaler = GradScaler('cuda')
         
         best_val_loss = float('inf')
         patience_counter = 0
         
         print("Beginning training")
         for epoch in range(num_epochs):
-            self.encoder.train()
-            self.decoder.train()
             total_train_loss = 0.0
-            optimizer.zero_grad()
+            num_train_batches = (len(train_lazy_data) + batch_size - 1) // batch_size
             
-            # Manual batching from train_lazy_data with tokenization
-            train_data_list = list(train_lazy_data)  # List[str]
-            random.shuffle(train_data_list)
-            num_train_batches = (len(train_data_list) + batch_size - 1) // batch_size
-            
-            for i in range(0, len(train_data_list), batch_size):
-                raw_batch = train_data_list[i:i + batch_size]
-                # Tokenize each group in the batch
+            for i in range(0, len(train_lazy_data), batch_size):
+                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]
                     for group in raw_batch
                 ]
-                group_batch = group_collate_fn(tokenized_batch)  # Process into padded tensors
+                group_batch = group_collate_fn(tokenized_batch)
                 
                 batch_loss = 0.0
-                with autocast():
+                with autocast('cuda'):
                     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)
                             force_single_vector = random.random() < single_vector_prob
                             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 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:
                     scaler.step(optimizer)
                     scaler.update()
                     optimizer.zero_grad()
                     torch.cuda.empty_cache()
                 
-                total_train_loss += batch_loss.item() * accum_steps
-                print(
-                    f"Batch {i // batch_size + 1}/{num_train_batches} - Train Loss: {batch_loss:.4f}")
+                print(f"Batch {i // batch_size + 1}/{num_train_batches} - Loss: {batch_loss.item():.4f}")
+                # Debugging
+                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')
             
@@ -557,13 +557,12 @@ class ThoughtVectors:
                 self.encoder.eval()
                 self.decoder.eval()
                 total_val_loss = 0.0
-                test_data_list = list(test_lazy_data)
-                num_val_batches = (len(test_data_list) + batch_size - 1) // batch_size
+                num_val_batches = (len(test_lazy_data) + batch_size - 1) // batch_size
+                print(f"Starting validation with {num_val_batches} batches")
                 
                 with torch.no_grad():
-                    for i in range(0, len(test_data_list), batch_size):
-                        raw_batch = test_data_list[i:i + batch_size]
-                        # Tokenize each group in the batch
+                    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
@@ -578,14 +577,17 @@ class ThoughtVectors:
                                 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}")
+            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
@@ -626,14 +628,14 @@ class ThoughtVectors:
         
         return self.encoder(input_tokens, force_single_vector)
     
-    def decode(self, thought_vectors: torch.Tensor, max_length: int = 50, beam_width: Optional[int] = 5,
+    def decode(self, thought_vectors: torch.Tensor, max_length: int = 50, beam_width: int = 0,
                temperature: float = 1.0) -> List[str]:
         """Decodes thought vectors into text sequences.
 
         Args:
             thought_vectors (torch.Tensor): Thought vectors, shape (batch_size, num_thoughts, d_model).
             max_length (int, optional): Maximum length of generated sequences. Defaults to 50.
-            beam_width (Optional[int], optional): Beam width for beam search; if None, uses greedy
+            beam_width (Optional[int], optional): Beam width for beam search; if 0, uses greedy
                 decoding. Defaults to 5.
             temperature (float, optional): Temperature for softmax sampling in greedy decoding.
                 Defaults to 1.0.
@@ -647,13 +649,14 @@ class ThoughtVectors:
         thought_vectors = thought_vectors.to(self.device)
         batch_size = thought_vectors.size(0)
         
-        if beam_width is not None and beam_width > 1:
+        if beam_width > 1:
             return self._beam_search_decode(thought_vectors, max_length, beam_width)
         
         target_tokens = torch.full((batch_size, 1), self.sp.bos_id(), dtype=torch.long, device=self.device)
         for _ in range(max_length - 1):
             logits = self.decoder(thought_vectors, target_tokens)
             logits = logits[:, -1, :] / temperature
+
             next_token = logits.softmax(dim=-1).multinomial(1)
             target_tokens = torch.cat([target_tokens, next_token], dim=1)
             if (next_token == self.sp.eos_id()).all():
@@ -702,7 +705,8 @@ class ThoughtVectors:
         results = []
         for i in range(batch_size):
             combined = sorted(finished[i] + beams[i], key=lambda x: x[0])
-            results.append(self.sp.DecodeIds(combined[0][1].squeeze(0).tolist()) if combined else "<empty>")
+            results.append(self.sp.DecodeIds(combined[0][1].squeeze(0).tolist()) if combined else self.sp.DecodeIds(
+                beams[i][0][1].squeeze(0).tolist()))
         
         return results
     
@@ -756,10 +760,11 @@ class ThoughtVectors:
             ThoughtVectors: Loaded model instance.
         """
         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")
+            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.")
@@ -778,6 +783,7 @@ class ThoughtVectors:
         translator.dropout = checkpoint['dropout']
         translator.max_len = checkpoint['max_len']
         translator.termination_threshold = checkpoint['termination_threshold']
+
         
         translator.encoder = ThoughtEncoder(
             vocab_size=translator.vocab_size, d_model=translator.d_model, max_thoughts=translator.max_thoughts,
@@ -796,16 +802,16 @@ class ThoughtVectors:
 
 
 if __name__ == "__main__":
-    csv.field_size_limit(8 ** 8)
     tv = ThoughtVectors()
     tv.train(
-        group_data="data.csv",
+        group_data="train.csv",
+        test_data="val.csv",
         num_epochs=1000,
-        batch_size=64,
-        accum_steps=4,
-        learning_rate=1e-4,
+        batch_size=128,
+        accum_steps=1,
+        learning_rate=5e-4,
         weight_decay=2e-5,
-        length_penalty=0.01,
+        length_penalty=0.001,
         single_vector_prob=0.1,
         save_path="thought_vectors_prototype.tar",
         spm_model_prefix="spm",
@@ -813,18 +819,18 @@ if __name__ == "__main__":
         d_model=512,
         encoder_nhead=8,
         decoder_nhead=8,
-        encoder_layers=2,
-        decoder_layers=2,
-        max_thoughts=32,
+        encoder_layers=4,
+        decoder_layers=4,
+        max_thoughts=16,
         dropout=0.1,
-        max_len=1024,
+        max_len=256,
         termination_threshold=0.8,
         patience=5
     )
     loaded_tv = ThoughtVectors.load("thought_vectors_prototype.tar")
     thought_vectors = loaded_tv.encode("AI is smart")
-    generated_text_greedy = loaded_tv.decode(thought_vectors, temperature=0.7)
-    generated_text_beam = loaded_tv.decode(thought_vectors, beam_width=3)
+    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)
     print(f"Greedy decoding: {generated_text_greedy}")
     print(f"Beam search decoding: {generated_text_beam}")