robo-lib 0.0.10__tar.gz → 1.0.0__tar.gz

robo_lib-1.0.0/.gitignore

@@ -0,0 +1 @@
+__pycache__/

{robo_lib-0.0.10 → robo_lib-1.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: robo_lib
-Version: 0.0.10
+Version: 1.0.0
 Summary: A package to create, configure, and train transformer models.
 Project-URL: Homepage, https://github.com/hamburgerfish/robo_pack
 Project-URL: Issues, https://github.com/hamburgerfish/robo_pack/issues

{robo_lib-0.0.10 → robo_lib-1.0.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "robo_lib"
-version = "0.0.10"
+version = "1.0.0"
 authors = [
   { name="Erik Papp", email="erik3papp@gmail.com" },
 ]

{robo_lib-0.0.10 → robo_lib-1.0.0}/robo_lib/__init__.py

@@ -1,8 +1,7 @@
 from .components import TokenizerConstructor as TokenizerConstructor
 from .components import create_mask as create_mask
-from .components import pad as pad
-from .components import process_row as process_row
-from .components import scan_max_block_size as scan_max_block_size
+from .components import pre_process_data as pre_process_data
+from .components import safe_stack as safe_stack
 from .components import DataProcessor as DataProcessor
 from .components import get_valid_samples as get_valid_samples
 from .components import get_batch as get_batch

{robo_lib-0.0.10 → robo_lib-1.0.0}/robo_lib/components.py

@@ -6,6 +6,8 @@ import numpy as np
 import random
 import pickle
 import itertools
+from pathlib import Path
+import os
 
 class TokenizerConstructor:
     '''
@@ -30,6 +32,7 @@ class TokenizerConstructor:
                  tokenizer_type:str="BPE",
                  pre_tokenizers:list[str]|str=["Whitespace"],
                  normalizers:list[str]|str=["Lowercase", "NFD", "StripAccents", "Strip"],
+                 vocab:dict[str,int] = {},
                  special_tokens:list[str]|str=[],
                  unknown_token_string:str="<unk>",
                  start_token_string:str="<sos>",
@@ -42,25 +45,28 @@ class TokenizerConstructor:
 
         if isinstance(special_tokens, str):
             special_tokens = [special_tokens]
-        self.special_tokens = special_tokens + [token for token in [unknown_token_string, start_token_string, end_token_string, pad_token_string, new_line_token_string] if token not in special_tokens and token != None]
-        self.unknown_token = self.special_tokens.index(unknown_token_string) if unknown_token_string != None else None
-        self.start_token = self.special_tokens.index(start_token_string) if start_token_string != None else None
-        self.end_token = self.special_tokens.index(end_token_string) if end_token_string != None else None
-        self.pad_token = self.special_tokens.index(pad_token_string) if pad_token_string != None else None
-        self.new_line_token = self.special_tokens.index(new_line_token_string) if new_line_token_string != None else None
+        self.special_tokens = special_tokens + [token for token in [unknown_token_string, start_token_string, end_token_string, pad_token_string, new_line_token_string] if token not in special_tokens and token is not None]
+        self.unknown_token = self.special_tokens.index(unknown_token_string) if unknown_token_string is not None else None
+        self.start_token = self.special_tokens.index(start_token_string) if start_token_string is not None else None
+        self.end_token = self.special_tokens.index(end_token_string) if end_token_string is not None else None
+        self.pad_token = self.special_tokens.index(pad_token_string) if pad_token_string is not None else None
+        self.new_line_token = self.special_tokens.index(new_line_token_string) if new_line_token_string is not None else None
 
         if tokenizer_type == "BPE":
             self.tokenizer_type = tokenizers.Tokenizer(tokenizers.models.BPE(unk_token=unknown_token_string))
             self.trainer = tokenizers.trainers.BpeTrainer(special_tokens=self.special_tokens, min_frequency=min_frequency, vocab_size=vocab_size)
         elif tokenizer_type == "WordLevel":
-            self.tokenizer_type = tokenizers.Tokenizer(tokenizers.models.WordLevel(unk_token=unknown_token_string))
+            self.tokenizer_type = tokenizers.Tokenizer(tokenizers.models.WordLevel(vocab = vocab, unk_token=unknown_token_string))
             self.trainer = tokenizers.trainers.WordLevelTrainer(special_tokens=self.special_tokens, min_frequency=min_frequency, vocab_size=vocab_size)
         elif tokenizer_type == "WordPiece":
-            self.tokenizer_type = tokenizers.Tokenizer(tokenizers.models.WordPiece(unk_token=unknown_token_string))
+            self.tokenizer_type = tokenizers.Tokenizer(tokenizers.models.WordPiece(vocab = vocab, unk_token=unknown_token_string))
             self.trainer = tokenizers.trainers.WordPieceTrainer(special_tokens=self.special_tokens, min_frequency=min_frequency, vocab_size=vocab_size)
         elif tokenizer_type == "Unigram":
-            self.tokenizer_type = tokenizers.Tokenizer(tokenizers.models.Unigram(unk_token=unknown_token_string))
+            self.tokenizer_type = tokenizers.Tokenizer(tokenizers.models.Unigram())
             self.trainer = tokenizers.trainers.UnigramTrainer(special_tokens=self.special_tokens, min_frequency=min_frequency, vocab_size=vocab_size)
+        
+        if self.pad_token is not None:
+            self.tokenizer_type.enable_padding(pad_id=self.pad_token, pad_token=pad_token_string)
 
         if isinstance(pre_tokenizers, str):
             pre_tokenizers = [pre_tokenizers]
@@ -122,6 +128,13 @@ class TokenizerConstructor:
     def encode(self, inp:str) -> list[int]:
         return self.tokenizer_type.encode(inp).ids
     
+    def encode_batch(self, inp:list[str], max_length:int=None) -> list[list[int]]:
+        if max_length is not None:
+            self.tokenizer_type.enable_truncation(max_length=max_length)
+        out = [row.ids for row in self.tokenizer_type.encode_batch(inp)]
+        self.tokenizer_type.no_truncation()
+        return out
+    
     def decode(self, inp:list[int]) -> str:
         return self.tokenizer_type.decode(inp)
     
@@ -136,38 +149,35 @@ def create_mask(row:list, block_size:int) -> list[bool]:
     mask = [1]*len(row) + [0]*(block_size - len(row))
     return mask
 
-def pad(row:list, block_size:int, pad_token:int) -> list[int]:
-    '''
-
-    returns padded row. Row is padded until length block_size with specified pad_token value
-
-    '''
-    row.extend([pad_token]*(block_size - len(row)))
-    return row
-
-def process_row(row:str, tokenizer:TokenizerConstructor) -> list[int]:
+def pre_process_data(data:str, start_token_string:str, end_token_string:str) -> list[int]:
     '''
     
-    returns tokenized row using specified tokenizer, and adds the tokenizer's start and end tokens if they exist
+    returns string row with the tokenizer's start and end tokens if they exist
     
     '''
-    processed_row = tokenizer.encode(row)
-    if tokenizer.start_token != None:
-        processed_row.insert(0, tokenizer.start_token)
-    if tokenizer.end_token != None:
-        processed_row.append(tokenizer.end_token)
-    
-    return processed_row
+    if start_token_string is None and end_token_string is None:
+        return data
+    else:
+        for i in range(len(data)):
+            if start_token_string is not None:
+                data[i] = start_token_string + data[i]
+            if end_token_string is not None:
+                data[i] = data[i] + end_token_string
+            
+    return data
 
-def scan_max_block_size(data:list[str], tokenizer:TokenizerConstructor) -> int:
+def safe_stack(tensor_list:list[torch.tensor]) -> torch.tensor:
     '''
     
-    returns max_block_size of given list of strings by taking the length of the longest process_row(row) in data
+    torch stack with check to ensure tensors are valid in input list
+
+    returns torch.stack(out_list) for all valid torch tensors in tensor_list. raises error if no valid tensors
     
     '''
-    lengths = [len(process_row(p, tokenizer)) for p in data]
-    max_block_size_scanner = max(lengths)
-    return max_block_size_scanner
+    out_list = [row for row in tensor_list if isinstance(row, torch.Tensor)]
+    if len(out_list) == 0:
+        raise ValueError("no valid tensors in list.")
+    return torch.stack(out_list)
 
 
 class DataProcessor:
@@ -196,93 +206,55 @@ class DataProcessor:
         self.enc_tokenizer = enc_tokenizer
 
     def process_list(self,
-                     save_path:str,
                      dec_data:list[str]|str,
                      dec_max_block_size:int=None,
                      dec_create_masks:bool=True,
-                     dec_block_size_exceeded_policy:str=None,
                      enc_data:list[str]=None,
                      enc_max_block_size:int=None,
                      enc_create_masks:bool=True,
-                     enc_block_size_exceeded_policy:str=None
+                     save_path:str = "."
                      ) -> None:
 
         if isinstance(dec_data, str):
             dec_data = [dec_data]
         dec_data_length = len(dec_data)
-        save_path = save_path.replace(".pt", "")
-
-        if dec_max_block_size == None:
-            dec_max_block_size = scan_max_block_size(dec_data, self.dec_tokenizer)
 
-        if enc_data != None:
-            self.enc_tokenizer = self.dec_tokenizer if self.enc_tokenizer == None else self.enc_tokenizer
+        if enc_data is not None:
+            if self.enc_tokenizer is None:
+                self.enc_tokenizer = self.dec_tokenizer
 
             enc_data_length = len(enc_data)
             if dec_data_length != enc_data_length:
-                raise Exception(f"decoder and encoder lengths do not match. decoder_data_length is {dec_data_length}, encoder_data_length is {enc_data_length}")
+                raise Exception(f"decoder data and encoder data lengths do not match. decoder_data_length is {dec_data_length}, encoder_data_length is {enc_data_length}")
 
-            if enc_max_block_size == None:
-                enc_max_block_size = scan_max_block_size(enc_data, self.enc_tokenizer)
-            
-            enc_out_list = [[]]*enc_data_length
-            enc_mask_list = [[]]*enc_data_length if enc_create_masks else []
-        else:
-            enc_out_list = []
-            enc_mask_list = []
-
-        dec_out_list = [[]]*dec_data_length
-        dec_mask_list = [[]]*dec_data_length if dec_create_masks else []
-        for index in range(len(dec_out_list)):
-            dec_processed_item = process_row(dec_data[index], self.dec_tokenizer)
-            if dec_max_block_size != None and len(dec_processed_item) > dec_max_block_size:
-                if dec_block_size_exceeded_policy == "trim":
-                    dec_processed_item = dec_processed_item[:dec_max_block_size]
-                elif dec_block_size_exceeded_policy == "skip":
-                    continue
-                elif dec_block_size_exceeded_policy == None:
-                    raise Exception(f"encountered item in dec_data larger than maximum block size ({dec_max_block_size})")
-            if dec_create_masks:
-                dec_mask = create_mask(dec_processed_item, dec_max_block_size)
-            dec_processed_item = pad(dec_processed_item, dec_max_block_size, self.dec_tokenizer.pad_token)
-                
-            if enc_data != None:
-                enc_processed_item = process_row(enc_data[index], self.enc_tokenizer)
-                if enc_max_block_size != None and len(enc_processed_item) > enc_max_block_size:
-                    if enc_block_size_exceeded_policy == "trim":
-                        enc_processed_item = enc_processed_item[:enc_max_block_size]
-                    elif enc_block_size_exceeded_policy == "skip":
-                        continue
-                    elif enc_block_size_exceeded_policy == None:
-                        raise Exception(f"encountered item in enc_data larger than maximum block size ({enc_max_block_size})")
-                if enc_create_masks:
-                    enc_mask = create_mask(enc_processed_item, enc_max_block_size)
-                enc_processed_item = pad(enc_processed_item, enc_max_block_size, self.enc_tokenizer.pad_token)
-                    
-            dec_out_list[index] = torch.tensor(dec_processed_item, dtype=torch.long)
-            if dec_create_masks:
-                dec_mask_list[index] = torch.tensor(dec_mask, dtype=torch.bool)
-
-            if enc_data != None:
-                enc_out_list[index] = torch.tensor(enc_processed_item, dtype=torch.long)
-                if enc_create_masks:
-                    enc_mask_list[index] = torch.tensor(enc_mask, dtype=torch.bool)
-
-        dec_out_list = torch.stack([row for row in dec_out_list if row != []])
-        torch.save(dec_out_list, save_path + "_decoder_data.pt")
+        print("processing data")
+        dec_out_list = self.dec_tokenizer.encode_batch(dec_data, max_length=dec_max_block_size)
+        if dec_create_masks:
+            mask_tokenizer = TokenizerConstructor(min_frequency=1, tokenizer_type="WordLevel", vocab={str(self.dec_tokenizer.pad_token): 0, "<unk>": 1}, special_tokens=["<pad>", "<unk>"], unknown_token_string="<unk>", start_token_string=None, end_token_string=None, pad_token_string=None)
+            dec_mask_list = mask_tokenizer.encode_batch([str(i).replace("[", "").replace("]", "").replace(",", "") for i in dec_out_list])
+
+        if enc_data is not None:
+            enc_out_list = self.enc_tokenizer.encode_batch(enc_data, max_length=enc_max_block_size)
+            if enc_create_masks:
+                mask_tokenizer = TokenizerConstructor(min_frequency=1, tokenizer_type="WordLevel", vocab={str(self.enc_tokenizer.pad_token): 0, "<unk>": 1}, special_tokens=["<pad>", "<unk>"], unknown_token_string="<unk>", start_token_string=None, end_token_string=None, pad_token_string=None)
+                enc_mask_list = mask_tokenizer.encode_batch([str(i).replace("[", "").replace("]", "").replace(",", "") for i in enc_out_list])
+
+        dec_out_list = torch.tensor(dec_out_list, dtype=torch.long)
+        Path(save_path).mkdir(parents=True, exist_ok=True)
+        torch.save(dec_out_list, os.path.join(save_path, "decoder_data.pt"))
         if dec_create_masks:
-            dec_mask_list = torch.stack([row for row in dec_mask_list if row != []])
-            torch.save(dec_mask_list, save_path + "_decoder_mask_data.pt")
-        if enc_data != None:
-            enc_out_list = torch.stack([row for row in enc_out_list if row != []])
-            torch.save(enc_out_list, save_path + "_encoder_data.pt")
+            dec_mask_list = torch.tensor(dec_mask_list, dtype=torch.long)
+            torch.save(dec_mask_list, os.path.join(save_path, "decoder_mask_data.pt"))
+        if enc_data is not None:
+            enc_out_list = torch.tensor(enc_out_list, dtype=torch.long)
+            torch.save(enc_out_list, os.path.join(save_path, "encoder_data.pt"))
             if enc_create_masks:
-                enc_mask_list = torch.stack([row for row in enc_mask_list if row != []])
-                torch.save(enc_mask_list, save_path + "_encoder_mask_data.pt")
+                enc_mask_list = torch.tensor(enc_mask_list, dtype=torch.long)
+                torch.save(enc_mask_list, os.path.join(save_path, "encoder_mask_data.pt"))
 
 
-def get_valid_samples(random_samples:torch.tensor,
-                      masks:torch.tensor,
+def get_valid_samples(random_samples:torch.Tensor,
+                      masks:torch.Tensor,
                       block_size:int
                       ) -> list[int]:
     '''
@@ -294,9 +266,9 @@ def get_valid_samples(random_samples:torch.tensor,
     valid_samples = [0 if sum(masks[row_num]) <= block_size else random.randint(0, sum(masks[row_num]) - block_size) for row_num in random_samples]
     return valid_samples
                 
-def get_batch(data:torch.tensor,
-                random_samples:torch.tensor,
-                masks:torch.tensor=None,
+def get_batch(data:torch.Tensor,
+                random_samples:torch.Tensor,
+                masks:torch.Tensor=None,
                 block_size:int=None,
                 get_offset:bool=True
                 ) -> tuple[torch.tensor]:
@@ -308,53 +280,78 @@ def get_batch(data:torch.tensor,
     
     '''
     batch_size = len(random_samples)
-    if block_size != None and block_size != data.shape[1]:
+    if block_size is not None and block_size != data.shape[1]:
         if block_size >= data.shape[1]:
             raise Exception(f"specified block size ({block_size}) is larger than input tensor length ({data.shape[1]})")
 
-        if masks != None:
+        if masks is not None:
             random_point = get_valid_samples(random_samples, masks, block_size)
         else:
             random_point = torch.randint(data.shape[1] - block_size, (batch_size,))
         batch_in = torch.stack([data[random_samples[i]][random_point[i]:random_point[i]+block_size-int(get_offset)] for i in range(batch_size)])
-        masks_in = torch.stack([masks[random_samples[i]][random_point[i]:random_point[i]+block_size-int(get_offset)] for i in range(batch_size)]) if masks != None else None
+        masks_in = torch.stack([masks[random_samples[i]][random_point[i]:random_point[i]+block_size-int(get_offset)] for i in range(batch_size)]) if masks is not None else None
         batch_out = torch.stack([data[random_samples[i]][1+random_point[i]:random_point[i]+block_size] for i in range(batch_size)]) if get_offset else None
     else:
         block_size = data.shape[1]
         batch_in = torch.stack([data[row_num][:block_size-int(get_offset)] for row_num in random_samples])
-        masks_in = torch.stack([masks[row_num][:block_size-int(get_offset)] for row_num in random_samples]) if masks != None else None
+        masks_in = torch.stack([masks[row_num][:block_size-int(get_offset)] for row_num in random_samples]) if masks is not None else None
         batch_out = torch.stack([data[row_num][1:block_size] for row_num in random_samples]) if get_offset else None
 
     return batch_in, batch_out, masks_in
 
-def top_kp_filter(logits:torch.tensor,
-                  top_k:int,
-                  top_p:float=None
-                  ) -> torch.tensor:
+def top_kp_filter(logits: torch.Tensor,
+                  top_k: int = None,
+                  top_p: float = None
+                  ) -> torch.Tensor:
     '''
+    Returns predicted token by filtering output logits using top_k and/or top_p (nucleus) filtering.
     
-    returns predicted token by filtering output logits using top_k and top_p
-    
-    '''
-    if top_p != None:
-        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
-        cumulative_probs = torch.cumsum(sorted_logits, dim=-1)
+    Args:
+        logits: (batch_size, vocab_size) tensor of raw logits.
+        top_k: keep only top_k tokens with highest logits.
+        top_p: keep the smallest set of tokens with cumulative probability >= top_p.
         
-        filter = cumulative_probs > top_p
-        filter[..., 1:] = filter[..., :-1].clone()
-        filter[..., 0] = 0
-        indices_to_remove = filter.scatter(1, sorted_indices, filter)
-        logits[indices_to_remove] = float("-inf")
+    Returns:
+        selected: tensor of selected token indices (batch_size,)
+    '''
+    logits = logits.clone()  # avoid modifying input logits in-place
+
+    # Apply top-p filtering if specified
+    if top_p is not None:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+        probs = F.softmax(sorted_logits, dim=-1)
+        cumulative_probs = torch.cumsum(probs, dim=-1)
+
+        # Remove tokens with cumulative probability above threshold (except first token)
+        sorted_mask = cumulative_probs > top_p
+        sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
+        sorted_mask[..., 0] = False
+
+        # Mask tokens to remove by setting logits to -inf
+        indices_to_remove = sorted_mask.scatter(1, sorted_indices, sorted_mask)
+        logits[indices_to_remove] = float('-inf')
+
+    # Apply top-k filtering if specified
+    if top_k is not None:
+        top_k = min(top_k, logits.size(-1))  # safety check
+        topk_logits, topk_indices = torch.topk(logits, top_k, dim=-1)
+        topk_probs = F.softmax(topk_logits, dim=-1).cpu().numpy()
+
+        # For each batch, sample from top_k candidates
+        selected = []
+        for i in range(topk_probs.shape[0]):
+            candidate = np.random.choice(topk_indices[i].cpu().numpy(), 1, p=topk_probs[i])
+            selected.append(candidate[0])
+        selected = torch.tensor(selected, dtype=torch.long)
 
-    if top_k != None:
-        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
-
-        sorted_logits = F.softmax(sorted_logits[:, :top_k], dim=-1)
-        sorted_indices = sorted_indices[:, :top_k].detach().cpu()
-        sorted_logits = sorted_logits.detach().cpu().numpy()
-        sorted_logits[0][0] += 1 - sum(sorted_logits[0])
-
-        selected = torch.tensor(np.random.choice(sorted_indices[0], 1, p=sorted_logits[0]), dtype=torch.long)
+    else:
+        # If only top_p is specified, sample from entire filtered logits
+        probs = F.softmax(logits, dim=-1).cpu().numpy()
+        selected = []
+        for i in range(probs.shape[0]):
+            candidate = np.random.choice(len(probs[i]), 1, p=probs[i])
+            selected.append(candidate[0])
+        selected = torch.tensor(selected, dtype=torch.long)
 
     return selected
 
@@ -387,10 +384,10 @@ class SelfAttention(nn.Module):
         self.dropout = nn.Dropout(dropout)
 
     def forward(self,
-                k:torch.tensor,
-                q:torch.tensor,
-                v:torch.tensor,
-                mask:torch.tensor=None
+                k:torch.Tensor,
+                q:torch.Tensor,
+                v:torch.Tensor,
+                mask:torch.Tensor=None
                 ) -> torch.tensor:
         '''
         
@@ -406,7 +403,7 @@ class SelfAttention(nn.Module):
         wei = q @ k.transpose(-2,-1) * k.shape[-1]**-0.5
         if self.triangle_mask and self.block_size >= 0:
             wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))
-        if mask != None:
+        if mask is not None:
             wei = wei.masked_fill(mask.unsqueeze(1)==0, float("-inf"))
         wei = F.softmax(wei, dim=-1)
         wei = self.dropout(wei)
@@ -438,10 +435,10 @@ class MultiHeadAttention(nn.Module):
         self.dropout = nn.Dropout(dropout)
 
     def forward(self,
-                k:torch.tensor,
-                q:torch.tensor,
-                v:torch.tensor,
-                mask:torch.tensor=None
+                k:torch.Tensor,
+                q:torch.Tensor,
+                v:torch.Tensor,
+                mask:torch.Tensor=None
                 ) -> torch.tensor:
         '''
         
@@ -475,7 +472,7 @@ class FeedForward(nn.Module):
         )
 
     def forward(self,
-                x:torch.tensor
+                x:torch.Tensor
                 ) -> torch.tensor:
         return self.net(x)
     
@@ -500,8 +497,8 @@ class EncoderBlock(nn.Module):
         self.ln2 = nn.LayerNorm(n_embed)
 
     def forward(self,
-                x:torch.tensor,
-                mask:torch.tensor=None
+                x:torch.Tensor,
+                mask:torch.Tensor=None
                 ) -> tuple[torch.tensor]:
         att = self.sa(x, x, x, mask=mask)
         x = self.ln1(att + x)
@@ -541,15 +538,15 @@ class DecoderBlock(nn.Module):
             self.ca = None
 
     def forward(self,
-                x:torch.tensor,
-                enc_k:torch.tensor,
-                enc_v:torch.tensor,
+                x:torch.Tensor,
+                enc_k:torch.Tensor,
+                enc_v:torch.Tensor,
                 mask_out:bool=None,
-                mask_in:torch.tensor=None
+                mask_in:torch.Tensor=None
                 ) -> tuple[torch.tensor]:
         att = self.sa(x, x, x, mask=mask_out)
         x = self.ln1(att + x)
-        if self.ca != None:
+        if self.ca is not None:
             catt = self.ca(enc_k, x, enc_v, mask=mask_in)
             x = self.ln3(catt + x)
         ff = self.ffwd(x)
@@ -615,6 +612,7 @@ class RoboConstructor(nn.Module):
                  enc_vocab_size:int=None,
                  enc_block_size:int=None,
                  enc_expansion_factor:int=4,
+                 enc_positional_encoding:bool=True,
                  dropout:float=0.1,
                  device:str=None
                  ) -> None:
@@ -627,7 +625,7 @@ class RoboConstructor(nn.Module):
         self.dec_expansion_factor = dec_expansion_factor
         self.dropout = dropout
         
-        if device == None:
+        if device is None:
             self.device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
         else:
             self.device = device
@@ -635,6 +633,7 @@ class RoboConstructor(nn.Module):
         self.dec_positional_embedding_table = nn.Embedding(dec_block_size, n_embed)
 
         if enc_n_blocks != 0:
+            self.enc_positional_encoding = enc_positional_encoding
             self.enc_n_blocks = enc_n_blocks
             self.enc_n_head = enc_n_head
             self.enc_expansion_factor = enc_expansion_factor
@@ -642,7 +641,8 @@ class RoboConstructor(nn.Module):
             self.enc_block_size = enc_block_size
             self.cross_attention = True
             self.enc_token_embedding_table = nn.Embedding(enc_vocab_size, n_embed)
-            self.enc_positional_embedding_table = nn.Embedding(enc_block_size, n_embed)
+            if enc_positional_encoding:
+                self.enc_positional_embedding_table = nn.Embedding(enc_block_size, n_embed)
             self.encoder_blocks = MySequential(*[EncoderBlock(n_embed, enc_n_head, enc_expansion_factor, dropout=dropout) for _ in range(enc_n_blocks)])
         else:
             self.cross_attention = False
@@ -670,13 +670,13 @@ class RoboConstructor(nn.Module):
             torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
 
     def forward(self,
-                dec_in:torch.tensor,
-                dec_mask:torch.tensor=None,
-                enc_in:torch.tensor=None,
-                enc_mask:torch.tensor=None
+                dec_in:torch.Tensor,
+                dec_mask:torch.Tensor=None,
+                enc_in:torch.Tensor=None,
+                enc_mask:torch.Tensor=None
                 ) -> torch.tensor:
         _, dec_T = dec_in.shape
-        if enc_in != None:
+        if enc_in is not None:
             _, enc_T = enc_in.shape
 
         dec_tok_emb = self.dec_token_embedding_table(dec_in)
@@ -685,8 +685,11 @@ class RoboConstructor(nn.Module):
 
         if self.cross_attention:
             enc_tok_emb = self.enc_token_embedding_table(enc_in)
-            enc_pos_emb = self.enc_positional_embedding_table(torch.arange(enc_T, device=self.device))
-            enc_x = enc_tok_emb + enc_pos_emb
+            if self.enc_positional_encoding:
+                enc_pos_emb = self.enc_positional_embedding_table(torch.arange(enc_T, device=self.device))
+                enc_x = enc_tok_emb + enc_pos_emb
+            else:
+                enc_x = enc_tok_emb
 
             enc_out, enc_mask = self.encoder_blocks(enc_x, enc_mask)
         else:
@@ -712,13 +715,13 @@ class RoboConstructor(nn.Module):
 
         dec_train_batch_in, dec_train_batch_out, dec_train_masks_in = get_batch(dec_data, random_samples, masks=dec_masks, block_size=dec_block_size, get_offset=True)
         dec_train_batch_in = dec_train_batch_in.to(self.device)
-        dec_train_batch_out = dec_train_batch_out.to(self.device) if dec_train_batch_out != None else None
-        dec_train_masks_in = dec_train_masks_in.to(self.device) if dec_train_masks_in != None else None
+        dec_train_batch_out = dec_train_batch_out.to(self.device) if dec_train_batch_out is not None else None
+        dec_train_masks_in = dec_train_masks_in.to(self.device) if dec_train_masks_in is not None else None
 
         if self.cross_attention:
             enc_train_batch_in, _, enc_train_masks_in = get_batch(enc_data, random_samples, masks=enc_masks, block_size=enc_block_size, get_offset=False)
             enc_train_batch_in = enc_train_batch_in.to(self.device)
-            enc_train_masks_in = enc_train_masks_in.to(self.device) if enc_train_masks_in != None else None
+            enc_train_masks_in = enc_train_masks_in.to(self.device) if enc_train_masks_in is not None else None
         else:
             enc_train_batch_in = None
             enc_train_masks_in = None
@@ -730,14 +733,8 @@ class RoboConstructor(nn.Module):
               max_iters:int,
               eval_interval:int,
               batch_size:int,
-              dec_training_path:str,
-              dec_eval_path:str=None,
-              dec_training_masks_path:str=None,
-              dec_eval_masks_path:str=None,
-              enc_training_path:str=None,
-              enc_eval_path:str=None,
-              enc_training_masks_path:str=None,
-              enc_eval_masks_path:str=None,
+              training_dir_path:str,
+              eval_dir_path:str,
               eval_iters:int=3,
               learning_rate:float=1e-4,
               pad_token:int=None,
@@ -746,21 +743,36 @@ class RoboConstructor(nn.Module):
               label_smoothing:float=0.1
               ) -> None:
         
+        dec_training_path = os.path.join(training_dir_path, "decoder_data.pt")
         dec_training_data = torch.load(dec_training_path, weights_only=True)
-        dec_eval_data = torch.load(dec_eval_path, weights_only=True) if dec_eval_path != None else None
-        dec_training_masks_data = torch.load(dec_training_masks_path, weights_only=True) if dec_training_masks_path != None else None
-        dec_eval_masks_data = torch.load(dec_eval_masks_path, weights_only=True) if dec_eval_masks_path != None else None
-        enc_training_data = torch.load(enc_training_path, weights_only=True) if enc_training_path != None else None
-        enc_eval_data = torch.load(enc_eval_path, weights_only=True) if enc_eval_path != None else None
-        enc_training_masks_data = torch.load(enc_training_masks_path, weights_only=True) if enc_training_masks_path != None else None
-        enc_eval_masks_data = torch.load(enc_eval_masks_path, weights_only=True) if enc_eval_masks_path != None else None
-
-        if pad_token == None and dec_tokenizer != None:
+
+        dec_eval_path = os.path.join(eval_dir_path, "decoder_data.pt")
+        dec_eval_data = torch.load(dec_eval_path, weights_only=True) if os.path.isfile(dec_eval_path) else None
+
+        dec_training_masks_path = os.path.join(training_dir_path, "decoder_mask_data.pt")
+        dec_training_masks_data = torch.load(dec_training_masks_path, weights_only=True) if os.path.isfile(dec_training_masks_path) else None
+
+        dec_eval_masks_path = os.path.join(eval_dir_path, "decoder_mask_data.pt")
+        dec_eval_masks_data = torch.load(dec_eval_masks_path, weights_only=True) if os.path.isfile(dec_eval_masks_path) else None
+
+        enc_training_path = os.path.join(training_dir_path, "encoder_data.pt")
+        enc_training_data = torch.load(enc_training_path, weights_only=True) if os.path.isfile(enc_training_path) else None
+
+        enc_eval_path = os.path.join(eval_dir_path, "encoder_data.pt")
+        enc_eval_data = torch.load(enc_eval_path, weights_only=True) if os.path.isfile(enc_eval_path) else None
+
+        enc_training_masks_path = os.path.join(training_dir_path, "encoder_mask_data.pt")
+        enc_training_masks_data = torch.load(enc_training_masks_path, weights_only=True) if os.path.isfile(enc_training_masks_path) else None
+
+        enc_eval_masks_path = os.path.join(eval_dir_path, "encoder_mask_data.pt")
+        enc_eval_masks_data = torch.load(enc_eval_masks_path, weights_only=True) if os.path.isfile(enc_eval_masks_path) else None
+
+        if pad_token is None and dec_tokenizer is not None:
             pad_token = dec_tokenizer.pad_token
 
         self.to(self.device)
 
-        if pad_token != None:
+        if pad_token is not None:
             loss_fn = nn.CrossEntropyLoss(ignore_index=pad_token, label_smoothing=label_smoothing).to(self.device)
         else:
             loss_fn = nn.CrossEntropyLoss(label_smoothing=label_smoothing).to(self.device)
@@ -776,7 +788,7 @@ class RoboConstructor(nn.Module):
                 proj_output = self.forward(dec_x, dec_mask, enc_x, enc_mask)
                 losses[k] = loss_fn(proj_output.view(-1, self.dec_vocab_size), dec_y.view(-1))
             out["train"] = losses.mean()
-            if dec_eval_data != None:
+            if dec_eval_data is not None:
                 for k in range(eval_iters):
                     dec_x, dec_y, dec_mask, enc_x, enc_mask = self.prep_data(batch_size, dec_eval_data, dec_masks=dec_eval_masks_data, dec_block_size=self.dec_block_size, enc_data=enc_eval_data, enc_masks=enc_eval_masks_data, enc_block_size=self.enc_block_size)
                     proj_output = self.forward(dec_x, dec_mask, enc_x, enc_mask)
@@ -792,7 +804,7 @@ class RoboConstructor(nn.Module):
             if iter % eval_interval == 0 or iter == max_iters-1:
                 losses = estimate_loss()
                 print(f"step {iter}: train loss {losses['train']:.4f}, eval loss {losses['eval']:.4f}")
-                if save_path != None:
+                if save_path is not None:
                     save_component(self, save_path=save_path)
 
             dec_x, dec_y, dec_mask, enc_x, enc_mask = self.prep_data(batch_size, dec_training_data, dec_masks=dec_training_masks_data, dec_block_size=self.dec_block_size, enc_data=enc_training_data, enc_masks=enc_training_masks_data, enc_block_size=self.enc_block_size)
@@ -819,25 +831,25 @@ class RoboConstructor(nn.Module):
                 top_k:int=None,
                 top_p:float=None
                 ) -> list[int]|str:
-        max_new_tokens = self.dec_block_size if max_new_tokens == None else max_new_tokens
+        max_new_tokens = self.dec_block_size if max_new_tokens is None else max_new_tokens
 
         if self.cross_attention:
-            if enc_tokenizer != None:
-                if enc_start_token == None:
+            if enc_tokenizer is not None:
+                if enc_start_token is None:
                     enc_start_token = enc_tokenizer.start_token
-                if enc_end_token == None:
+                if enc_end_token is None:
                     enc_end_token = enc_tokenizer.end_token
                 if isinstance(inputs, str):
                     inputs = enc_tokenizer.encode(inputs)
 
-        if dec_tokenizer != None:
-            if dec_start_token == None:
+        if dec_tokenizer is not None:
+            if dec_start_token is None:
                 dec_start_token = dec_tokenizer.start_token
-            if dec_end_token == None:
+            if dec_end_token is None:
                 dec_end_token = dec_tokenizer.end_token
-            if new_line_token == None:
+            if new_line_token is None:
                 new_line_token = dec_tokenizer.new_line_token
-        if self.cross_attention == False and isinstance(inputs, str):
+        if not self.cross_attention and isinstance(inputs, str):
             inputs = dec_tokenizer.encode(inputs)
 
 
@@ -846,7 +858,7 @@ class RoboConstructor(nn.Module):
             idx = torch.tensor([[dec_start_token]], dtype=torch.long, device=self.device)
         else:
             enc_input = None
-            if separator_token != None:
+            if separator_token is not None:
                 idx = torch.tensor([[dec_start_token] + inputs + [separator_token]], dtype=torch.long, device=self.device)
             else:
                 idx = torch.tensor([[dec_start_token] + inputs], dtype=torch.long, device=self.device)
@@ -860,7 +872,7 @@ class RoboConstructor(nn.Module):
             logits = proj_output[:, -1, :]
             probabilities = F.log_softmax(logits/temperature, dim=-1)
 
-            if top_k == None and top_p == None:
+            if top_k is None and top_p is None:
                 idx_next = torch.max(probabilities, dim=-1).indices.unsqueeze(0)
             else:
                 idx_next = top_kp_filter(probabilities, top_k=top_k, top_p=top_p).unsqueeze(0).to(self.device)
@@ -868,10 +880,10 @@ class RoboConstructor(nn.Module):
             if idx_next[0] == dec_end_token:
                 break
         
-        if dec_tokenizer == None:
+        if dec_tokenizer is None:
             return idx[0].tolist()
         else:
-            if new_line_token != None:
+            if new_line_token is not None:
                 return "\n".join([dec_tokenizer.decode(list(y)) for x, y in itertools.groupby(idx[0].tolist(), lambda z: z == 0) if not x])
             else:
                 return dec_tokenizer.decode(idx[0].tolist())

robo_lib-1.0.0/tests/test_data_processor.py

@@ -0,0 +1,82 @@
+import os
+import shutil
+import torch
+import pytest
+from robo_lib.components import DataProcessor, TokenizerConstructor
+
+
+@pytest.fixture
+def temp_save_path():
+    path = "temp_test_dir"
+    yield path
+    if os.path.exists(path):
+        shutil.rmtree(path)
+
+
+@pytest.fixture
+def dummy_tokenizer():
+    tokenizer = TokenizerConstructor(
+        tokenizer_type="WordLevel",
+        pre_tokenizers="Whitespace",
+        special_tokens=["<pad>", "<unk>"],
+        vocab={"hello": 0, "world": 1, "<pad>": 2, "<unk>": 3},
+        unknown_token_string="<unk>",
+        pad_token_string="<pad>",
+        start_token_string=None,
+        end_token_string=None,
+        new_line_token_string=None
+    )
+    # Faking "training" so encode works with the vocab
+    tokenizer.tokenizer_type.add_tokens(["hello", "world"])
+    return tokenizer
+
+
+def test_data_processor_initialization(dummy_tokenizer):
+    processor = DataProcessor(dec_tokenizer=dummy_tokenizer)
+    assert processor.dec_tokenizer is dummy_tokenizer
+    assert processor.enc_tokenizer is None
+
+
+def test_process_list_decoder_only(dummy_tokenizer, temp_save_path):
+    processor = DataProcessor(dec_tokenizer=dummy_tokenizer)
+    data = ["hello world", "world hello"]
+    
+    processor.process_list(
+        dec_data=data,
+        dec_max_block_size=10,
+        save_path=temp_save_path
+    )
+
+    assert os.path.exists(os.path.join(temp_save_path, "decoder_data.pt"))
+    assert os.path.exists(os.path.join(temp_save_path, "decoder_mask_data.pt"))
+
+    tensor = torch.load(os.path.join(temp_save_path, "decoder_data.pt"))
+    assert isinstance(tensor, torch.Tensor)
+    assert tensor.shape[0] == len(data)
+
+
+def test_process_list_encoder_decoder(dummy_tokenizer, temp_save_path):
+    processor = DataProcessor(dec_tokenizer=dummy_tokenizer, enc_tokenizer=dummy_tokenizer)
+    data = ["hello world", "world hello"]
+
+    processor.process_list(
+        dec_data=data,
+        enc_data=data,
+        dec_max_block_size=10,
+        enc_max_block_size=10,
+        save_path=temp_save_path
+    )
+
+    assert os.path.exists(os.path.join(temp_save_path, "decoder_data.pt"))
+    assert os.path.exists(os.path.join(temp_save_path, "encoder_data.pt"))
+    assert os.path.exists(os.path.join(temp_save_path, "decoder_mask_data.pt"))
+    assert os.path.exists(os.path.join(temp_save_path, "encoder_mask_data.pt"))
+
+
+def test_process_list_mismatched_lengths_raises(dummy_tokenizer):
+    processor = DataProcessor(dec_tokenizer=dummy_tokenizer)
+    dec_data = ["hello world"]
+    enc_data = ["world hello", "extra row"]
+
+    with pytest.raises(Exception, match="decoder data and encoder data lengths do not match"):
+        processor.process_list(dec_data=dec_data, enc_data=enc_data)

robo_lib-1.0.0/tests/test_functions.py

@@ -0,0 +1,176 @@
+import pytest
+import torch
+import numpy as np
+import random
+from robo_lib import create_mask, pre_process_data, safe_stack, get_valid_samples, get_batch, top_kp_filter
+
+def test_create_mask_basic():
+    row = [1, 2, 3]
+    block_size = 5
+    expected = [1, 1, 1, 0, 0]
+    assert create_mask(row, block_size) == expected
+
+def test_create_mask_equal_length():
+    row = [1, 2, 3, 4]
+    block_size = 4
+    expected = [1, 1, 1, 1]
+    assert create_mask(row, block_size) == expected
+
+def test_create_mask_empty_row():
+    row = []
+    block_size = 3
+    expected = [0, 0, 0]
+    assert create_mask(row, block_size) == expected
+
+def test_pre_process_data_none_tokens():
+    data = ["hello", "world"]
+    start_token = None
+    end_token = None
+    # Should return the input unchanged
+    assert pre_process_data(data.copy(), start_token, end_token) == data
+
+def test_pre_process_data_start_token_only():
+    data = ["hello", "world"]
+    start_token = "<s>"
+    end_token = None
+    expected = ["<s>hello", "<s>world"]
+    assert pre_process_data(data.copy(), start_token, end_token) == expected
+
+def test_pre_process_data_end_token_only():
+    data = ["hello", "world"]
+    start_token = None
+    end_token = "</s>"
+    expected = ["hello</s>", "world</s>"]
+    assert pre_process_data(data.copy(), start_token, end_token) == expected
+
+def test_pre_process_data_both_tokens():
+    data = ["hello", "world"]
+    start_token = "<s>"
+    end_token = "</s>"
+    expected = ["<s>hello</s>", "<s>world</s>"]
+    assert pre_process_data(data.copy(), start_token, end_token) == expected
+
+def test_safe_stack_valid_tensors():
+    t1 = torch.tensor([1, 2])
+    t2 = torch.tensor([3, 4])
+    tensor_list = [t1, t2]
+    stacked = safe_stack(tensor_list)
+    assert isinstance(stacked, torch.Tensor)
+    assert stacked.shape == (2, 2)
+
+def test_safe_stack_ignore_non_tensors():
+    t1 = torch.tensor([1, 2])
+    not_tensor = [1, 2, 3]
+    tensor_list = [t1, not_tensor]
+    stacked = safe_stack(tensor_list)
+    assert stacked.shape == (1, 2)
+
+def test_safe_stack_raises_for_empty():
+    with pytest.raises(ValueError):
+        safe_stack(["not a tensor", 123, None])
+
+
+# For reproducibility
+random.seed(0)
+torch.manual_seed(0)
+np.random.seed(0)
+
+def test_get_valid_samples_all_masked_less_than_block():
+    masks = torch.tensor([[1, 0, 0], [1, 1, 0]])
+    random_samples = torch.tensor([0, 1])
+    block_size = 2
+    result = get_valid_samples(random_samples, masks, block_size)
+    # For first row sum(masks) = 1 <= block_size => should return 0
+    # For second row sum(masks) = 2 <= block_size => 0
+    assert result == [0, 0]
+
+def test_get_valid_samples_some_greater_than_block():
+    masks = torch.tensor([[1, 1, 1], [1, 1, 0]])
+    random_samples = torch.tensor([0, 1])
+    block_size = 2
+    result = get_valid_samples(random_samples, masks, block_size)
+    # first sum = 3 > 2, so random index in [0, 1]
+    # second sum = 2 <= 2, so 0
+    assert result[1] == 0
+    assert 0 <= result[0] <= 1
+
+def test_get_batch_no_masks_get_offset_true():
+    data = torch.arange(30).view(5, 6)  # 5 rows, 6 cols
+    random_samples = torch.tensor([0, 1, 2])
+    block_size = 4
+    batch_in, batch_out, masks_in = get_batch(data, random_samples, masks=None, block_size=block_size, get_offset=True)
+    assert batch_in.shape == (3, block_size-1)
+    assert batch_out.shape == (3, block_size-1)
+    assert masks_in is None
+
+def test_get_batch_with_masks_get_offset_false():
+    data = torch.arange(30).view(5, 6)
+    masks = torch.ones_like(data)
+    random_samples = torch.tensor([0, 1])
+    block_size = 5
+    batch_in, batch_out, masks_in = get_batch(data, random_samples, masks=masks, block_size=block_size, get_offset=False)
+    assert batch_in.shape == (2, block_size)
+    assert batch_out is None
+    assert masks_in.shape == (2, block_size)
+
+def test_get_batch_block_size_larger_than_data_length_raises():
+    data = torch.arange(20).view(4, 5)
+    random_samples = torch.tensor([0])
+    block_size = 6
+    with pytest.raises(Exception):
+        get_batch(data, random_samples, block_size=block_size)
+
+
+def test_top_kp_filter_top_k_only():
+    # Create dummy logits batch (2 samples, vocab size 5)
+    logits = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0],
+                           [5.0, 4.0, 3.0, 2.0, 1.0]])
+    top_k = 3
+    selected = top_kp_filter(logits, top_k=top_k, top_p=None)
+
+    assert selected.shape == (2,)
+    # Selected indices must be in top_k tokens
+    for i, sel in enumerate(selected):
+        topk_indices = torch.topk(logits[i], top_k).indices.tolist()
+        assert sel.item() in topk_indices
+
+def test_top_kp_filter_top_p_only():
+    # Dummy logits with clear probabilities
+    logits = torch.tensor([[0.1, 0.2, 0.3, 0.4],
+                           [0.4, 0.3, 0.2, 0.1]])
+    top_p = 0.7
+    selected = top_kp_filter(logits, top_k=None, top_p=top_p)
+
+    assert selected.shape == (2,)
+    # Selected indices must be in vocab range
+    for sel in selected:
+        assert 0 <= sel.item() < logits.shape[1]
+
+def test_top_kp_filter_top_k_and_top_p():
+    logits = torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5],
+                           [0.5, 0.4, 0.3, 0.2, 0.1]])
+    top_k = 2
+    top_p = 0.6
+    selected = top_kp_filter(logits, top_k=top_k, top_p=top_p)
+
+    assert selected.shape == (2,)
+    for i, sel in enumerate(selected):
+        # With both filters, selected index should be in top_k indices
+        topk_indices = torch.topk(logits[i], top_k).indices.tolist()
+        assert sel.item() in topk_indices
+
+def test_top_kp_filter_no_filter():
+    logits = torch.tensor([[0.1, 0.2, 0.3],
+                           [0.3, 0.2, 0.1]])
+    selected = top_kp_filter(logits, top_k=None, top_p=None)
+
+    assert selected.shape == (2,)
+    for sel in selected:
+        assert 0 <= sel.item() < logits.shape[1]
+
+def test_top_kp_filter_empty_logits():
+    # Edge case: logits empty or zero size
+    logits = torch.empty((0, 0))
+    with pytest.raises(IndexError):
+        _ = top_kp_filter(logits, top_k=1, top_p=0.5)
+

