Update tokenization_xgen.py

tokenization_xgen.py

@@ -60,9 +60,18 @@ def tiktoken_tokenizer(base="gpt2", pad_token=None, add_special=True):
         ]
         return fim_tokens
 
+    def include_additional_tokens():
+        tokens = []
+        tokens += [f"<dummy_{i}>" for i in range(4)]
+        tokens.append("<sep>")  # 50317
+        tokens.append("<eom>")  # 50318
+        tokens += [f"<mask_{i}>" for i in reversed(range(1, 51199-50318+1))]
+        return tokens        
+
     add_whitespaces = include_whitespace(n_min=2, n_max=32)
     add_tabs = include_tabs(n_min=2, n_max=10)
     fim_tokens = include_fim_tokens()
+    additional_tokens = include_additional_tokens()
 
     tokenizer = tiktoken.get_encoding(base)
 
@@ -82,6 +91,9 @@ def tiktoken_tokenizer(base="gpt2", pad_token=None, add_special=True):
     for sp in fim_tokens:
         special_tokens[sp] = idx
         idx += 1
+    for sp in additional_tokens:
+        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