add pad token and default eos token (#7)

- add pad token and default eos token (b5bab0d91365fa7bfe453985b1db4f2cc073d83d)


Co-authored-by: Tian Xie <[email protected]>

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,19 +115,22 @@ 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):
@@ -142,6 +148,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 +225,4 @@ class XgenTokenizer(PreTrainedTokenizer):
         if token_ids_1 is not None:
             output += [1] * len(token_ids_1 + eos_token_id)
 
-        return output
+        return output