Merge pull request EleutherAI#779 from EleutherAI/tiktoken

Add tiktoken Tokenizer support

configs/neox_arguments.md

@@ -111,7 +111,7 @@ Logging Arguments
 
 - **git_hash**: str
 
-    Default = 62e9440
+    Default = 3463ca4
 
     current git hash of repository
 
@@ -987,11 +987,11 @@ Tokenizer Arguments
 
 
 
-- **tokenizer_type**: typing.Literal['GPT2BPETokenizer', 'HFTokenizer', 'HFGPT2Tokenizer', 'SPMTokenizer', 'CharLevelTokenizer']
+- **tokenizer_type**: typing.Literal['GPT2BPETokenizer', 'HFTokenizer', 'HFGPT2Tokenizer', 'SPMTokenizer', 'CharLevelTokenizer', 'TiktokenTokenizer']
 
     Default = GPT2BPETokenizer
 
-    Type of tokenizer to use - should be one of ["GPT2BPETokenizer", "HFTokenizer", "HFGPT2Tokenizer", "SPMTokenizer", "CharLevelTokenizer"]
+    Type of tokenizer to use - should be one of ["GPT2BPETokenizer", "HFTokenizer", "HFGPT2Tokenizer", "SPMTokenizer", "CharLevelTokenizer", "TiktokenTokenizer"]
 
 
 

megatron/neox_arguments/neox_args.py

@@ -658,9 +658,10 @@ class NeoXArgsTokenizer(NeoXArgsTemplate):
         "HFGPT2Tokenizer",
         "SPMTokenizer",
         "CharLevelTokenizer",
+        "TiktokenTokenizer",
     ] = "GPT2BPETokenizer"
     """
-    Type of tokenizer to use - should be one of ["GPT2BPETokenizer", "HFTokenizer", "HFGPT2Tokenizer", "SPMTokenizer", "CharLevelTokenizer"]
+    Type of tokenizer to use - should be one of ["GPT2BPETokenizer", "HFTokenizer", "HFGPT2Tokenizer", "SPMTokenizer", "CharLevelTokenizer", "TiktokenTokenizer"]
     """
 
     padded_vocab_size: int = None

megatron/tokenizer/tokenizer.py

@@ -24,6 +24,7 @@
 from transformers import GPT2Tokenizer, GPT2TokenizerFast
 import numpy as np
 import sentencepiece as spm
+import tiktoken
 from typing import List, Union
 from .gpt2_tokenization import GPT2Tokenizer
 
@@ -52,6 +53,9 @@ def build_tokenizer(args):
         tokenizer = HFGPT2Tokenizer(args.vocab_file)
     elif args.tokenizer_type.lower() == "CharLevelTokenizer".lower():
         tokenizer = CharLevelTokenizer(vocab_size=512)
+    elif args.tokenizer_type.lower() == "TiktokenTokenizer".lower():
+        assert args.vocab_file is not None
+        tokenizer = TiktokenTokenizer(args.vocab_file)
     else:
         raise NotImplementedError(
             "{} tokenizer is not " "implemented.".format(args.tokenizer_type)
@@ -345,3 +349,47 @@ def detokenize(self, token_ids):
     @property
     def eod(self):
         return self.eod_id
+
+
+class TiktokenTokenizer(AbstractTokenizer):
+    """Tokenizer from OpenAI's tiktoken implementation"""
+
+    def __init__(self, vocab_file):
+        name = "TiktokenTokenizer"
+        super().__init__(name)
+
+        self.tokenizer = tiktoken.get_encoding(vocab_file)
+        self.eod_id = self.tokenizer.eot_token
+        self.pad_id = None
+
+    @property
+    def vocab_size(self):
+        return self.tokenizer.n_vocab
+
+    @property
+    def vocab(self):
+        raise NotImplementedError("TiktokenTokenizer does not implement vocabulary access.")
+
+    @property
+    def inv_vocab(self):
+        raise NotImplementedError("TiktokenTokenizer does not implement vocabulary access. \
+                To get the idx-th token in vocabulary, use tokenizer.decode([idx]) .")
+
+    def tokenize(self, text: str):
+        return self.tokenizer.encode(text) #,  allowed_special="all")
+
+    def tokenize_batch(self, text_batch: List[str]):
+        return self.tokenizer.encode_batch(text_batch, allowed_special="all")
+
+    def detokenize(self, token_ids):
+        return self.tokenizer.decode(tokens=token_ids, errors="strict")
+
+    @property
+    def eod(self):
+        return self.eod_id
+
+    @property
+    def pad(self):
+        raise NotImplementedError
+
+

prepare_data.py

@@ -20,6 +20,7 @@
     "HFTokenizer",
     "GPT2BPETokenizer",
     "CharLevelTokenizer",
+    "TiktokenTokenizer",
 ]
 DATASET_CHOICES = [i for i in DATA_DOWNLOADERS.keys() if i != "pass"]
 

requirements/requirements.txt

@@ -10,6 +10,7 @@ pybind11==2.6.2
 regex
 sentencepiece
 six
+tiktoken==0.1.2
 tokenizers==0.12.1
 transformers~=4.24.0
 wandb==0.10.28

tools/preprocess_data.py

@@ -93,6 +93,7 @@ def get_args():
             "HFTokenizer",
             "GPT2BPETokenizer",
             "CharLevelTokenizer",
+            "TiktokenTokenizer",
         ],
         help="What type of tokenizer to use.",
     )