Octa / tokenization_octagon.py
Georg4000's picture
Create tokenization_octagon.py
2ae8b59 verified
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)