sds

tokenizeConfig.py

@@ -3,26 +3,45 @@ from tokenizers import Tokenizer, models, pre_tokenizers, trainers, decoders
 import json
 from typing import List, Optional, Union, Dict
 from transformers.tokenization_utils_base import EncodedInput, BatchEncoding
-from transformers import PreTrainedTokenizer
-from transformers.utils import logging, PaddingStrategy
-
+from transformers.utils import PaddingStrategy
 
 class OBITokenizer(PreTrainedTokenizer):
-    def __init__(self, auto_map={"AutoTokenizer": ["tokenizeConfig.OBITokenizer"]},tokenizer_class= "OBITokenizer",**kwargs):
-         # Initialize your tokenizer with the auto_map parameter if needed
-        self.auto_map=auto_map
-        self.tokenizer_class=tokenizer_class
+    def __init__(
+        self,
+        vocab_file,
+        unk_token="<unk>",
+        bos_token="<s>",
+        eos_token="</s>",
+        pad_token=None,
+        add_bos_token=True,
+        add_eos_token=False,
+        clean_up_tokenization_spaces=False,
+        auto_map={"AutoTokenizer": ["tokenizeConfig.OBITokenizer"]},
+        tokenizer_class="OBITokenizer",
+        **kwargs,
+    ):
+        super().__init__(
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            add_bos_token=add_bos_token,
+            add_eos_token=add_eos_token,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+
         # Initialize a BPE model for tokenization
         bpe_model = models.BPE()
-        # Initialize the tokenizer
         self.tokenizer = Tokenizer(bpe_model)
+
         # Add pre-tokenization and decoding steps if needed
         self.tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel()
         self.tokenizer.decoder = decoders.ByteLevel()
-        super().__init__(**kwargs)
 
         # Set the padding token
         self.pad_token = "[PAD]"
+
         # Set the special tokens
         self.cls_token = "[CLS]"
         self.sep_token = "[SEP]"
@@ -31,30 +50,29 @@ class OBITokenizer(PreTrainedTokenizer):
         self.bos_token = "[CLS]"
         self.eos_token = "[SEP]"
         self.pad_token = "[PAD]"
-      
+
+        # Load the vocabulary file
+        self.tokenizer.get_vocab().add_special_tokens([self.cls_token, self.sep_token, self.unk_token, self.mask_token])
 
     def _pad(
-            self,
-            encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
-            max_length: Optional[int] = None,
-            padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
-            pad_to_multiple_of: Optional[int] = None,
-            return_attention_mask: Optional[bool] = None,
+        self,
+        encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
+        max_length: Optional[int] = None,
+        padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
+        pad_to_multiple_of: Optional[int] = None,
+        return_attention_mask: Optional[bool] = None,
     ) -> dict:
         # Modify the _pad method as needed for OBITokenizer
         # You can inherit the implementation from ChatGLMTokenizer and customize it further
         return super()._pad(encoded_inputs, max_length, padding_strategy, pad_to_multiple_of, return_attention_mask)
-        
 
-    def train(self, files,save_path):
+    def train(self, files, save_path):
         # Training: Fit the tokenizer on your text data
         trainer = trainers.BpeTrainer(special_tokens=["[PAD]", "[CLS]", "[SEP]", "[MASK]", "[UNK]"])
         self.tokenizer.train(trainer=trainer, files=files)
         # Save the trained tokenizer to a file
         self.tokenizer.save(save_path)
 
-
-
     def save_config(self, config_file):
         # Serialize the tokenizer's config to a JSON file
         config_dict = {
@@ -62,7 +80,7 @@ class OBITokenizer(PreTrainedTokenizer):
             "vocab_size": self.tokenizer.get_vocab_size(),
             "tokenizer_class": "OBITokenizer",
             "auto_map": {
-                "AutoTokenizer": ["tokenizeConfig.OBITokenizer","null"]
+                "AutoTokenizer": ["tokenizeConfig.OBITokenizer", "null"]
             },
             "bos_token": "[CLS]",
             "eos_token": "[SEP]",
@@ -81,4 +99,4 @@ class OBITokenizer(PreTrainedTokenizer):
 
     def decode(self, ids):
         # Decode IDs to text using the custom tokenizer
-        return self.tokenizer.decode(ids)
+        return self.tokenizer.decode(ids)