Updated to handle big data

You can now pass just a csv file for training data. It will only read the first column of the data and ignore all others

ThoughtVectors.py

@@ -1,13 +1,15 @@
 # thought_vectors.py
-import math
 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, random_split
-from typing import List, Union, Optional, Iterator
+from torch.utils.data import Dataset, DataLoader
+from torch.cuda.amp import autocast, GradScaler
+from typing import List, Union, Generator, Sequence, Optional
 import os
+import csv
 import tempfile
 import tarfile
 import shutil
@@ -42,7 +44,6 @@ class PositionalEncoding(nn.Module):
             # Precompute positional encodings
             pe = torch.zeros(max_len, d_model)
             position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
-            # Adjust div_term for odd d_model
             div_term = torch.exp(torch.arange(0, d_model - (d_model % 2), 2, dtype=torch.float) *
                                  (-math.log(10000.0) / d_model))
             
@@ -104,7 +105,7 @@ class ThoughtEncoder(nn.Module):
     
     def __init__(self, vocab_size: int, d_model: int = 256, max_thoughts: int = 16,
                  nhead: int = 8, num_layers: int = 2, dropout: float = 0.1, pad_id: int = 0,
-                 termination_threshold: float = 0.75):
+                 termination_threshold: float = 0.75, max_len: int = 5000):
         super().__init__()
         assert nhead % 2 == 0, "nhead must be even for thought_attention"
         self.d_model = d_model
@@ -112,7 +113,7 @@ class ThoughtEncoder(nn.Module):
         self.pad_id = pad_id
         self.termination_threshold = termination_threshold
         self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=pad_id)
