encode_dataset.py

import numpy as np
from tokenizers import Tokenizer
import os

TOKENIZER_PATH = "./tokenizer/tokenizer.json"
DATA_DIR       = "./data"
OUT_DIR        = "./data/encoded"
os.makedirs(OUT_DIR, exist_ok=True)

DTYPE = np.uint16

CHUNK_SIZE=10000 #encode this much then flush to disk

tokenizer = Tokenizer.from_file(TOKENIZER_PATH)

tokenizer.post_processor = None #dispable post processing, we will handle bos/eos ourselves for training

def encode_file(txt_path, out_path):
    print(f"Encoding: {txt_path} -> {out_path}")

    total_tokens = 0
    tmp_path = out_path + ".tmp"
    with open(txt_path, "r", encoding="utf-8") as f_in, open(tmp_path, "wb") as f_tmp:
        batch = []
        for line in f_in:
            line = line.strip()
            if line:
                batch.append(line)
            if len(batch) >= CHUNK_SIZE:
                encodings = tokenizer.encode_batch(batch)
                for enc in encodings:
                    ids = np.array(enc.ids, dtype=DTYPE)
                    ids.tofile(f_tmp)
                    total_tokens += len(ids)
                batch = []
        #Flush remaining lines
        if batch:
            encodings = tokenizer.encode_batch(batch)
            for enc in encodings:
                ids = np.array(enc.ids, dtype=DTYPE)
                ids.tofile(f_tmp)
                total_tokens += len(ids)
    #Rename tmp to final
    os.rename(tmp_path, out_path)
    print(f"  {total_tokens} tokens written to {out_path}")
    return total_tokens

for split in ["train", "validation", "test"]:
    txt_path = os.path.join(DATA_DIR, f"{split}.txt")
    out_path = os.path.join(OUT_DIR, f"{split}.bin")
    n = encode_file(txt_path, out_path)

    with open(out_path + ".count", "w") as f:
        f.write(str(n))

print("\nAll splits encoded")