nb-t5-base / train_tokenizer.py
from datasets import load_dataset, concatenate_datasets
from tokenizers import trainers, Tokenizer, normalizers
from t5_tokenizer_model import SentencePieceUnigramTokenizer
vocab_size = 50_000
input_sentence_size = None
model_dir = "./" # ${MODEL_DIR}
# Initialize a dataset
dataset = load_dataset("json", data_files=["/mnt/disks/flaxdisk/corpus/norwegian_colossal_corpus_validation.json","/mnt/disks/flaxdisk/corpus/special_chars.json"], split='train')
tokenizer = SentencePieceUnigramTokenizer(unk_token="<unk>", eos_token="</s>", pad_token="<pad>")
# Build an iterator over this dataset
def batch_iterator(input_sentence_size=None):
if input_sentence_size is None:
input_sentence_size = len(dataset)
batch_length = 100
for i in range(0, input_sentence_size, batch_length):
yield dataset[i: i + batch_length]["text"]
# Train tokenizer
tokenizer.train_from_iterator(
iterator=batch_iterator(input_sentence_size=input_sentence_size),
vocab_size=vocab_size,
show_progress=True,
)
# Save files to disk
tokenizer.save(f"{model_dir}/tokenizer.json")