add 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,46 @@ 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
+
+ @property
+ def inv_vocab(self):
+ raise NotImplementedError
+
+ 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.",
  )