{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":121150,"databundleVersionId":15833391}],"dockerImageVersionId":31259,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\"\"\"\nprepare_data.py — Run this LOCALLY (on a machine with internet) to produce\ntrain.bin and val.bin for the IGLO comparison notebook.\n\nRequirements (install once):\n    pip install datasets tiktoken numpy tqdm\n\nWhat it does:\n    1. Downloads TinyStories from Hugging Face (~2 GB raw text).\n    2. Tokenizes with the GPT-2 BPE tokenizer (50,257 vocab).\n    3. Writes train.bin and val.bin as uint16 numpy memmap files.\n       Total output size: roughly 1 GB.\n\nAfter running:\n    Upload train.bin and val.bin to Kaggle as a dataset named\n    'tinystories-gpt2-bin' (or update DATA_DIR in the notebook).\n\"\"\"\n\nimport os\nimport numpy as np\nfrom tqdm import tqdm\nfrom datasets import load_dataset\nimport tiktoken\n\nOUT_DIR = \"./tinystories_bin\"\nNUM_PROC = 4  # adjust to your CPU\nSHARD_BATCH = 1024  # how many docs to process at once when writing\n\nos.makedirs(OUT_DIR, exist_ok=True)\n\nprint(\"Loading TinyStories dataset...\")\nds = load_dataset(\"roneneldan/TinyStories\")\nprint(f\"  train: {len(ds['train']):,} docs\")\nprint(f\"  val  : {len(ds['validation']):,} docs\")\n\nenc = tiktoken.get_encoding(\"gpt2\")\nEOT = enc.eot_token  # 50256\n\ndef tokenize(example):\n    ids = enc.encode_ordinary(example[\"text\"])\n    ids.append(EOT)\n    return {\"ids\": ids, \"len\": len(ids)}\n\nprint(\"\\nTokenizing (this takes a few minutes)...\")\ntokenized = ds.map(\n    tokenize,\n    remove_columns=[\"text\"],\n    desc=\"tokenizing\",\n    num_proc=NUM_PROC,\n)\n\nfor split, dset in tokenized.items():\n    out_name = \"train.bin\" if split == \"train\" else \"val.bin\"\n    path = os.path.join(OUT_DIR, out_name)\n    total = int(np.sum(dset[\"len\"], dtype=np.int64))\n    print(f\"\\n{split}: writing {total:,} tokens to {path}\")\n    arr = np.memmap(path, dtype=np.uint16, mode=\"w+\", shape=(total,))\n    idx = 0\n    for batch_start in tqdm(range(0, len(dset), SHARD_BATCH), desc=f\"writing {split}\"):\n        batch = dset.select(range(batch_start, min(batch_start + SHARD_BATCH, len(dset))))\n        flat = np.concatenate([np.asarray(x, dtype=np.uint16) for x in batch[\"ids\"]])\n        arr[idx : idx + len(flat)] = flat\n        idx += len(flat)\n    arr.flush()\n    size_mb = os.path.getsize(path) / 1e6\n    print(f\"  done. file size: {size_mb:.1f} MB\")\n\nprint(\"\\n\" + \"=\" * 60)\nprint(\"DONE. Now:\")\nprint(\"  1. Go to https://www.kaggle.com/datasets\")\nprint(\"  2. Click 'New Dataset'\")\nprint(f\"  3. Upload {OUT_DIR}/train.bin and {OUT_DIR}/val.bin\")\nprint(\"  4. Name it 'tinystories-gpt2-bin' (or update DATA_DIR in notebook)\")\nprint(\"=\" * 60)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}