File size: 1,786 Bytes
2ae8b59 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 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) |