import os
import numpy as np
from datasets import load_dataset
from transformers import AutoTokenizer
from tqdm import tqdm

# --- CONFIG ---
NUM_TOKENS_TO_COLLECT = 20_000_000_000 # 20 billion
VAL_TOKENS = 2_000_000                 # 2 million
TOKENIZER_NAME = "gpt2"
DATASET_NAME = "HuggingFaceFW/fineweb-edu"
DATASET_CONFIG = "sample-100BT"

print(f"Loading tokenizer: {TOKENIZER_NAME}...")
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
# Sicherheitscheck: EOS Token garantieren
if tokenizer.eos_token is None:
    tokenizer.add_special_tokens({'eos_token': '<|endoftext|>'})

print("Loading FineWeb-Edu 100BT (Streaming)...")
dataset = load_dataset(DATASET_NAME, name=DATASET_CONFIG, split="train", streaming=True)

train_file = "train.bin"
val_file = "val.bin"

for f in [train_file, val_file]:
    if os.path.exists(f):
        os.remove(f)

tokens_collected = 0
val_collected = 0
row_counter = 0
VAL_FREQUENCY = 100

print(f"Starting tokenization. Target: {NUM_TOKENS_TO_COLLECT:,} tokens. (Data Type: uint32)")

dtype = np.uint32 

with open(train_file, "wb") as f_train, open(val_file, "wb") as f_val:
    for row in tqdm(dataset):
        train_done = tokens_collected >= NUM_TOKENS_TO_COLLECT
        val_done = val_collected >= VAL_TOKENS
        if train_done and val_done:
            print("\nTarget reached! Stopping streaming.")
            break

        tokens = tokenizer.encode(row["text"], add_special_tokens=False)
        tokens.append(tokenizer.eos_token_id)
        tokens_arr = np.array(tokens, dtype=dtype)

        is_val_row = (row_counter % VAL_FREQUENCY == 0) and not val_done
        if is_val_row:
            f_val.write(tokens_arr.tobytes())
            val_collected += len(tokens)
        elif not train_done:
            f_train.write(tokens_arr.tobytes())
            tokens_collected += len(tokens)

        row_counter += 1

print(f"Done! Train tokens: {tokens_collected:,} | Val tokens: {val_collected:,}")
print("Ready for train.py!")