Update tokenizer.py

tokenizer.py

@@ -75,14 +75,25 @@ class HFAutoTokenizer:
 
 class ByteTokenizer(PreTrainedTokenizer):
     """UTF-8 Encoder."""
-
     def __init__(self):
-        self.vocab_size = 512
-        self.eod_id = 0
-        self.eos_id = 0
-        self.eos_token = 0
-        self.eos_token_id = 0
-        self.pad_id = 1
+        super().__init__(
+            bos_token=self.decode_token(2),
+            eos_token=self.decode_token(0),
+            unk_token=self.decode_token(0),
+            pad_token=self.decode_token(1),
+            mask_token=self.decode_token(3),
+        )
+    
+    @property
+    def vocab_size(self) -> int:
+        return 512
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        return cls()
+
+    def get_vocab(self):
+        return {str(i): i for i in range(512)}
 
     def clamp(self, n):
         return max(32, min(n, self.vocab_size))
@@ -90,12 +101,15 @@ class ByteTokenizer(PreTrainedTokenizer):
     def decode_token(self, token: int):
         return str(chr(self.clamp(token)))
 
-    def __call__(self, text: str, *args, **kwargs):
+    def __call__(self, text: str, return_tensors: bool = False, *args, **kwargs):
         ids = torch.tensor(self.tokenize(text), dtype=torch.long).unsqueeze(0)
-        return {"input_ids": ids}
-    
+        return {"input_ids": ids} if return_tensors == False else ids
+
+    def _tokenize(self, text: str):
+        return np.frombuffer(text.encode('utf-8'), dtype=np.uint8)
+        
     def tokenize(self, text: str):
-        return list(np.fromstring(text, dtype=np.uint8))
+        return self._tokenize(text).tolist()
 
     def tokenize_batch(self, text_batch: Union[List[str], str]):
         if isinstance(text_batch, list):
@@ -109,7 +123,7 @@ class ByteTokenizer(PreTrainedTokenizer):
     def decode_batch(self, token_ids: Union[List[str], str]):
         if isinstance(token_ids, list):
             return [self.decode(s) for s in token_ids]
-        # elif if tensor, convert to list first
+        
         elif isinstance(token_ids, torch.Tensor):
             return [self.decode(s) for s in token_ids.tolist()]
         else: