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")