add pad token and default eos token

tokenization_xgen.py

@@ -25,7 +25,7 @@ MAX_MODEL_INPUT_SIZES = {
 }
 
 
-def tiktoken_tokenizer(base="gpt2", add_special=True):
+def tiktoken_tokenizer(base="gpt2", pad_token=None, add_special=True):
     if not add_special:
         return tiktoken.get_encoding(base)
 
@@ -83,6 +83,9 @@ def tiktoken_tokenizer(base="gpt2", add_special=True):
         special_tokens[sp] = idx
         idx += 1
 
+    if pad_token and pad_token not in tokenizer._special_tokens and pad_token not in special_tokens:
+        special_tokens[pad_token] = idx
+        idx += 1
     # In production, load the arguments directly instead of accessing private attributes
     # See openai_public.py for examples of arguments for specific encodings
     enc = tiktoken.Encoding(
@@ -112,25 +115,40 @@ class XgenTokenizer(PreTrainedTokenizer):
     def __init__(
             self,
             pad_token=None,
+            eos_token="<|endoftext|>",
             add_eos_token=False,
             add_special_tokens=True,
             **kwargs,
     ):
-        pad_token = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_token
+        pad_token_added = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_token
+        eos_token_added = AddedToken(eos_token, lstrip=False, rstrip=False) if isinstance(eos_token, str) else eos_token
         super().__init__(
-            pad_token=pad_token,
+            pad_token=pad_token_added,
+            eos_token=eos_token_added,
             add_eos_token=add_eos_token,
             add_special_tokens=add_special_tokens,
             **kwargs,
         )
         self.add_eos_token = add_eos_token
-        self.encoder = tiktoken_tokenizer(base="gpt2", add_special=add_special_tokens)
+        self.encoder = tiktoken_tokenizer(base="gpt2", pad_token=pad_token, add_special=add_special_tokens)
 
     @property
     def vocab_size(self):
         """Returns vocab size"""
         return self.encoder.n_vocab
 
+    @property
+    def eos_token_id(self):
+        if self.eos_token is not None:
+            return self.encoder.encode(self.eos_token, allowed_special="all")[0]
+        return None
+
+    @property
+    def pad_token_id(self):
+        if self.pad_token is not None:
+            return self.encoder.encode(self.pad_token, allowed_special="all")[0]
+        return None
+
     def get_vocab(self):
         """Returns vocab as a dict"""
         vocab = {self._convert_id_to_token(i): i for i in range(self.vocab_size)}
@@ -142,6 +160,9 @@ class XgenTokenizer(PreTrainedTokenizer):
 
     def _convert_token_to_id(self, token):
         """Converts a token (str) in an id using the vocab."""
+        if isinstance(token, str):
+            ids = self._tokenize(token)
+            return ids[0]
         return token
 
     def _convert_id_to_token(self, index):
@@ -216,4 +237,4 @@ class XgenTokenizer(PreTrainedTokenizer):
         if token_ids_1 is not None:
             output += [1] * len(token_ids_1 + eos_token_id)
 
-        return output
+        return output