robo_lib-1.0.0/tests/test_robo_constructor.py

@@ -0,0 +1,130 @@
+import pytest
+import torch
+import tempfile
+import os
+from types import SimpleNamespace
+from unittest.mock import patch, MagicMock
+from robo_lib import RoboConstructor, save_component, load_component
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+# ---------- FIXTURES AND MOCKS ----------
+
+@pytest.fixture
+def mock_encoder_block():
+    return MagicMock()
+
+@pytest.fixture
+def mock_decoder_block():
+    return MagicMock()
+
+@pytest.fixture
+def mock_my_sequential():
+    class DummySequential(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+        def forward(self, *args, **kwargs):
+            return args[0], None, None, None, None
+    return DummySequential
+
+@pytest.fixture
+def dummy_tokenizer():
+    return SimpleNamespace(
+        start_token=1,
+        end_token=2,
+        pad_token=0,
+        new_line_token=3,
+        encode=lambda s: [4, 5, 6],
+        decode=lambda tokens: "decoded"
+    )
+
+@pytest.fixture
+def dummy_data():
+    return torch.randint(0, 10, (8, 32)).to(device)  # 8 samples of 32 tokens
+
+@pytest.fixture
+def robo_decoder_only(mock_my_sequential):
+    with patch("robo_lib.DecoderBlock", return_value=MagicMock()), \
+         patch("robo_lib.MySequential", mock_my_sequential):
+        return RoboConstructor(
+            n_embed=16,
+            dec_n_blocks=2,
+            dec_n_head=2,
+            dec_vocab_size=50,
+            dec_block_size=32
+        ).to(device)
+
+@pytest.fixture
+def robo_enc_dec(mock_my_sequential):
+    with patch("robo_lib.DecoderBlock", return_value=MagicMock()), \
+         patch("robo_lib.EncoderBlock", return_value=MagicMock()), \
+         patch("robo_lib.MySequential", mock_my_sequential):
+        return RoboConstructor(
+            n_embed=16,
+            dec_n_blocks=2,
+            dec_n_head=2,
+            dec_vocab_size=50,
+            dec_block_size=32,
+            enc_n_blocks=2,
+            enc_n_head=2,
+            enc_vocab_size=50,
+            enc_block_size=32
+        ).to(device)
+
+# ---------- TESTS ----------
+
+def test_decoder_only_init(robo_decoder_only):
+    assert not robo_decoder_only.cross_attention
+    assert robo_decoder_only.decoder_blocks is not None
+    assert robo_decoder_only.encoder_blocks is None
+
+def test_encoder_decoder_init(robo_enc_dec):
+    assert robo_enc_dec.cross_attention
+    assert robo_enc_dec.encoder_blocks is not None
+
+def test_forward_decoder_only(robo_decoder_only):
+    input_tensor = torch.randint(0, 50, (2, 32)).to(device)
+    output = robo_decoder_only(dec_in=input_tensor)
+    assert output.shape[:2] == (2, 32)
+
+def test_forward_encoder_decoder(robo_enc_dec):
+    dec_input = torch.randint(0, 50, (2, 32)).to(device)
+    enc_input = torch.randint(0, 50, (2, 32)).to(device)
+    output = robo_enc_dec(dec_in=dec_input, enc_in=enc_input)
+    assert output.shape[:2] == (2, 32)
+
+@patch("robo_lib.get_batch")
+def test_prep_data_decoder_only(mock_get_batch, robo_decoder_only, dummy_data):
+    mock_get_batch.return_value = (dummy_data[:2], dummy_data[:2], dummy_data[:2])
+    out = robo_decoder_only.prep_data(batch_size=2, dec_data=dummy_data, dec_block_size=32)
+    assert len(out) == 5
+    assert out[0].shape[0] == 2
+
+@patch("robo_lib.get_batch")
+def test_prep_data_encoder_decoder(mock_get_batch, robo_enc_dec, dummy_data):
+    mock_get_batch.side_effect = [
+        (dummy_data[:2], dummy_data[:2], dummy_data[:2]),  # decoder
+        (dummy_data[:2], None, dummy_data[:2])             # encoder
+    ]
+    out = robo_enc_dec.prep_data(batch_size=2, dec_data=dummy_data, dec_block_size=32, enc_data=dummy_data, enc_block_size=32)
+    assert len(out) == 5
+    assert out[3].shape[0] == 2  # encoder input
+
+@patch("robo_lib.top_kp_filter", return_value=torch.tensor([2]))
+def test_generate_decoder_only(mock_top_kp, robo_decoder_only, dummy_tokenizer):
+    out = robo_decoder_only.generate(inputs="hello", dec_tokenizer=dummy_tokenizer, max_new_tokens=3, dec_start_token=1, dec_end_token=2)
+    assert isinstance(out, str)
+
+@patch("robo_lib.top_kp_filter", return_value=torch.tensor([2]))
+def test_generate_encoder_decoder(mock_top_kp, robo_enc_dec, dummy_tokenizer):
+    out = robo_enc_dec.generate(inputs="hello", enc_tokenizer=dummy_tokenizer, dec_tokenizer=dummy_tokenizer,
+                                max_new_tokens=3, enc_start_token=1, enc_end_token=2, dec_start_token=1, dec_end_token=2)
+    assert isinstance(out, str)
+
+def test_save_and_load_component(robo_decoder_only):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        path = os.path.join(tmpdir, "test_model")
+        save_component(robo_decoder_only, path)
+        loaded = load_component(path)
+        assert isinstance(loaded, RoboConstructor)

robo_lib-1.0.0/tests/test_tokenizer_constructor.py

@@ -0,0 +1,89 @@
+import pytest
+import os
+from tempfile import NamedTemporaryFile
+from robo_lib import TokenizerConstructor
+
+
+@pytest.fixture
+def training_file():
+    with NamedTemporaryFile(mode="w+", delete=False) as f:
+        f.write("Hello world\nThis is a test\nTokenizer test\n")
+        f.flush()
+        yield f.name
+    os.remove(f.name)
+
+
+def test_tokenizer_creation():
+    tokenizer = TokenizerConstructor(
+        tokenizer_type="BPE",
+        pre_tokenizers="Whitespace",
+        normalizers=["Lowercase"],
+        special_tokens=["<unk>", "<pad>"],
+        vocab_size=100
+    )
+    assert tokenizer is not None
+    assert "<unk>" in tokenizer.special_tokens
+    assert tokenizer.vocab_size is None  # Untrained tokenizer should have vocab_size None
+
+
+def test_tokenizer_train(training_file):
+    tokenizer = TokenizerConstructor(
+        tokenizer_type="WordLevel",
+        pre_tokenizers="Whitespace",
+        normalizers=["Lowercase"],
+        special_tokens=["<unk>", "<pad>"],
+        vocab_size=50
+    )
+    tokenizer.train(training_file)
+    assert tokenizer.vocab_size is not None
+    assert tokenizer.vocab_size > 0
+
+
+def test_tokenizer_encode_decode(training_file):
+    tokenizer = TokenizerConstructor(
+        tokenizer_type="BPE",
+        pre_tokenizers="Whitespace",
+        normalizers=["Lowercase"],
+        special_tokens=["<unk>", "<pad>"],
+        vocab_size=50
+    )
+    tokenizer.train(training_file)
+    encoded = tokenizer.encode("This is a test")
+    assert isinstance(encoded, list)
+    assert all(isinstance(i, int) for i in encoded)
+
+    decoded = tokenizer.decode(encoded)
+    assert isinstance(decoded, str)
+    assert len(decoded) > 0
+
+
+def test_tokenizer_encode_batch(training_file):
+    tokenizer = TokenizerConstructor(
+        tokenizer_type="BPE",
+        pre_tokenizers="Whitespace",
+        normalizers=["Lowercase"],
+        special_tokens=["<unk>", "<pad>"],
+        vocab_size=50
+    )
+    tokenizer.train(training_file)
+    batch = ["This is a test", "Hello world"]
+    encoded_batch = tokenizer.encode_batch(batch)
+    assert isinstance(encoded_batch, list)
+    assert len(encoded_batch) == len(batch)
+    assert all(isinstance(seq, list) for seq in encoded_batch)
+
+    encoded_truncated = tokenizer.encode_batch(batch, max_length=3)
+    assert all(len(seq) <= 3 for seq in encoded_truncated)
+
+
+def test_special_token_indexes():
+    tokenizer = TokenizerConstructor(
+        tokenizer_type="BPE",
+        pre_tokenizers="Whitespace",
+        special_tokens=["<unk>", "<sos>", "<eos>", "<pad>", "\n"]
+    )
+    assert tokenizer.unknown_token == tokenizer.special_tokens.index("<unk>")
+    assert tokenizer.start_token == tokenizer.special_tokens.index("<sos>")
+    assert tokenizer.end_token == tokenizer.special_tokens.index("<eos>")
+    assert tokenizer.pad_token == tokenizer.special_tokens.index("<pad>")
+    assert tokenizer.new_line_token == tokenizer.special_tokens.index("\n")