Upload ThoughtVectors.py

ThoughtVectors.py

@@ -0,0 +1,762 @@
+# thought_vectors.py
+import math
+import random
+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
+import os
+import tempfile
+import tarfile
+import shutil
+
+
+class PositionalEncoding(nn.Module):
+    """Adds positional information to token embeddings using sine and cosine functions.
+
+    This module can either precompute positional encodings up to a specified `max_len` or compute
+    them dynamically based on the input sequence length. If `max_len` is 0, encodings are computed
+    on-the-fly in the forward pass; otherwise, they are precomputed during initialization.
+
+    Args:
+        d_model (int): Dimensionality of the model embeddings.
+        dropout (float, optional): Dropout probability applied after adding encodings. Defaults to 0.1.
+        max_len (int, optional): Maximum sequence length for precomputed encodings; if 0, computes
+            dynamically. Defaults to 5000.
+
+    Attributes:
+        dropout (nn.Dropout): Dropout layer for regularization.
+        pe (torch.Tensor, optional): Precomputed positional encodings, shape (1, max_len, d_model),
+            present only if max_len > 0.
+    """
+    
+    def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):
+        super().__init__()
+        self.d_model = d_model
+        self.dropout = nn.Dropout(dropout)
+        self.max_len = max_len
+        
+        if max_len > 0:
+            # 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))
+            
+            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 = pe.unsqueeze(0)
+            self.register_buffer('pe', pe)
+        else:
+            self.pe = None
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Applies positional encodings to the input tensor.
+
+        Args:
+            x (torch.Tensor): Input embeddings, shape (batch_size, seq_len, d_model).
+
+        Returns:
+            torch.Tensor: Embeddings with positional encodings added and dropout applied.
+
+        Raises:
+            ValueError: If max_len > 0 and seq_len exceeds max_len.
+        """
+        seq_len = x.size(1)
+        
+        if self.max_len > 0:
+            if seq_len > self.max_len:
+                raise ValueError(f"Sequence length {seq_len} exceeds max_len {self.max_len}")
+            pe = self.pe[:, :seq_len]
+        else:
+            pe = torch.zeros(seq_len, self.d_model, device=x.device)
+            position = torch.arange(0, seq_len, dtype=torch.float, device=x.device).unsqueeze(1)
+            div_term = torch.exp(torch.arange(0, self.d_model - (self.d_model % 2), 2, dtype=torch.float,
+                                              device=x.device) * (-math.log(10000.0) / self.d_model))
+            
+            pe[:, 0::2] = torch.sin(position * div_term)
+            pe[:, 1::2] = torch.cos(position * div_term[:self.d_model // 2 + 1] if self.d_model % 2 else div_term)
+            
+            pe = pe.unsqueeze(0)
+        
+        x = x + pe
+        return self.dropout(x)
+
+
+class ThoughtEncoder(nn.Module):
+    """Encodes text into thought vectors, representing abstract ideas or summaries.
+
+    Args:
+        vocab_size (int): Size of the vocabulary for token embeddings.
+        d_model (int, optional): Dimensionality of the model embeddings. Defaults to 256.
+        max_thoughts (int, optional): Maximum number of thought vectors to generate. Defaults to 16.
+        nhead (int, optional): Number of attention heads in the transformer encoder. Defaults to 8.
+        num_layers (int, optional): Number of transformer encoder layers. Defaults to 2.
+        dropout (float, optional): Dropout probability. Defaults to 0.1.
+        pad_id (int, optional): ID for padding tokens. Defaults to 0.
+        termination_threshold (float, optional): Threshold for stopping thought vector generation.
+            Defaults to 0.75.
+    """
+    
+    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):
+        super().__init__()
+        assert nhead % 2 == 0, "nhead must be even for thought_attention"
+        self.d_model = d_model
+        self.max_thoughts = max_thoughts
+        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)
+        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)
+        self.thought_rnn = nn.GRU(d_model, d_model, batch_first=True)
+        self.fc_thought = nn.Linear(2 * d_model, d_model)
+        self.fc_terminate = nn.Linear(d_model, 1)
+    
+    def forward(self, input_tokens: torch.Tensor, force_single_vector: bool = False) -> torch.Tensor:
+        """Encodes input tokens into a sequence of thought vectors.
+
+        Args:
+            input_tokens (torch.Tensor): Token IDs, shape (batch_size, seq_len).
+            force_single_vector (bool, optional): If True, generates only one thought vector.
+                Defaults to False.
+
+        Returns:
+            torch.Tensor: Thought vectors, shape (batch_size, num_thoughts, d_model).
+        """
+        batch_size, seq_len = input_tokens.size()
+        max_thoughts = 1 if force_single_vector else self.max_thoughts
+        
+        src_key_padding_mask = (input_tokens == self.pad_id)
+        x = self.embedding(input_tokens)
+        x = self.positional_encoding(x)
+        encoded_text = self.text_encoder(x, src_key_padding_mask=src_key_padding_mask)
+        
+        valid_mask = (~src_key_padding_mask).unsqueeze(-1).float()
+        text_context = (encoded_text * valid_mask).sum(dim=1) / valid_mask.sum(dim=1).clamp(min=1.0)
+        
+        thought_vectors_list = []
+        thought_hidden = text_context.unsqueeze(0)
+        finished = torch.zeros(batch_size, dtype=torch.bool, device=input_tokens.device)
+        
+        for t in range(max_thoughts):
+            if t > 0:
+                prev_thoughts = torch.stack(thought_vectors_list, dim=1)
+                attn_output, _ = self.thought_attention(prev_thoughts, prev_thoughts, prev_thoughts)
+                prev_context = attn_output.mean(dim=1)
+            else:
+                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]
+            termination_logit = self.fc_terminate(thought_hidden.squeeze(0))
+            termination_score = torch.sigmoid(termination_logit).squeeze(-1)
+            
+            thought_vectors_list.append(next_thought)
+            finished = finished | (termination_score > self.termination_threshold)
+            if finished.all():
+                break
+        
+        return torch.stack(thought_vectors_list, dim=1)
+
+
+class ThoughtDecoder(nn.Module):
+    """Decodes thought vectors back into token sequences using a transformer decoder.
+
+    Args:
+        vocab_size (int): Size of the vocabulary for token embeddings.
+        d_model (int, optional): Dimensionality of the model embeddings. Defaults to 256.
+        num_layers (int, optional): Number of transformer decoder layers. Defaults to 2.
+        nhead (int, optional): Number of attention heads in the transformer decoder. Defaults to 8.
+        dropout (float, optional): Dropout probability. Defaults to 0.1.
+    """
+    
+    def __init__(self, vocab_size: int, d_model: int = 256, num_layers: int = 2,
+                 nhead: int = 8, dropout: float = 0.1):
+        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)
+        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)
+    
+    def generate_square_subsequent_mask(self, sz: int, device: torch.device) -> torch.Tensor:
+        """Generates a causal mask for autoregressive decoding.
+
+        Args:
+            sz (int): Size of the mask (sequence length).
+            device (torch.device): Device to create the mask on.
+
+        Returns:
+            torch.Tensor: Mask tensor, shape (sz, sz).
+        """
+        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
+    
+    def forward(self, thought_vectors: torch.Tensor, target_tokens: torch.Tensor) -> torch.Tensor:
+        """Decodes thought vectors into token logits.
+
+        Args:
+            thought_vectors (torch.Tensor): Thought vectors, shape (batch_size, num_thoughts, d_model).
+            target_tokens (torch.Tensor): Target token IDs, shape (batch_size, tgt_seq_len).
+
+        Returns:
+            torch.Tensor: Logits over vocabulary, shape (batch_size, tgt_seq_len, vocab_size).
+        """
+        tgt_emb = self.embedding(target_tokens)
+        tgt_emb = self.positional_encoding(tgt_emb)
+        tgt_mask = self.generate_square_subsequent_mask(tgt_emb.size(1), tgt_emb.device)
+        output = self.transformer_decoder(tgt=tgt_emb, memory=thought_vectors, tgt_mask=tgt_mask)
+        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.
+
+    Args:
+        batch (List[List[List[int]]]): Batch of tokenized groups.
+
+    Returns:
+        List[torch.Tensor]: Padded tensors for each group, shape (num_sentences, max_seq_len).
+    """
+    collated_groups = []
+    for group in batch:
+        if not group:
+            collated_groups.append(torch.tensor([]))
+            continue
+        group_tensors = [torch.tensor(seq, dtype=torch.long) for seq in group]
+        padded_group = nn.utils.rnn.pad_sequence(group_tensors, batch_first=True, padding_value=0)
+        collated_groups.append(padded_group)
+    return collated_groups
+
+
+class ThoughtVectors:
+    """Main class for a thought vector model, encoding text into latent vectors and decoding them back.
+
+    Attributes:
+        device (torch.device): Computation device (CUDA if available, else CPU).
+        encoder (ThoughtEncoder): Encoder module.
+        decoder (ThoughtDecoder): Decoder module.
+        sp (spm.SentencePieceProcessor): SentencePiece processor.
+    """
+    
+    def __init__(self):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.encoder = None
+        self.decoder = None
+        self.vocab_size = None
+        self.sp = None
+        self.d_model = None
+        self.max_thoughts = None
+        self.encoder_nhead = None
+        self.decoder_nhead = None
+        self.encoder_layers = None
+        self.decoder_layers = None
+        self.dropout = None
+        self.max_len = None
+        self.termination_threshold = None
+
+    
+    def train(self,
+              group_data: Union[List[List[str]], Iterator[List[str]]],
+              model: str = None,
+              num_epochs: int = 30,
+              batch_size: int = 1,
+              val_split: float = 0.8,
+              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,
+              d_model: int = 256,
+              encoder_nhead: int = 8,
+              decoder_nhead: int = 8,
+              encoder_layers: int = 2,
+              decoder_layers: int = 2,
+              max_thoughts: int = 16,
+              dropout: float = 0.1,
+              max_len: int = 5000,
+              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
+            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.
+            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.
+            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.
+            dropout (float, optional): Dropout probability. Defaults to 0.1.
+            max_len (int, optional): Maximum sequence length for positional encoding. Defaults to 5000.
+            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
+        self.decoder_nhead = decoder_nhead
+        self.encoder_layers = encoder_layers
+        self.decoder_layers = decoder_layers
+        self.dropout = dropout
+        self.max_len = max_len
+        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()
+        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
+            ).to(self.device)
+            self.decoder = ThoughtDecoder(
+                vocab_size=self.vocab_size, d_model=d_model, num_layers=decoder_layers,
+                nhead=decoder_nhead, dropout=dropout
+            ).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
+            
+            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
+                    input_sentence_size=8_388_608,
+                    train_extremely_large_corpus=True  # Optimize for large datasets
+                )
+                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)
+                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
+            ).to(self.device)
+            self.decoder = ThoughtDecoder(
+                vocab_size=self.vocab_size, d_model=d_model, num_layers=decoder_layers,
+                nhead=decoder_nhead, dropout=dropout
+            ).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())
+        
+        best_val_loss = float('inf')
+        patience_counter = 0
+        
+        # Training loop
+        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()
+            
+            avg_train_loss = total_train_loss / len(train_loader)
+            
+            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 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),
+                                             group_tensor[:, 1:].reshape(-1))
+                        batch_val_loss += val_loss.item()
+                    
+                    total_val_loss += batch_val_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}")
+            
+            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:
+                patience_counter += 1
+                if patience_counter >= patience:
+                    print(f"Early stopping triggered after {epoch + 1} epochs.")
+                    break
+            
+            scheduler.step()
+        
+        print(f"Best 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.
+
+        Args:
+            text (Union[str, List[str]]): Input text as a string or list of strings.
+            force_single_vector (bool, optional): If True, generates only one thought vector.
+                Defaults to False.
+
+        Returns:
+            torch.Tensor: Thought vectors, shape (batch_size, num_thoughts, d_model).
+        """
+        if self.encoder is None or self.sp is None:
+            raise RuntimeError("Model not trained or loaded.")
+        
+        if isinstance(text, str):
+            text = [text]
+        
+        tokenized = [[self.sp.bos_id()] + self.sp.EncodeAsIds(t) + [self.sp.eos_id()] for t in text]
+        input_tokens = nn.utils.rnn.pad_sequence(
+            [torch.tensor(seq, dtype=torch.long) for seq in tokenized],
+            batch_first=True, padding_value=self.sp.pad_id()
+        ).to(self.device)
+        
+        return self.encoder(input_tokens, force_single_vector)
+    
+    def decode(self, thought_vectors: torch.Tensor, max_length: int = 50, beam_width: Optional[int] = 5,
+               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
+                decoding. Defaults to 5.
+            temperature (float, optional): Temperature for softmax sampling in greedy decoding.
+                Defaults to 1.0.
+
+        Returns:
+            List[str]: Decoded text sequences.
+        """
+        if self.decoder is None or self.sp is None:
+            raise RuntimeError("Model not trained or loaded.")
+        
+        thought_vectors = thought_vectors.to(self.device)
+        batch_size = thought_vectors.size(0)
+        
+        if beam_width is not None and 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():
+                break
+        
+        return [self.sp.DecodeIds(seq.tolist()) for seq in target_tokens]
+    
+    def _beam_search_decode(self, thought_vectors: torch.Tensor, max_length: int, beam_width: int) -> List[str]:
+        """Performs beam search decoding of thought vectors.
+
+        Args:
+            thought_vectors (torch.Tensor): Thought vectors, shape (batch_size, num_thoughts, d_model).
+            max_length (int): Maximum length of generated sequences.
+            beam_width (int): Number of beams to maintain during search.
+
+        Returns:
+            List[str]: Decoded text sequences, one per batch item.
+        """
+        batch_size = thought_vectors.size(0)
+        start_tokens = torch.full((batch_size, 1), self.sp.bos_id(), dtype=torch.long, device=self.device)
+        beams = [[] for _ in range(batch_size)]
+        finished = [[] for _ in range(batch_size)]
+        
+        for i in range(batch_size):
+            beams[i].append((0.0, start_tokens[i].unsqueeze(0)))
+        
+        for _ in range(max_length - 1):
+            new_beams = [[] for _ in range(batch_size)]
+            for i in range(batch_size):
+                for score, seq in beams[i]:
+                    if seq[:, -1].item() == self.sp.eos_id():
+                        finished[i].append((score, seq))
+                        continue
+                    logits = self.decoder(thought_vectors[i:i + 1], seq)[:, -1, :]
+                    probs, next_tokens = logits.softmax(dim=-1).topk(beam_width, dim=-1)
+                    for p, t in zip(probs[0], next_tokens[0]):
+                        new_score = score - p.log().item()
+                        new_seq = torch.cat([seq, t.unsqueeze(0).unsqueeze(-1)], dim=1)
+                        new_beams[i].append((new_score, new_seq))
+                
+                beams[i] = sorted(new_beams[i], key=lambda x: x[0])[:beam_width]
+            
+            if all(len(finished[i]) >= beam_width for i in range(batch_size)):
+                break
+        
+        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>")
+        
+        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.
+        """
+        if self.encoder is None or self.decoder is None:
+            raise RuntimeError("No model to save.")
+        
+        if not os.path.exists(f"{spm_model_prefix}.model") or not os.path.exists(f"{spm_model_prefix}.vocab"):
+            raise FileNotFoundError(f"SentencePiece files ({spm_model_prefix}.model or .vocab) not found.")
+        
+        with tarfile.open(path, "w") as tar:
+            torch.save({
+                'encoder_state_dict': self.encoder.state_dict(),
+                'decoder_state_dict': self.decoder.state_dict(),
+                'vocab_size': self.vocab_size,
+                'd_model': self.d_model,
+                'max_thoughts': self.max_thoughts,
+                'encoder_nhead': self.encoder_nhead,
+                'decoder_nhead': self.decoder_nhead,
+                'encoder_layers': self.encoder_layers,
+                'decoder_layers': self.decoder_layers,
+                'dropout': self.dropout,
+                'max_len': self.max_len,
+                'termination_threshold': self.termination_threshold
+            }, "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':
+        """Loads a trained model from a tar file.
+
+        Args:
+            path (str): Path to the tar file containing the model.
+
+        Returns:
+            ThoughtVectors: Loaded model instance.
+        """
+        translator = cls()
+        try:
+            with tarfile.open(path, "r") as tar:
+                tar.extractall()
+            checkpoint = torch.load("model.pth")
+            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)
+        
+        translator.vocab_size = checkpoint['vocab_size']
+        translator.d_model = checkpoint['d_model']
+        translator.max_thoughts = checkpoint['max_thoughts']
+        translator.encoder_nhead = checkpoint['encoder_nhead']
+        translator.decoder_nhead = checkpoint['decoder_nhead']
+        translator.encoder_layers = checkpoint['encoder_layers']
+        translator.decoder_layers = checkpoint['decoder_layers']
+        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,
+            nhead=translator.encoder_nhead, num_layers=translator.encoder_layers, dropout=translator.dropout,
+            termination_threshold=translator.termination_threshold
+        ).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
+        ).to(translator.device)
+        
+        translator.encoder.load_state_dict(checkpoint['encoder_state_dict'])
+        translator.decoder.load_state_dict(checkpoint['decoder_state_dict'])
+        
+        return translator
+
+
+# Example usage with a shorter dataset
+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]]
+    
+    
+    data = data_generator("data.csv")
+    
+    # Define the ThoughtVectors training and usage process
+    tv = ThoughtVectors()
+    tv.train(
+        group_data=data,
+        num_epochs=1000,
+        batch_size=128,
+        val_split=0.8,
+        learning_rate=2e-4,
+        weight_decay=5e-5,
+        length_penalty=0.01,
+        single_vector_prob=0.2,
+        save_path="thought_vectors_prototype.tar",
+        spm_model_prefix="spm",
+        vocab_size=8192,
+        d_model=1024,
+        encoder_nhead=8,
+        decoder_nhead=8,
+        encoder_layers=4,
+        decoder_layers=4,
+        max_thoughts=32,
+        dropout=0.1,
+        max_len=256,
+        termination_threshold=0.9,
+        patience=10
+    )
+    
+    # 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)
+    
+    print(f"Greedy decoding: {generated_text_greedy}")
+    print(f"Beam search decoding: {generated_text_beam}")
+