pino-bigbird-roberta-base / train_tokenizer.py
dat
Saving weights and logs at step 1252
f291f93
raw
history blame
1.36 kB
import glob
import random
from tokenizers import ByteLevelBPETokenizer
from datasets import load_dataset
data_files = []
def add_jsonlines_dir(path, filespec):
global data_files
data_files += glob.glob(f"{path}/{filespec}")
data_files = list(set(data_files))
print(f"Number of files {len(data_files)} after adding {path} glob {filespec}")
add_jsonlines_dir(f"/data/c4_cleaned2", "*.gz")
add_jsonlines_dir(f"/data/nrc_uniq_cleaned_20210223", "*.gz")
add_jsonlines_dir(f"/data/nu_uniq_cleaned_20210225", "*.gz")
random.Random(42).shuffle(data_files)
total = len(data_files)
print(total)
perc = 0.05
val_size = int(perc * total)
train_size = total - val_size
train = data_files[:train_size]
val = data_files[train_size:]
print(f"Got {len(train)} training files and {perc * 100} % {len(val)} validation files")
assert list(set(train) & set(val)) == [], "Train overlaps with test"
datasets = load_dataset('json', data_files={'train': train, 'validation': val})
tokenizer = ByteLevelBPETokenizer()
def batch_iterator(batch_size=1000):
for i in range(0, len(datasets), batch_size):
yield datasets["train"][i: i + batch_size]["text"]
tokenizer.train_from_iterator(batch_iterator(), vocab_size=50358, min_frequency=2, special_tokens=[
"<s>",
"<pad>",
"</s>",
"<unk>",
"<mask>",
])
tokenizer.save("tokenizer.json")