-        self.positional_encoding = PositionalEncoding(d_model, dropout)
+        self.positional_encoding = PositionalEncoding(d_model, dropout, max_len=max_len)
         encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, dropout=dropout, batch_first=True)
         self.text_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
         self.thought_attention = nn.MultiheadAttention(d_model, num_heads=nhead // 2, dropout=dropout, batch_first=True)
@@ -181,12 +182,12 @@ class ThoughtDecoder(nn.Module):
     """
     
     def __init__(self, vocab_size: int, d_model: int = 256, num_layers: int = 2,
-                 nhead: int = 8, dropout: float = 0.1):
+                 nhead: int = 8, dropout: float = 0.1, max_len: int = 5000):
         super().__init__()
         self.vocab_size = vocab_size
         self.d_model = d_model
         self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
-        self.positional_encoding = PositionalEncoding(d_model, dropout)
+        self.positional_encoding = PositionalEncoding(d_model, dropout, max_len)
         decoder_layer = nn.TransformerDecoderLayer(d_model=d_model, nhead=nhead, dropout=dropout, batch_first=True)
         self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
         self.fc_out = nn.Linear(d_model, vocab_size)
@@ -222,30 +223,6 @@ class ThoughtDecoder(nn.Module):
         return self.fc_out(output)
 
 
-class GroupTextDataset(Dataset):
-    """Dataset class for handling groups of text sentences.
-
-    Args:
-        groups (List[List[str]]): List of sentence groups.
-        sp (spm.SentencePieceProcessor): SentencePiece processor for tokenization.
-    """
-    
-    def __init__(self, groups: List[List[str]], sp: spm.SentencePieceProcessor):
-        self.groups = groups
-        self.sp = sp
-    
-    def __len__(self) -> int:
-        return len(self.groups)
-    
-    def __getitem__(self, idx: int) -> List[List[int]]:
-        group = self.groups[idx]
-        tokenized = []
-        for sentence in group:
-            token_ids = [self.sp.bos_id()] + self.sp.EncodeAsIds(sentence) + [self.sp.eos_id()]
-            tokenized.append(token_ids)
-        return tokenized
-
-
 def group_collate_fn(batch: List[List[List[int]]]) -> List[torch.Tensor]:
     """Collates a batch of tokenized groups into padded tensors.
 
@@ -266,6 +243,111 @@ def group_collate_fn(batch: List[List[List[int]]]) -> List[torch.Tensor]:
     return collated_groups
 
 
+class LazyList(Sequence):
+    def __init__(self, data: Union[str, List[List[str]]], sample_prob: float = 0.75):
+        """Initialize LazyList with data source and sampling probability.
+
+        Args:
+            data (Union[str, List[List[str]]]): A filepath to a CSV or a list of sentence groups.
+            sample_prob (float, optional): Probability of returning a row in __getitem__ (0.0 to 1.0).
+                Defaults to 0.75 (75% chance to return, 25% to skip).
+        """
+        if isinstance(data, str):
+            self.filepath = data
+            self.data_list = None
+        elif isinstance(data, list):
+            self.filepath = None
+            self.data_list = data
+        else:
+            raise ValueError("Data must be a filepath (str) or a list of lists of strings")
+        self.sample_prob = max(0.0, min(1.0, sample_prob))  # Clamp between 0 and 1
+        self._length = None
+        self._file_handle = None  # Persistent file handle for sequential reading
+        self._reader = None  # Persistent CSV reader
+
+    def _open_file(self):
+        """Open the file and initialize the reader if not already open."""
+        if self.filepath and os.path.exists(self.filepath) and self._file_handle is None:
+            self._file_handle = open(self.filepath, 'r', encoding='utf-8', newline='')
+            self._reader = csv.reader(self._file_handle)
+
+    def __iter__(self) -> Generator[List[str], None, None]:
+        """Iterate through all rows sequentially without sampling."""
+        if self.data_list is not None:
+            yield from self.data_list
+        elif self.filepath and os.path.exists(self.filepath):
+            with open(self.filepath, 'r', encoding='utf-8', newline='') as f:
+                reader = csv.reader(f)
+                for row in reader:
+                    if row:
+                        yield row
+        else:
+            dummy_data = [
+                ["Hello world", "AI is cool"],
+                ["This is a test", "Another sentence"],
+                ["Python is fun", "Coding rocks"],
+                ["Short sentence", "Quick test"],
+            ]
+            yield from dummy_data
+
+    def __len__(self) -> int:
+        """Compute and cache the total number of rows."""
+        if self._length is None:
+            if self.data_list is not None:
+                self._length = len(self.data_list)
+            elif self.filepath and os.path.exists(self.filepath):
+                with open(self.filepath, 'r', encoding='utf-8', newline='') as file:
+                    reader = csv.reader(file)
+                    self._length = sum(1 for row in reader if row)
+            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.
+        """
+        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
+                        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 = [
+            ["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."""
+        if self._file_handle is not None:
+            self._file_handle.close()
+            self._file_handle = None
+            self._reader = None
+
+
 class ThoughtVectors:
     """Main class for a thought vector model, encoding text into latent vectors and decoding them back.
 
@@ -291,58 +373,58 @@ class ThoughtVectors:
         self.dropout = None
         self.max_len = None
         self.termination_threshold = None
-
     
     def train(self,
-              group_data: Union[List[List[str]], Iterator[List[str]]],
+              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 = 1,
-              val_split: float = 0.8,
+              batch_size: int = 8,
+              accum_steps: int = 4,
               learning_rate: float = 2e-4,
               weight_decay: float = 1e-5,
               length_penalty: float = 0.01,
               single_vector_prob: float = 0.1,
               save_path: str = "thought_vectors.tar",
               spm_model_prefix: str = "spm",
-              vocab_size: int = 2000,
+              vocab_size: int = 4096,
               d_model: int = 256,
               encoder_nhead: int = 8,
               decoder_nhead: int = 8,
               encoder_layers: int = 2,
               decoder_layers: int = 2,
-              max_thoughts: int = 16,
+              max_thoughts: int = 8,
               dropout: float = 0.1,
-              max_len: int = 5000,
+              max_len: int = 1024,
               termination_threshold: float = 0.75,
               patience: int = 5) -> None:
         """Trains the thought vector model on grouped text data.
 
         Args:
-            group_data (Union[List[List[str]], Iterator[List[str]]]): Either a list of sentence groups or a generator yielding sentence groups for training.
-            model (str): Path to pre-existing model to load. If None, creates a new model. Defaults to None
+            group_data (Union[str, List[List[str]]]): Filepath to a CSV or list of sentence groups for training.
+            test_data (Union[str, List[List[str]]], optional): Filepath to a CSV or list of sentence groups for validation. If None, no validation is performed. Defaults to None.
+            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 1.
-            val_split (float, optional): Fraction of data for training. Defaults to 0.8.
+            batch_size (int, optional): Batch size. 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.
             length_penalty (float, optional): Penalty per additional thought vector. Defaults to 0.01.
             single_vector_prob (float, optional): Probability of forcing a single thought vector. Defaults to 0.1.
             save_path (str, optional): Path to save the model tar file. Defaults to "thought_vectors.tar".
             spm_model_prefix (str, optional): Prefix for SentencePiece files. Defaults to "spm".
-            vocab_size (int, optional): Vocabulary size for tokenization. Defaults to 2000.
+            vocab_size (int, optional): Vocabulary size for tokenization. Defaults to 4096.
             d_model (int, optional): Embedding dimensionality. Defaults to 256.
             encoder_nhead (int, optional): Number of encoder attention heads. Defaults to 8.
             decoder_nhead (int, optional): Number of decoder attention heads. Defaults to 8.
             encoder_layers (int, optional): Number of encoder layers. Defaults to 2.
             decoder_layers (int, optional): Number of decoder layers. Defaults to 2.
-            max_thoughts (int, optional): Maximum number of thought vectors. Defaults to 16.
+            max_thoughts (int, optional): Maximum number of thought vectors. Defaults to 8.
             dropout (float, optional): Dropout probability. Defaults to 0.1.
-            max_len (int, optional): Maximum sequence length for positional encoding. Defaults to 5000.
+            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.
         """
-        # Set model attributes
         self.d_model = d_model
         self.max_thoughts = max_thoughts
         self.encoder_nhead = encoder_nhead
@@ -354,152 +436,156 @@ class ThoughtVectors:
         self.termination_threshold = termination_threshold
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
-        # Load or create SentencePiece model
         self.sp = spm.SentencePieceProcessor()
+        train_lazy_data = LazyList(group_data)
+        test_lazy_data = LazyList(test_data) if test_data is not None else None
+        
         if model and os.path.exists(model):
-            # Load existing model checkpoint
             checkpoint = torch.load(model, map_location=self.device)
             self.sp.Load(checkpoint['spm_model_path'])
             self.vocab_size = self.sp.GetPieceSize()
-            
-            # Initialize and load encoder/decoder with saved parameters
             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
+                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
+                nhead=decoder_nhead, dropout=dropout, max_len=self.max_len
             ).to(self.device)
-            
             self.encoder.load_state_dict(checkpoint['encoder_state_dict'])
             self.decoder.load_state_dict(checkpoint['decoder_state_dict'])
             print(f"Loaded pre-existing model from {model}")
         else:
-            # Create new SentencePiece model
-            temp_dataset_path = None
-            if isinstance(group_data, (list, tuple)):
-                # Handle list input as before
-                all_sentences = [sentence for group in group_data for sentence in group]
-                with tempfile.NamedTemporaryFile(mode="w", encoding="utf8", delete=False) as temp_file:
-                    for line in all_sentences:
-                        temp_file.write(line.strip() + "\n")
-                    temp_dataset_path = temp_file.name
-            else:
-                # Handle generator input
-                with tempfile.NamedTemporaryFile(mode="w", encoding="utf8", delete=False) as temp_file:
-                    for group in group_data:
-                        for sentence in group:
-                            temp_file.write(sentence.strip() + "\n")
-                    temp_dataset_path = temp_file.name
+            with tempfile.NamedTemporaryFile(mode="w", encoding="utf8", delete=False) as temp_file:
+                for group in train_lazy_data:
+                    for sentence in group:
+                        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")
-            
             try:
                 spm.SentencePieceTrainer.Train(
                     input=temp_dataset_path,
                     model_prefix=temp_spm_prefix,
                     vocab_size=vocab_size,
                     model_type="bpe",
-                    pad_id=0,
-                    unk_id=1,
-                    bos_id=2,
-                    eos_id=3,
-                    max_sentence_length=8192,  # Added to handle long sequences
+                    pad_id=0, unk_id=1, bos_id=2, eos_id=3,
+                    max_sentence_length=self.max_len,
                     input_sentence_size=8_388_608,
-                    train_extremely_large_corpus=True  # Optimize for large datasets
+                    train_extremely_large_corpus=True
                 )
                 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:
-                if temp_dataset_path:
-                    os.remove(temp_dataset_path)
+                os.remove(temp_dataset_path)
                 shutil.rmtree(temp_dir, ignore_errors=True)
             
             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
+                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
+                nhead=decoder_nhead, dropout=dropout, max_len=self.max_len
             ).to(self.device)
         
-        # Rest of the training setup
-        dataset = GroupTextDataset(group_data, self.sp)
-        train_size = int(val_split * len(dataset))
-        val_size = len(dataset) - train_size
-        train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
-        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=group_collate_fn)
-        val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, collate_fn=group_collate_fn)
-        
         optimizer = optim.Adam(list(self.encoder.parameters()) + list(self.decoder.parameters()),
                                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()
         
         best_val_loss = float('inf')
         patience_counter = 0
         
-        # Training loop
+        print("Beginning training")
         for epoch in range(num_epochs):
             self.encoder.train()
             self.decoder.train()
             total_train_loss = 0.0
-            for group_batch in train_loader:
-                optimizer.zero_grad()
-                batch_loss = 0.0
-                for group_tensor in group_batch:
-                    if group_tensor.numel() == 0:
-                        continue
-                    group_tensor = group_tensor.to(self.device)
-                    force_single_vector = random.random() < single_vector_prob
-                    thought_vectors = self.encoder(group_tensor, force_single_vector=force_single_vector)
-                    
-                    decoder_input = group_tensor[:, :-1]
-                    target_tokens = group_tensor[:, 1:]
-                    output_logits = self.decoder(thought_vectors, decoder_input)
-                    
-                    loss = criterion(output_logits.reshape(-1, self.decoder.vocab_size), target_tokens.reshape(-1))
-                    if not force_single_vector:
-                        loss += length_penalty * thought_vectors.shape[1]
-                    batch_loss += loss
-                
-                if batch_loss > 0:
-                    batch_loss.backward()
-                    optimizer.step()
-                    total_train_loss += batch_loss.item()
+            optimizer.zero_grad()
             
-            avg_train_loss = total_train_loss / len(train_loader)
+            # 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
             
-            self.encoder.eval()
-            self.decoder.eval()
-            total_val_loss = 0.0
-            with torch.no_grad():
-                for group_batch in val_loader:
-                    batch_val_loss = 0.0
+            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
+                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
+                
+                batch_loss = 0.0
+                with autocast():
                     for group_tensor in group_batch:
-                        if group_tensor.numel() == 0:
-                            continue
-                        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.decoder.vocab_size),
+                        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])
+                            loss = criterion(output_logits.reshape(-1, self.vocab_size),
                                              group_tensor[:, 1:].reshape(-1))
-                        batch_val_loss += val_loss.item()
-                    
-                    total_val_loss += batch_val_loss
+                            if not force_single_vector:
+                                loss += length_penalty * thought_vectors.shape[1]
+                            batch_loss += loss / accum_steps
+                
+                scaler.scale(batch_loss).backward()
+                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}")
+            
+            avg_train_loss = total_train_loss / num_train_batches if num_train_batches > 0 else float('inf')
+            
+            # Validation with test_lazy_data
+            if test_lazy_data is not None:
+                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
+                
+                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
+                        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))
+                                batch_val_loss += val_loss.item()
+                        total_val_loss += batch_val_loss
+                
+                avg_val_loss = total_val_loss / num_val_batches if num_val_batches > 0 else float('inf')
+            else:
+                avg_val_loss = avg_train_loss
             
-            avg_val_loss = total_val_loss / len(val_loader)
-            print(f"Epoch {epoch + 1}/{num_epochs} - Train Loss: {avg_train_loss:.4f} - "
-                  f"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}")
             
             if avg_val_loss < best_val_loss:
                 best_val_loss = avg_val_loss
@@ -647,7 +733,8 @@ class ThoughtVectors:
                 'decoder_layers': self.decoder_layers,
                 'dropout': self.dropout,
                 'max_len': self.max_len,
-                'termination_threshold': self.termination_threshold
+                'termination_threshold': self.termination_threshold,
+                'spm_model_path': f"{spm_model_prefix}.model"  # Added to match loading logic
             }, "model.pth")
             tar.add("model.pth")
             tar.add(f"{spm_model_prefix}.model")
@@ -695,11 +782,11 @@ class ThoughtVectors:
         translator.encoder = ThoughtEncoder(
             vocab_size=translator.vocab_size, d_model=translator.d_model, max_thoughts=translator.max_thoughts,
             nhead=translator.encoder_nhead, num_layers=translator.encoder_layers, dropout=translator.dropout,
-            termination_threshold=translator.termination_threshold
+            termination_threshold=translator.termination_threshold, max_len=translator.max_len
         ).to(translator.device)
         translator.decoder = ThoughtDecoder(
             vocab_size=translator.vocab_size, d_model=translator.d_model, num_layers=translator.decoder_layers,
-            nhead=translator.decoder_nhead, dropout=translator.dropout
+            nhead=translator.decoder_nhead, dropout=translator.dropout, max_len=translator.max_len
         ).to(translator.device)
         
         translator.encoder.load_state_dict(checkpoint['encoder_state_dict'])
@@ -708,55 +795,36 @@ class ThoughtVectors:
         return translator
 
 
-# Example usage
 if __name__ == "__main__":
-    import csv
-    
-    csv.field_size_limit(8 ** 8)  # You can adjust this to a higher limit (e.g., 1MB)
-    
-    # Generator function to read data from a CSV and format it
-    def data_generator(csv_file_path):
-        with open(csv_file_path, mode='r', newline='', encoding='utf-8') as file:
-            reader = csv.reader(file)
-            for row in reader:
-                if row:
-                    yield [row[0]]  # assuming first collumn is sentences
-    
-    
-    data = data_generator("path/to/csv")
-    
-    # ThoughtVectors training and usage process
+    csv.field_size_limit(8 ** 8)
     tv = ThoughtVectors()
     tv.train(
-        group_data=data,
+        group_data="data.csv",
         num_epochs=1000,
-        batch_size=128,
-        val_split=0.8,
-        learning_rate=2e-4,
-        weight_decay=5e-5,
+        batch_size=64,
+        accum_steps=4,
+        learning_rate=1e-4,
+        weight_decay=2e-5,
         length_penalty=0.01,
-        single_vector_prob=0.2,
+        single_vector_prob=0.1,
         save_path="thought_vectors_prototype.tar",
         spm_model_prefix="spm",
         vocab_size=8192,
-        d_model=1024,
+        d_model=512,
         encoder_nhead=8,
         decoder_nhead=8,
-        encoder_layers=4,
-        decoder_layers=4,
+        encoder_layers=2,
+        decoder_layers=2,
         max_thoughts=32,
         dropout=0.1,
-        max_len=256,
-        termination_threshold=0.9,
-        patience=10
+        max_len=1024,
+        termination_threshold=0.8,
+        patience=5
     )
-    
-    # Load and test
     loaded_tv = ThoughtVectors.load("thought_vectors_prototype.tar")
-    loaded_thought_vectors = loaded_tv.encode("AI is smart")
-    generated_text_greedy = tv.decode(loaded_thought_vectors, temperature=0.7)
-    generated_text_beam = tv.decode(loaded_thought_vectors, beam_width=3)
-    
+    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)
     print(f"Greedy decoding: {generated_text_greedy}")
     print(f"Beam search decoding: {generated_text_beam}")