Create model/tokenizer.py

model/tokenizer.py

@@ -0,0 +1,24 @@
+from tokenizers import Tokenizer, models, trainers, pre_tokenizers, processors, decoders
+from tokenizers.normalizers import NFC
+from transformers import PreTrainedTokenizerFast
+from pathlib import Path
+
+SPECIAL_TOKENS = [
+    "<|pad|>", "<|bos|>", "<|eos|>", "<|unk|>", "<|sep|>",
+    "<|fim_prefix|>", "<|fim_middle|>", "<|fim_suffix|>",
+    "<|python|>", "<|javascript|>", "<|typescript|>",
+    "<|cpp|>", "<|rust|>", "<|go|>", "<|java|>", "<|bash|>",
+]
+
+def get_gpt2_tokenizer_for_code():
+    from transformers import AutoTokenizer
+    tok = AutoTokenizer.from_pretrained("gpt2")
+    tok.pad_token = tok.eos_token
+    tok.add_special_tokens({"additional_special_tokens": [
+        "<|fim_prefix|>", "<|fim_middle|>", "<|fim_suffix|>",
+        "<|python|>", "<|javascript|>", "<|rust|>", "<|go|>",
+    ]})
+    return tok
+
+def load_tokenizer(save_dir="./tokenizer"):
+    return PreTrainedTokenizerFast.from_pretrained(save_dir)