| from tokenizers import Tokenizer | |
| from tokenizers.models import WordLevel, BPE | |
| from tokenizers.trainers import WordLevelTrainer, BpeTrainer | |
| from tokenizers.pre_tokenizers import Whitespace, ByteLevel | |
| from pathlib import Path | |
| def get_all_sentences(ds, lang: str): | |
| for item in ds: | |
| yield item['translation'][lang] | |
| def get_or_build_local_tokenizer(config, ds, lang: str, tokenizer_type: str, force_build: bool = False) -> Tokenizer: | |
| tokenizer_path = Path(config['dataset']['tokenizer_file'].format(lang)) | |
| if not Path.exists(tokenizer_path) or force_build: | |
| if ds is None: | |
| raise ValueError("Cannot find local tokenizer, dataset given is None") | |
| if tokenizer_type == "WordLevel": | |
| tokenizer = Tokenizer(WordLevel(unk_token='<unk>')) | |
| tokenizer.pre_tokenizer = Whitespace() | |
| trainer = WordLevelTrainer(special_tokens=['<unk>', '<pad>', '<sos>', '<eos>'], min_frequency=2) | |
| elif tokenizer_type == "BPE": | |
| tokenizer = Tokenizer(BPE(unk_token='<unk>')) | |
| tokenizer.pre_tokenizer = Whitespace() | |
| trainer = BpeTrainer(special_tokens=['<unk>', '<pad>', '<sos>', '<eos>'], min_frequency=2) | |
| else: | |
| raise ValueError("Unsupported Tokenizer type") | |
| tokenizer.train_from_iterator( | |
| get_all_sentences(ds, lang), trainer=trainer | |
| ) | |
| tokenizer.save(str(tokenizer_path)) | |
| else: | |
| tokenizer = Tokenizer.from_file(str(tokenizer_path)) | |
| return tokenizer |