t5-base-dutch / train_tokenizer.py
yhavinga's picture
Retrain tokenizer for case sensitive
a37447a
raw
history blame
2 kB
from datasets import load_dataset
from t5_tokenizer_model import SentencePieceUnigramTokenizer
# from tokenizers import trainers, Tokenizer, normalizers, ByteLevelBPETokenizer
data_dir = "/home/yeb"
data_files = []
def train_val_files():
import glob
import random
SEED = 12345
def add_jsonlines_dir(path, filespec):
global data_files
data_files += glob.glob(f"{path}/{filespec}")
print(f"Number of files {len(data_files)} after adding {path}")
# add_jsonlines_dir(f"{data_dir}/oscar_nl_cleaned")
add_jsonlines_dir(f"{data_dir}/c4_cleaned2", "*47*.gz")
add_jsonlines_dir(f"{data_dir}/nrc_uniq_cleaned_20210223", "*.gz")
add_jsonlines_dir(f"{data_dir}/nu_uniq_cleaned_20210225", "*.gz")
random.Random(SEED).shuffle(data_files)
print(data_files)
total = len(data_files)
print(total)
perc = 0.01
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"
return train, val
train, val = train_val_files()
dataset = load_dataset('json', data_files={'train': train, 'validation': val}, split='train')
vocab_size = 32000
input_sentence_size = None
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("./tokenizer.json")