Create tokenization_octagon.py

tokenization_octagon.py

@@ -0,0 +1,54 @@
+from transformers import PreTrainedTokenizerFast
+from tokenizers import Tokenizer, normalizers, pre_tokenizers, trainers, models
+from tokenizers.normalizers import Lowercase, NFD, StripAccents
+from tokenizers.pre_tokenizers import Whitespace
+from typing import Optional, List, Union
+
+class OctagonTokenizer(PreTrainedTokenizerFast):
+    def __init__(
+        self,
+        vocab_file=None,
+        merges_file=None,
+        tokenizer_file=None,
+        unk_token="[UNK]",
+        sep_token="[SEP]",
+        pad_token="[PAD]",
+        cls_token="[CLS]",
+        mask_token="[MASK]",
+        **kwargs
+    ):
+        super().__init__(
+            tokenizer_file=tokenizer_file,
+            unk_token=unk_token,
+            sep_token=sep_token,
+            pad_token=pad_token,
+            cls_token=cls_token,
+            mask_token=mask_token,
+            **kwargs
+        )
+    
+    @classmethod
+    def train_tokenizer(cls, texts: List[str], vocab_size: int = 30522, save_path: Optional[str] = None):
+        # Initialize a tokenizer
+        tokenizer = Tokenizer(models.BPE())
+        
+        # Normalizer
+        tokenizer.normalizer = normalizers.Sequence([NFD(), Lowercase(), StripAccents()])
+        
+        # Pre-tokenizer
+        tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()
+        
+        # Trainer
+        trainer = trainers.BpeTrainer(
+            vocab_size=vocab_size,
+            special_tokens=["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"]
+        )
+        
+        # Train the tokenizer
+        tokenizer.train_from_iterator(texts, trainer=trainer)
+        
+        # Save if path is provided
+        if save_path:
+            tokenizer.save(save_path)
+        
+        return cls(tokenizer_file=save_path) if save_path else cls(tokenizer_object=tokenizer)