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

DATASET_NAME = "HuggingFaceFW/fineweb-edu"
SAMPLE_NAME = "sample-10BT"
TARGET_TOKENS = 100_000_000
NUM_PROC = 8 

enc = tiktoken.get_encoding("gpt2")

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

if __name__ == "__main__":
    print(f"Loading streaming dataset {DATASET_NAME}...")
    dataset = load_dataset(DATASET_NAME, name=SAMPLE_NAME, split='train', streaming=True)
    
    all_tokens = []
    total_tokens = 0
    
    pbar = tqdm(total=TARGET_TOKENS, desc="Collecting tokens")
    for example in dataset:
        tokens = process(example)['ids']
        all_tokens.extend(tokens)
        total_tokens += len(tokens)
        pbar.update(len(tokens))
        if total_tokens >= TARGET_TOKENS:
            break
    pbar.close()

    n = len(all_tokens)
    train_data = all_tokens[:int(n*0.95)]
    val_data = all_tokens[int(n*0.95):]

    for name, d in [('train', train_data), ('val', val_data)]:
        arr = np.array(d, dtype=np.uint16)
        filename = f"{name}.bin"
        arr.tofile(filename)
        print(f"Saved {filename} with {len(d):,} tokens.")

    print("\nDone!")