Update cognitivess_model/tokenization_cognitivess.py

cognitivess_model/tokenization_cognitivess.py

@@ -3,8 +3,8 @@ import json
 import os
 
 class CognitivessTokenizer(PreTrainedTokenizer):
-    def __init__(self, vocab_file, merges_file, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self, vocab_file, merges_file, **kwargs):
+        super().__init__(**kwargs)
         self.vocab_file = vocab_file
         self.merges_file = merges_file
         self.encoder = self.load_vocab(vocab_file)
@@ -15,6 +15,12 @@ class CognitivessTokenizer(PreTrainedTokenizer):
         bpe_merges = [tuple(merge_str.split()) for merge_str in bpe_data.split("\n")[1:-1]]
         self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
 
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *init_inputs, **kwargs):
+        vocab_file = os.path.join(pretrained_model_name_or_path, "vocab.json")
+        merges_file = os.path.join(pretrained_model_name_or_path, "merges.txt")
+        return cls(vocab_file, merges_file, **kwargs)
+
     @property
     def vocab_size(self):
         return len(self.encoder)
@@ -54,4 +60,4 @@ class CognitivessTokenizer(PreTrainedTokenizer):
 
     def load_vocab(self, vocab_file):
         with open(vocab_file, "r", encoding="utf-8") as f:
-            return json.load(f)
+            return json.load(f)