load from here

tokenizer/tokenizer.py

@@ -0,0 +1,147 @@
+import json
+import re
+from collections import Counter
+import pickle
+import argparse
+
+class Tokenizer:
+    def __init__(self):
+        self.special_tokens = ["[PAD]", "[MASK]"]
+        self.vocab = {}
+        self.token_to_id = {}
+        self.id_to_token = {}
+
+    def tokenize(self, text):
+        # Match words, numbers, periods, and commas as separate tokens
+        tokens = re.findall(r'\w+|[.,]|\[mask\]|\[pad\]', text.lower())
+        # Restore MASK and PAD to all caps
+        modified_list = []
+        for s in tokens:
+            modified_s = s.replace("[mask]", "[MASK]").replace("[pad]", "[PAD]")
+            modified_list.append(modified_s)
+        return modified_list
+
+    def pad_sequence(self, tokens, length):
+        """Pads tokenized sequences to length with a padding token (assumed to be '[PAD]')."""
+        if len(tokens) > length:
+            raise ValueError(f"Token sequence length {len(tokens)} exceeds specified length {length}.")
+        
+        pad_token = self.token_to_id["[PAD]"]
+        return tokens + [pad_token] * (length - len(tokens))
+
+    def build_vocab(self, dataset_path, min_freq=1):
+        token_counter = Counter()
+
+        with open(dataset_path, 'r') as f:
+            data = json.load(f)
+            for entry in data:
+                caption = entry['caption']
+                tokens = self.tokenize(caption)
+                token_counter.update(tokens)
+
+        # Keep tokens that meet the min frequency
+        tokens = [tok for tok, count in token_counter.items() if count >= min_freq]
+
+        # Ensure special tokens are always included
+        all_tokens = self.special_tokens + sorted(tokens)
+        
+        # Build vocab dictionaries
+        self.vocab = {tok: idx for idx, tok in enumerate(all_tokens)}
+        self.token_to_id = self.vocab
+        self.id_to_token = {idx: tok for tok, idx in self.vocab.items()}
+
+        print(f"Vocabulary size: {len(self.vocab)}")
+
+    def encode(self, text):
+        tokens = self.tokenize(text)
+        encoded = []
+        for tok in tokens:
+            if tok not in self.token_to_id:
+                raise ValueError(f"Unknown token encountered: {tok} in {text}")
+            encoded.append(self.token_to_id[tok])
+        return encoded
+
+    def encode_batch(self, texts, pad_to_length=None):
+        """
+        Encode a batch of texts into token IDs with padding to ensure uniform length.
+    
+        Args:
+            texts (list): A list of strings to encode
+            pad_to_length (int, optional): Length to pad all sequences to. If None,
+                                          will pad to the length of the longest sequence.
+        
+        Returns:
+            list: A list of lists, where each inner list contains the token IDs for a text
+        """
+        # Get the padding token ID
+        pad_token = self.token_to_id["[PAD]"]
+    
+        # First encode all texts
+        encoded_texts = []
+        for text in texts:
+            try:
+                encoded = self.encode(text)
+                encoded_texts.append(encoded)
+            except ValueError as e:
+                raise ValueError(f"Error encoding text: {text}. {str(e)}")
+    
+        # Determine padding length
+        if pad_to_length is None:
+            pad_to_length = max(len(seq) for seq in encoded_texts)
+    
+        # Pad sequences to uniform length
+        padded_texts = []
+        for seq in encoded_texts:
+            if len(seq) > pad_to_length:
+                # Truncate if too long
+                padded_texts.append(seq[:pad_to_length])
+            else:
+                # Pad if too short
+                padding = [pad_token] * (pad_to_length - len(seq))
+                padded_texts.append(seq + padding)
+    
+        return padded_texts
+
+    def decode(self, token_ids):
+        return ' '.join(self.id_to_token[tok_id] for tok_id in token_ids)
+
+    def save(self, path):
+        with open(path, 'wb') as f:
+            pickle.dump({'vocab': self.vocab}, f)
+
+    def load(self, path):
+        with open(path, 'rb') as f:
+            data = pickle.load(f)
+            self.vocab = data['vocab']
+            self.token_to_id = self.vocab
+            self.id_to_token = {idx: tok for tok, idx in self.vocab.items()}
+
+    def get_vocab(self):
+        return sorted(self.vocab.keys())
+
+    def get_vocab_size(self):
+        return len(self.vocab)
+
+if __name__ == "__main__":
+    tokenizer = Tokenizer()
+
+    parser = argparse.ArgumentParser(description="Tokenizer utility for saving and loading vocabularies.")
+    parser.add_argument("action", choices=["save", "load"], help="Action to perform: 'save' or 'load'.")
+    parser.add_argument("--json_file", type=str, default='Mario_LevelsAndCaptions.json', help="Path to the JSON file containing the dataset (required for 'save').")
+    parser.add_argument("--pkl_file", type=str, default='Mario_Tokenizer.pkl', help="Path to the pickle file to save/load the tokenizer.")
+
+    args = parser.parse_args()
+
+    if args.action == "save":
+        if not args.json_file:
+            raise ValueError("The --json_file argument is required for the 'save' action.")
+        tokenizer.build_vocab(args.json_file)
+        tokenizer.save(args.pkl_file)
+    elif args.action == "load":
+        tokenizer.load(args.pkl_file)
+
+    # Example usage
+    #print(tokenizer.encode("floor with one gap. one enemy."))
+    #print(tokenizer.get_vocab())
+    #for id, token in tokenizer.id_to_token.items():
+    #    print(id,":",token)