import atexit import random import math import torch import torch.nn as nn import torch.optim as optim import sentencepiece as spm from torch.amp import autocast, GradScaler from typing import List, Union, Generator, Sequence import os import csv 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) 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[:self.d_model // 2 + (self.d_model % 2)]) 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, max_len: int = 5000): super().__init__() 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 self.termination_threshold = termination_threshold self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=pad_id) 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) 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) 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) thought_vectors_list.append(next_thought) finished = finished | (termination_score > self.termination_threshold) if finished.all() and not force_single_vector: 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, 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, 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) 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). """ 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 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) 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 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: 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: 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: return row except StopIteration: self._file_handle.seek(0) self._reader = csv.reader(self._file_handle) 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"], ] 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. 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 # 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 = 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, 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 = 4096, d_model: int = 256, encoder_nhead: int = 8, decoder_nhead: int = 8, encoder_layers: int = 2, decoder_layers: int = 2, max_thoughts: int = 8, dropout: float = 0.1, 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[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 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. 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 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 8. 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): BATCHES to wait for improvement before early stopping. Defaults to 5. """ 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") 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): print("loading model...") loaded_instance = self.load(model) self.__dict__.update(loaded_instance.__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: temp_file.write(sentence.strip() + "\n") temp_dataset_path = temp_file.name # 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)") 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=self.max_len, 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) # 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( 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) 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('cuda') best_val_loss = float('inf') patience_counter = 0 best_batch = 0 # Track best batch number print("Beginning training") try: for epoch in range(num_epochs): total_train_loss = 0.0 num_train_batches = (len(train_lazy_data) + batch_size - 1) // batch_size 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] for group in raw_batch ] group_batch = group_collate_fn(tokenized_batch) batch_loss = 0.0 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]) 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 (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() # 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") 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() 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 in ("best", "b"): break elif saving in ("current", "c", "curr"): self.save(save_path, spm_model_prefix) break 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"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: 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 0, 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 > 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 self.sp.DecodeIds( beams[i][0][1].squeeze(0).tolist())) return results 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.") 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, '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") @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() translator.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # load based on available device, not model train device 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'] 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, 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, max_len=translator.max_len ).to(translator.device) translator.encoder.load_state_dict(checkpoint['encoder_state_dict']) translator.decoder.load_state_dict(checkpoint['decoder_state_dict']) return translator 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="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=2e-4, weight_decay=1e-5, length_penalty=0.001, single_vector_prob=0.1, save_path="thought_vectors_prototype-0.2.0.tar", spm_model_prefix="spm", vocab_size=8192, d_model=512, encoder_nhead=8, decoder_nhead=8, encoder_layers=4, decoder_layers=4, max_thoughts=16, dropout=0.1, max_len=256, termination_threshold=0.8, patience=10 ) 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}")