import os
from tqdm import tqdm
import numpy as np
import tiktoken
from datasets import load_dataset

enc = tiktoken.get_encoding("gpt2")

if __name__ == '__main__':
    dataset = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train")
    split_dataset = dataset.train_test_split(test_size=0.0005, seed=2357, shuffle=True)
    split_dataset['val'] = split_dataset.pop('test')

    def process(example):
        ids = enc.encode_ordinary(example['text'])
        ids.append(enc.eot_token)
        return {'ids': ids, 'len': len(ids)}

    tokenized = split_dataset.map(process, remove_columns=['text'], desc="tokenizing", num_proc=8)

    for split, dset in tokenized.items():
        arr_len = np.sum(dset['len'], dtype=np.int64)
        filename = os.path.join(os.path.dirname(__file__), f'{split}.bin')
        arr = np.memmap(filename, dtype=np.uint16, mode='w+', shape=(arr_len,))
        idx = 0
        for example in tqdm(dset, desc=f"writing {filename}"):
            arr[idx : idx + example['len']] = example['ids']
            idx += example['len']
        arr.flush()