NbAiLabArchive/test_w5_long_dataset

1from datasets import load_dataset, concatenate_datasets2from tokenizers import trainers, Tokenizer, normalizers, ByteLevelBPETokenizer3 4model_dir = "./"  # ${MODEL_DIR}5 6# load dataset7dataset = load_dataset("json", data_files=["/mnt/disks/flaxdisk/corpus/norwegian_colossal_corpus_validation.json","/mnt/disks/flaxdisk/corpus/special_chars.json"], split='train')8 9 10# Instantiate tokenizer11tokenizer = ByteLevelBPETokenizer()12 13def batch_iterator(batch_size=1000):14    for i in range(0, len(dataset), batch_size):15        yield dataset[i: i + batch_size]["text"]16 17# Customized training18tokenizer.train_from_iterator(batch_iterator(), vocab_size=50265, min_frequency=2, special_tokens=[19    "<s>",20    "<pad>",21    "</s>",22    "<unk>",23    "<mask>",24])25 26 27# Save files to disk28tokenizer.save(f"{model_dir}/tokenizer.json")29