{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SLM Training — Curriculum Learning\n",
    "All training commands live here. All logic is in the modules.\n",
    "Run cells in order for a full curriculum run, or jump to any stage individually."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install torch tokenizers datasets pyyaml --upgrade -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# Login using e.g. `huggingface-cli login` to access this dataset\n",
    "ds = load_dataset(\"BabyLM-community/babylm-eng\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import login\n",
    "login(\"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\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Train Tokenizer A — Corpus-derived vocab (~32-40k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[tokenizer] Training TokenizerA (corpus) | vocab_size=40,000\n",
      "[00:00:00] Pre-processing sequences       ██████████████████ 0        /        0[tokenizer] Sampling 500,000 docs from each source (2,000,000 total)\n",
      "Resolving data files: 100%|██████████████| 2410/2410 [00:00<00:00, 84041.61it/s]\n",
      "\u001b[2K[00:00:01] Tokenize words                 ██████████████████ 2260621  /  2260621[00:00:00] Tokenize words                 ██████████████████ 0        /        0\n",
      "\u001b[2K[00:00:01] Count pairs                    ██████████████████ 2260621  /  2260621\n",
      "\u001b[2K[00:00:05] Compute merges                 ██████████████████ 39786    /    39786\n",
      "[tokenizer] Saved TokenizerA (corpus) → tokenizers/tokenizer_corpus.json\n",
      "[tokenizer] Actual vocab size: 40,000\n",
      "\n",
      "[tokenizer] Done. Files saved in: tokenizers/\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.chdir('/home/user20/NLP/slm')  # adjust to your path\n",
    "\n",
    "!python tokenizer.py \\\n",
    "    --output_dir tokenizers/ \\\n",
    "    --sample_size 2_000_000 \\\n",
    "    --which corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A:\tOkay.\n",
      "A:\tSo, What kind of experience do you, do you have, then with child car\n",
      "Well it's just that, you know, a pound, or a hundred pounds today, is not the sa\n",
      "You want me to start again?\n",
      "Yeah.\n",
      "Right erm, could you tell me about how you lef\n",
      "Is this on yet?\n",
      "Yeah.\n",
      "Oh.\n",
      "Okay well, good morning.\n",
      "Erm I have a er an important \n",
      "Come in, good morning.\n",
      "Hello, well what's your mum been doing to you this mornin\n"
     ]
    }
   ],
   "source": [
    "from tokenizer import iter_babylm\n",
    "\n",
    "for text in iter_babylm(max_docs=5):\n",
    "    print(text[:80])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vocab size: 40000\n",
      "Tokens: ['ĠThe', 'Ġquick', 'Ġbrown', 'Ġfox', 'Ġjumps', 'Ġover', 'Ġthe', 'Ġlazy', 'Ġdog', '.']\n",
      "IDs: [281, 1772, 3733, 4510, 12121, 671, 217, 9818, 1117, 17]\n"
     ]
    }
   ],
   "source": [
    "from tokenizers import Tokenizer\n",
    "\n",
    "# Load your new tokenizer\n",
    "tokenizer = Tokenizer.from_file(\"tokenizers/tokenizer_corpus.json\")\n",
    "\n",
    "test_text = \"The quick brown fox jumps over the lazy dog.\"\n",
    "encoded = tokenizer.encode(test_text)\n",
    "\n",
    "print(f\"Vocab size: {tokenizer.get_vocab_size()}\")\n",
    "print(f\"Tokens: {encoded.tokens}\")\n",
    "print(f\"IDs: {encoded.ids}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Train Tokenizer B — Fixed 50k vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[tokenizer] Training TokenizerB (50k) | vocab_size=50,000\n",
      "[00:00:00] Pre-processing sequences       ██████████████████ 0        /        0[tokenizer] Sampling 500,000 docs from each source (2,000,000 total)\n",
      "Resolving data files: 100%|██████████████| 2410/2410 [00:00<00:00, 84284.07it/s]\n",
      "\u001b[2K[00:00:01] Tokenize words                 ██████████████████ 2260621  /  2260621[00:00:00] Tokenize words                 ██████████████████ 0        /        0\n",
      "\u001b[2K[00:00:01] Count pairs                    ██████████████████ 2260621  /  2260621\n",
      "\u001b[2K[00:00:05] Compute merges                 ██████████████████ 49786    /    49786\n",
      "[tokenizer] Saved TokenizerB (50k) → tokenizers/tokenizer_50k.json\n",
      "[tokenizer] Actual vocab size: 50,000\n",
      "\n",
      "[tokenizer] Done. Files saved in: tokenizers/\n"
     ]
    }
   ],
   "source": [
    "!python tokenizer.py \\\n",
    "    --output_dir tokenizers/ \\\n",
    "    --sample_size 2_000_000 \\\n",
    "    --which fixed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Inspect tokenizer vocab sizes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Corpus: vocab_size=40,000\n",
      "  Sample encode: [772, 756, 216, 487, 15, 216, 566, 720, 2422, 650, 218, 2021, 17]\n",
      "  Decode back  :  Once upon a time, a little girl walked into the forest.\n",
      "\n",
      "50k: vocab_size=50,000\n",
      "  Sample encode: [772, 756, 216, 487, 15, 216, 566, 720, 2422, 650, 218, 2021, 17]\n",
      "  Decode back  :  Once upon a time, a little girl walked into the forest.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from tokenizers import Tokenizer\n",
    "\n",
    "for name, path in [('Corpus', 'tokenizers/tokenizer_corpus.json'),\n",
    "                   ('50k',    'tokenizers/tokenizer_50k.json')]:\n",
    "    tok = Tokenizer.from_file(path)\n",
    "    print(f'{name}: vocab_size={tok.get_vocab_size():,}')\n",
    "    sample = 'Once upon a time, a little girl walked into the forest.'\n",
    "    enc    = tok.encode(sample)\n",
    "    print(f'  Sample encode: {enc.ids}')\n",
    "    print(f'  Decode back  : {tok.decode(enc.ids)}')\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Inspect model param counts for both pos_type options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tokenizer=50k | pos=learnable | total=51.03M\n",
      "  tok_emb   : 25.60M\n",
      "  pos_emb   : 0.26M\n",
      "  rope      : 0.00M\n",
      "  layers    : 25.17M\n",
      "  norm_out  : 0.00M\n",
      "  lm_head   : 0.00M\n",
      "\n",
      "tokenizer=50k | pos=rope | total=50.77M\n",
      "  tok_emb   : 25.60M\n",
      "  pos_emb   : 0.00M\n",
      "  rope      : 0.00M\n",
      "  layers    : 25.17M\n",
      "  norm_out  : 0.00M\n",
      "  lm_head   : 0.00M\n",
      "\n",
      "tokenizer=corpus | pos=learnable | total=45.91M\n",
      "  tok_emb   : 20.48M\n",
      "  pos_emb   : 0.26M\n",
      "  rope      : 0.00M\n",
      "  layers    : 25.17M\n",
      "  norm_out  : 0.00M\n",
      "  lm_head   : 0.00M\n",
      "\n",
      "tokenizer=corpus | pos=rope | total=45.65M\n",
      "  tok_emb   : 20.48M\n",
      "  pos_emb   : 0.00M\n",
      "  rope      : 0.00M\n",
      "  layers    : 25.17M\n",
      "  norm_out  : 0.00M\n",
      "  lm_head   : 0.00M\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from tokenizers import Tokenizer\n",
    "from model import SLM, SLMConfig\n",
    "\n",
    "for tok_name, tok_path in [('50k', 'tokenizers/tokenizer_50k.json'),\n",
    "                            ('corpus', 'tokenizers/tokenizer_corpus.json')]:\n",
    "    tok = Tokenizer.from_file(tok_path)\n",
    "    for pos in ('learnable', 'rope'):\n",
    "        cfg   = SLMConfig(vocab_size=tok.get_vocab_size(), pos_type=pos)\n",
    "        model = SLM(cfg)\n",
    "        bd    = model.param_breakdown()\n",
    "        print(f'tokenizer={tok_name} | pos={pos} | total={bd[\"total\"]/1e6:.2f}M')\n",
    "        for k, v in bd.items():\n",
    "            if k != 'total':\n",
    "                print(f'  {k:<10}: {v/1e6:.2f}M')\n",
    "        print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Stage 0 Training — TinyStories\n",
    "Choose your tokenizer and pos_type, then run.\n",
    "Change `--pos_type rope` to use RoPE instead of learnable embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[train] Stage 0 | device=cuda | seq_len=256\n",
      "[train] Model params: 45.7M\n",
      "[dataset] Loading cached train_stage0\n",
      "[dataset] Loaded 778,211 chunks\n",
      "[dataset] Loading cached val_stage0 from cache/val_stage0_seq256.pkl\n",
      "[dataset] Loading cached val_stage1 from cache/val_stage1_seq384.pkl\n",
      "[dataset] Loading cached val_stage2 from cache/val_stage2_seq512.pkl\n",
      "[train] max_steps=24,414  tokens/step=8,192\n",
      "[logger] Logging to logs/stage0_20260405_050310.csv\n",
      "[Stage 0]:   2%|▍                       | 500/24414 [01:20<1:04:01,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=2.9554)\n",
      "[stage0] step=   500  tok=   4.1M  train=2.9058  val_s0=2.9554  val_s1=8.7646  val_s2=8.4479  lr=3.0e-04\n",
      "[Stage 0]:   4%|▉                      | 1000/24414 [02:52<1:02:43,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=2.4682)\n",
      "[stage0] step=  1000  tok=   8.2M  train=2.3591  val_s0=2.4682  val_s1=8.8909  val_s2=8.3747  lr=3.0e-04\n",
      "[Stage 0]:   6%|█▍                     | 1500/24414 [04:24<1:01:21,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=2.2675)\n",
      "[stage0] step=  1500  tok=  12.3M  train=2.2584  val_s0=2.2675  val_s1=8.7414  val_s2=8.2340  lr=3.0e-04\n",
      "[Stage 0]:   8%|██                       | 2000/24414 [05:55<59:59,  6.23step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=2.1491)\n",
      "[stage0] step=  2000  tok=  16.4M  train=2.2747  val_s0=2.1491  val_s1=8.6993  val_s2=8.2744  lr=3.0e-04\n",
      "[Stage 0]:  10%|██▌                      | 2500/24414 [07:27<58:39,  6.23step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=2.0718)\n",
      "[stage0] step=  2500  tok=  20.5M  train=2.0581  val_s0=2.0718  val_s1=8.8742  val_s2=8.3755  lr=3.0e-04\n",
      "[Stage 0]:  12%|███                      | 3000/24414 [08:59<57:20,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=2.0149)\n",
      "[stage0] step=  3000  tok=  24.6M  train=1.8970  val_s0=2.0149  val_s1=8.9037  val_s2=8.4539  lr=2.9e-04\n",
      "[Stage 0]:  14%|███▌                     | 3500/24414 [10:31<56:00,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.9682)\n",
      "[stage0] step=  3500  tok=  28.7M  train=1.9449  val_s0=1.9682  val_s1=8.8163  val_s2=8.3171  lr=2.9e-04\n",
      "[Stage 0]:  16%|████                     | 4000/24414 [12:03<54:38,  6.23step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.9363)\n",
      "[stage0] step=  4000  tok=  32.8M  train=1.8734  val_s0=1.9363  val_s1=8.6370  val_s2=8.2168  lr=2.9e-04\n",
      "[Stage 0]:  18%|████▌                    | 4500/24414 [13:35<53:19,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.9097)\n",
      "[stage0] step=  4500  tok=  36.9M  train=1.9686  val_s0=1.9097  val_s1=8.8109  val_s2=8.3591  lr=2.8e-04\n",
      "[Stage 0]:  20%|█████                    | 5000/24414 [15:07<52:00,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.8855)\n",
      "[stage0] step=  5000  tok=  41.0M  train=1.9416  val_s0=1.8855  val_s1=8.5115  val_s2=8.0946  lr=2.8e-04\n",
      "[Stage 0]:  23%|█████▋                   | 5500/24414 [16:39<50:38,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.8638)\n",
      "[stage0] step=  5500  tok=  45.1M  train=1.7275  val_s0=1.8638  val_s1=8.6389  val_s2=8.2657  lr=2.7e-04\n",
      "[Stage 0]:  25%|██████▏                  | 6000/24414 [18:11<49:17,  6.23step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.8476)\n",
      "[stage0] step=  6000  tok=  49.2M  train=1.8883  val_s0=1.8476  val_s1=8.5566  val_s2=8.2463  lr=2.7e-04\n",
      "[Stage 0]:  27%|██████▋                  | 6500/24414 [19:43<47:56,  6.23step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.8308)\n",
      "[stage0] step=  6500  tok=  53.3M  train=1.8003  val_s0=1.8308  val_s1=8.4476  val_s2=8.1404  lr=2.6e-04\n",
      "[Stage 0]:  29%|███████▏                 | 7000/24414 [21:15<46:37,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.8131)\n",
      "[stage0] step=  7000  tok=  57.4M  train=1.6856  val_s0=1.8131  val_s1=8.4028  val_s2=8.1105  lr=2.5e-04\n",
      "[Stage 0]:  31%|███████▋                 | 7500/24414 [22:46<45:16,  6.23step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.7995)\n",
      "[stage0] step=  7500  tok=  61.4M  train=1.7266  val_s0=1.7995  val_s1=8.4178  val_s2=8.0647  lr=2.5e-04\n",
      "[Stage 0]:  33%|████████▏                | 8000/24414 [24:18<43:55,  6.23step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.7857)\n",
      "[stage0] step=  8000  tok=  65.5M  train=1.7736  val_s0=1.7857  val_s1=8.4972  val_s2=8.0946  lr=2.4e-04\n",
      "[Stage 0]:  35%|████████▋                | 8500/24414 [25:50<42:38,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.7740)\n",
      "[stage0] step=  8500  tok=  69.6M  train=1.7707  val_s0=1.7740  val_s1=8.5627  val_s2=8.0705  lr=2.3e-04\n",
      "[Stage 0]:  37%|█████████▏               | 9000/24414 [27:22<41:16,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.7632)\n",
      "[stage0] step=  9000  tok=  73.7M  train=1.7982  val_s0=1.7632  val_s1=8.6377  val_s2=8.1868  lr=2.2e-04\n",
      "[Stage 0]:  39%|█████████▋               | 9500/24414 [28:54<39:54,  6.23step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.7540)\n",
      "[stage0] step=  9500  tok=  77.8M  train=1.7653  val_s0=1.7540  val_s1=8.5986  val_s2=8.1109  lr=2.2e-04\n",
      "[Stage 0]:  41%|█████████▊              | 10000/24414 [30:26<38:35,  6.23step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.7408)\n",
      "[stage0] step= 10000  tok=  81.9M  train=1.6761  val_s0=1.7408  val_s1=8.5158  val_s2=8.0448  lr=2.1e-04\n",
      "[Stage 0]:  43%|██████████▎             | 10500/24414 [31:58<37:14,  6.23step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.7314)\n",
      "[stage0] step= 10500  tok=  86.0M  train=1.6593  val_s0=1.7314  val_s1=8.7288  val_s2=8.1726  lr=2.0e-04\n",
      "[Stage 0]:  45%|██████████▊             | 11000/24414 [33:30<35:54,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.7204)\n",
      "[stage0] step= 11000  tok=  90.1M  train=1.7813  val_s0=1.7204  val_s1=8.4062  val_s2=7.9594  lr=1.9e-04\n",
      "[Stage 0]:  47%|███████████▎            | 11500/24414 [35:02<34:34,  6.23step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.7133)\n",
      "[stage0] step= 11500  tok=  94.2M  train=1.7505  val_s0=1.7133  val_s1=8.3821  val_s2=7.9584  lr=1.8e-04\n",
      "[Stage 0]:  49%|███████████▊            | 12000/24414 [36:34<33:13,  6.23step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.7013)\n",
      "[stage0] step= 12000  tok=  98.3M  train=1.7825  val_s0=1.7013  val_s1=8.4702  val_s2=8.0006  lr=1.7e-04\n",
      "[Stage 0]:  51%|████████████▎           | 12500/24414 [38:06<31:54,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6917)\n",
      "[stage0] step= 12500  tok= 102.4M  train=1.7525  val_s0=1.6917  val_s1=8.4782  val_s2=8.0714  lr=1.6e-04\n",
      "[Stage 0]:  53%|████████████▊           | 13000/24414 [39:38<30:33,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6842)\n",
      "[stage0] step= 13000  tok= 106.5M  train=1.5805  val_s0=1.6842  val_s1=8.5747  val_s2=8.0857  lr=1.6e-04\n",
      "[Stage 0]:  54%|████████████▉           | 13142/24414 [40:13<34:29,  5.45step/s]\n",
      "\n",
      "[logger] Stage 0 exit at step 13142 (107.7M tokens): loss_spike\n",
      "\n",
      "[train] Stage 0 complete. Best val: 1.6842\n",
      "[train] Best checkpoint: checkpoints/stage0_best.pt\n"
     ]
    }
   ],
   "source": [
    "TOKENIZER = 'tokenizers/tokenizer_corpus.json'   # or tokenizer_corpus.json\n",
    "POS_TYPE  = 'rope'                        # or 'rope'\n",
    "\n",
    "!python train.py \\\n",
    "    --stage 0 \\\n",
    "    --config configs/stage0.yaml \\\n",
    "    --tokenizer {TOKENIZER} \\\n",
    "    --pos_type {POS_TYPE} \\\n",
    "    --checkpoint_dir checkpoints/ \\\n",
    "    --log_dir logs/ \\\n",
    "    --cache_dir cache/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Resume Stage 0 (if interrupted)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[train] Stage 0 | device=cuda | seq_len=256\n",
      "[train] Model params: 45.7M\n",
      "[train] Resumed from checkpoints/stage0_best.pt  (step=13000, val=1.6842)\n",
      "[dataset] Loading cached train_stage0\n",
      "[dataset] Loaded 778,211 chunks\n",
      "[dataset] Loading cached val_stage0 from cache/val_stage0_seq256.pkl\n",
      "[dataset] Loading cached val_stage1 from cache/val_stage1_seq384.pkl\n",
      "[dataset] Loading cached val_stage2 from cache/val_stage2_seq512.pkl\n",
      "[train] max_steps=24,414  tokens/step=8,192\n",
      "[logger] Logging to logs/stage0_20260405_060014.csv\n",
      "[Stage 0]:  53%|█████████████████               | 13000/24414 [00:00<?, ?step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6837)\n",
      "[stage0] step= 13000  tok= 106.5M  train=1.6035  val_s0=1.6837  val_s1=8.5629  val_s2=8.0885  lr=1.6e-04\n",
      "[Stage 0]:  55%|█████████████▎          | 13500/24414 [01:32<29:15,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6759)\n",
      "[stage0] step= 13500  tok= 110.6M  train=1.6030  val_s0=1.6759  val_s1=8.4677  val_s2=7.9823  lr=1.5e-04\n",
      "[Stage 0]:  57%|█████████████▊          | 14000/24414 [03:04<27:54,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6667)\n",
      "[stage0] step= 14000  tok= 114.7M  train=1.6712  val_s0=1.6667  val_s1=8.5672  val_s2=8.0426  lr=1.4e-04\n",
      "[Stage 0]:  59%|██████████████▎         | 14500/24414 [04:36<26:34,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6583)\n",
      "[stage0] step= 14500  tok= 118.8M  train=1.5603  val_s0=1.6583  val_s1=8.6355  val_s2=8.1378  lr=1.3e-04\n",
      "[Stage 0]:  61%|██████████████▋         | 15000/24414 [06:08<25:15,  6.21step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6531)\n",
      "[stage0] step= 15000  tok= 122.9M  train=1.5726  val_s0=1.6531  val_s1=8.4124  val_s2=7.9499  lr=1.2e-04\n",
      "[Stage 0]:  63%|███████████████▏        | 15500/24414 [07:40<23:53,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6456)\n",
      "[stage0] step= 15500  tok= 127.0M  train=1.7267  val_s0=1.6456  val_s1=8.3302  val_s2=7.9189  lr=1.1e-04\n",
      "[Stage 0]:  66%|███████████████▋        | 16000/24414 [09:12<22:33,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6400)\n",
      "[stage0] step= 16000  tok= 131.1M  train=1.5696  val_s0=1.6400  val_s1=8.4188  val_s2=7.9356  lr=1.0e-04\n",
      "[Stage 0]:  68%|████████████████▏       | 16500/24414 [10:44<21:13,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6323)\n",
      "[stage0] step= 16500  tok= 135.2M  train=1.5312  val_s0=1.6323  val_s1=8.3179  val_s2=7.9415  lr=9.7e-05\n",
      "[Stage 0]:  70%|████████████████▋       | 17000/24414 [12:20<29:35,  4.18step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6278)\n",
      "[stage0] step= 17000  tok= 139.3M  train=1.4691  val_s0=1.6278  val_s1=8.3390  val_s2=7.8624  lr=8.9e-05\n",
      "[Stage 0]:  72%|█████████████████▏      | 17500/24414 [13:53<18:32,  6.21step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6204)\n",
      "[stage0] step= 17500  tok= 143.4M  train=1.5974  val_s0=1.6204  val_s1=8.4206  val_s2=7.9288  lr=8.2e-05\n",
      "[Stage 0]:  74%|█████████████████▋      | 18000/24414 [15:26<17:11,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6154)\n",
      "[stage0] step= 18000  tok= 147.5M  train=1.5915  val_s0=1.6154  val_s1=8.2098  val_s2=7.8690  lr=7.5e-05\n",
      "[Stage 0]:  76%|██████████████████▏     | 18500/24414 [17:00<15:51,  6.21step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6084)\n",
      "[stage0] step= 18500  tok= 151.6M  train=1.4574  val_s0=1.6084  val_s1=8.2393  val_s2=7.8641  lr=6.9e-05\n",
      "[Stage 0]:  78%|██████████████████▋     | 19000/24414 [18:33<14:30,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.6053)\n",
      "[stage0] step= 19000  tok= 155.7M  train=1.7011  val_s0=1.6053  val_s1=8.2315  val_s2=7.8786  lr=6.3e-05\n",
      "[Stage 0]:  80%|███████████████████▏    | 19500/24414 [20:06<13:11,  6.21step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.5996)\n",
      "[stage0] step= 19500  tok= 159.8M  train=1.5699  val_s0=1.5996  val_s1=8.2321  val_s2=7.8438  lr=5.7e-05\n",
      "[Stage 0]:  82%|███████████████████▋    | 20000/24414 [21:41<11:51,  6.21step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.5961)\n",
      "[stage0] step= 20000  tok= 163.9M  train=1.6109  val_s0=1.5961  val_s1=8.2143  val_s2=7.8210  lr=5.2e-05\n",
      "[Stage 0]:  84%|████████████████████▏   | 20500/24414 [23:17<10:30,  6.21step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.5924)\n",
      "[stage0] step= 20500  tok= 168.0M  train=1.4827  val_s0=1.5924  val_s1=8.2494  val_s2=7.8202  lr=4.7e-05\n",
      "[Stage 0]:  86%|████████████████████▋   | 21000/24414 [24:51<09:09,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.5878)\n",
      "[stage0] step= 21000  tok= 172.0M  train=1.4441  val_s0=1.5878  val_s1=8.2117  val_s2=7.8232  lr=4.3e-05\n",
      "[Stage 0]:  88%|█████████████████████▏  | 21500/24414 [26:24<07:48,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.5851)\n",
      "[stage0] step= 21500  tok= 176.1M  train=1.4431  val_s0=1.5851  val_s1=8.2389  val_s2=7.8500  lr=4.0e-05\n",
      "[Stage 0]:  90%|█████████████████████▋  | 22000/24414 [27:56<06:28,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.5806)\n",
      "[stage0] step= 22000  tok= 180.2M  train=1.4599  val_s0=1.5806  val_s1=8.1978  val_s2=7.7960  lr=3.7e-05\n",
      "[Stage 0]:  92%|██████████████████████  | 22500/24414 [29:28<05:07,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.5795)\n",
      "[stage0] step= 22500  tok= 184.3M  train=1.5260  val_s0=1.5795  val_s1=8.2558  val_s2=7.8620  lr=3.4e-05\n",
      "[Stage 0]:  94%|██████████████████████▌ | 23000/24414 [31:00<03:47,  6.21step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.5762)\n",
      "[stage0] step= 23000  tok= 188.4M  train=1.5518  val_s0=1.5762  val_s1=8.1968  val_s2=7.8414  lr=3.2e-05\n",
      "[Stage 0]:  96%|███████████████████████ | 23500/24414 [32:33<02:27,  6.21step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.5744)\n",
      "[stage0] step= 23500  tok= 192.5M  train=1.5555  val_s0=1.5744  val_s1=8.2230  val_s2=7.8338  lr=3.1e-05\n",
      "[Stage 0]:  98%|███████████████████████▌| 24000/24414 [34:09<01:06,  6.22step/s][train] Checkpoint saved → checkpoints/stage0_best.pt  (val=1.5726)\n",
      "[stage0] step= 24000  tok= 196.6M  train=1.6689  val_s0=1.5726  val_s1=8.2236  val_s2=7.8677  lr=3.0e-05\n",
      "[Stage 0]: 100%|████████████████████████| 24414/24414 [35:27<00:00,  5.36step/s]\n",
      "\n",
      "[logger] Stage 0 exit at step 24414 (200.0M tokens): token_budget\n",
      "\n",
      "[train] Stage 0 complete. Best val: 1.5726\n",
      "[train] Best checkpoint: checkpoints/stage0_best.pt\n"
     ]
    }
   ],
   "source": [
    "!python train.py \\\n",
    "    --stage 0 \\\n",
    "    --config configs/stage0.yaml \\\n",
    "    --tokenizer {TOKENIZER} \\\n",
    "    --pos_type {POS_TYPE} \\\n",
    "    --checkpoint_dir checkpoints/ \\\n",
    "    --log_dir logs/ \\\n",
    "    --cache_dir cache/ \\\n",
    "    --resume"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Stage 0b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre-build Dataset Caches (ONE TIME ONLY)\n",
    "Run this cell ONCE before training any stage. Caches data so training starts instantly.\n",
    "Takes ~30-45 minutes total but all future training is instant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======================================================================\n",
      "PRE-BUILDING DATASET CACHES FOR ALL STAGES\n",
      "======================================================================\n",
      "✓ tinystories     @ seq_len=256:    129 MB (cached)\n",
      "✓ babylm_easy     @ seq_len=384:    200 MB (cached)\n",
      "✓ tinystories     @ seq_len=384:    129 MB (cached)\n",
      "⏳ babylm_hard     @ seq_len=512... [dataset] Building train 'babylm_hard' (seq=512, max_tokens=60,000,000)...\n",
      "`trust_remote_code` is not supported anymore.\n",
      "Please check that the Hugging Face dataset 'BabyLM-community/babylm-eng' isn't based on a loading script and remove `trust_remote_code`.\n",
      "If the dataset is based on a loading script, please ask the dataset author to remove it and convert it to a standard format like Parquet.\n",
      "[dataset] WARNING: BabyLM has no source field — yielding all text\n",
      "[dataset] Saved 116,960 chunks → cache/train_babylm_hard_seq512.pkl\n",
      " 116,960 chunks →    152 MB ✓\n",
      "⏳ babylm_easy     @ seq_len=512... [dataset] Building train 'babylm_easy' (seq=512, max_tokens=50,000,000)...\n",
      "`trust_remote_code` is not supported anymore.\n",
      "Please check that the Hugging Face dataset 'BabyLM-community/babylm-eng' isn't based on a loading script and remove `trust_remote_code`.\n",
      "If the dataset is based on a loading script, please ask the dataset author to remove it and convert it to a standard format like Parquet.\n",
      "[dataset] WARNING: BabyLM has no source field — yielding all text\n",
      "[dataset] Saved 97,466 chunks → cache/train_babylm_easy_seq512.pkl\n",
      "  97,466 chunks →    129 MB ✓\n",
      "✓ tinystories     @ seq_len=512:    129 MB (cached)\n",
      "⏳ simplewiki      @ seq_len=384... [dataset] Building train 'simplewiki' (seq=384, max_tokens=220,000,000)...\n",
      "[dataset] Saved 178,373 chunks → cache/train_simplewiki_seq384.pkl\n",
      " 178,373 chunks →    176 MB ✓\n",
      "⏳ babylm_hard     @ seq_len=384... [dataset] Building train 'babylm_hard' (seq=384, max_tokens=60,000,000)...\n",
      "`trust_remote_code` is not supported anymore.\n",
      "Please check that the Hugging Face dataset 'BabyLM-community/babylm-eng' isn't based on a loading script and remove `trust_remote_code`.\n",
      "If the dataset is based on a loading script, please ask the dataset author to remove it and convert it to a standard format like Parquet.\n",
      "[dataset] WARNING: BabyLM has no source field — yielding all text\n",
      "[dataset] Saved 155,845 chunks → cache/train_babylm_hard_seq384.pkl\n",
      " 155,845 chunks →    153 MB ✓\n",
      "⏳ fineweb_edu     @ seq_len=512... [dataset] Building train 'fineweb_edu' (seq=512, max_tokens=500,000,000)...\n",
      "Resolving data files: 100%|██████████████| 2410/2410 [00:00<00:00, 83893.74it/s]\n",
      "[dataset] Saved 974,659 chunks → cache/train_fineweb_edu_seq512.pkl\n",
      " 974,659 chunks →   1295 MB ✓\n",
      "⏳ simplewiki      @ seq_len=512... [dataset] Building train 'simplewiki' (seq=512, max_tokens=220,000,000)...\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/user20/NLP/slm/prebuild_cache.py\", line 60, in <module>\n",
      "    ds = StreamingStageDataset().build(\n",
      "         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/user20/NLP/slm/dataset.py\", line 406, in build\n",
      "    max_tokens,\n",
      "               \n",
      "  File \"/home/user20/NLP/slm/dataset.py\", line 222, in tokenize_and_chunk\n",
      "    chunks.append(buffer[: seq_len + 1])\n",
      "                 ^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "^C\n"
     ]
    }
   ],
   "source": [
    "!python prebuild_cache.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[train] Stage 0 | device=cuda | seq_len=384\n",
      "[train] Model params: 45.7M\n",
      "[train] Loading weights from prev stage: checkpoints/stage0_best_1stfull.pt\n",
      "[dataset] Loading cached train 'babylm_easy' @ seq=384\n",
      "[dataset] Loaded 207,793 chunks\n",
      "[dataset] Loading replay source 'tinystories' from cache/train_tinystories_seq384.pkl (exact match)\n",
      "^C\n"
     ]
    }
   ],
   "source": [
    "!python train.py \\\n",
    "    --stage 0 \\\n",
    "    --config configs/stage0b.yaml \\\n",
    "    --tokenizer tokenizers/tokenizer_corpus.json \\\n",
    "    --pos_type rope \\\n",
    "    --checkpoint_dir checkpoints/ \\\n",
    "    --log_dir logs/ \\\n",
    "    --cache_dir cache/ \\\n",
    "    --prev_checkpoint checkpoints/stage0_best_1stfull.pt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Stage 1 Training — SimpleWiki + BabyLM\n",
    "Loads Stage 0 best checkpoint as starting weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python train.py \\\n",
    "    --stage 1 \\\n",
    "    --config configs/stage1.yaml \\\n",
    "    --tokenizer {TOKENIZER} \\\n",
    "    --pos_type {POS_TYPE} \\\n",
    "    --checkpoint_dir checkpoints/ \\\n",
    "    --log_dir logs/ \\\n",
    "    --cache_dir cache/ \\\n",
    "    --prev_checkpoint checkpoints/stage0_best.pt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Resume Stage 1 (if interrupted)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python train.py \\\n",
    "    --stage 1 \\\n",
    "    --config configs/stage1.yaml \\\n",
    "    --tokenizer {TOKENIZER} \\\n",
    "    --pos_type {POS_TYPE} \\\n",
    "    --checkpoint_dir checkpoints/ \\\n",
    "    --log_dir logs/ \\\n",
    "    --cache_dir cache/ \\\n",
    "    --resume"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Stage 2 Training — FineWeb-Edu\n",
    "Loads Stage 1 best checkpoint as starting weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python train.py \\\n",
    "    --stage 2 \\\n",
    "    --config configs/stage2.yaml \\\n",
    "    --tokenizer {TOKENIZER} \\\n",
    "    --pos_type {POS_TYPE} \\\n",
    "    --checkpoint_dir checkpoints/ \\\n",
    "    --log_dir logs/ \\\n",
    "    --cache_dir cache/ \\\n",
    "    --prev_checkpoint checkpoints/stage1_best.pt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. (Alternative) Run full curriculum automatically\n",
    "Runs all stages in sequence, passing checkpoints forward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python curriculum.py \\\n",
    "    --tokenizer {TOKENIZER} \\\n",
    "    --pos_type {POS_TYPE} \\\n",
    "    --checkpoint_dir checkpoints/ \\\n",
    "    --log_dir logs/ \\\n",
    "    --cache_dir cache/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Plot training curves — all stages, all val losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, axes = plt.subplots(1, 3, figsize=(18, 5))\n",
    "\n",
    "for stage_idx, ax in enumerate(axes):\n",
    "    files = sorted(glob.glob(f'logs/stage{stage_idx}_*.csv'))\n",
    "    if not files:\n",
    "        ax.set_title(f'Stage {stage_idx} — no logs yet')\n",
    "        continue\n",
    "\n",
    "    df = pd.read_csv(files[-1])   # most recent run\n",
    "    df = df.dropna(subset=['train_loss'])\n",
    "    df['val_s0'] = pd.to_numeric(df['val_s0'], errors='coerce')\n",
    "    df['val_s1'] = pd.to_numeric(df['val_s1'], errors='coerce')\n",
    "    df['val_s2'] = pd.to_numeric(df['val_s2'], errors='coerce')\n",
    "    df['train_loss'] = pd.to_numeric(df['train_loss'], errors='coerce')\n",
    "\n",
    "    ax.plot(df['step'], df['train_loss'], label='train',  alpha=0.7, linewidth=1.5)\n",
    "    ax.plot(df['step'], df['val_s0'],    label='val_s0', alpha=0.9, linewidth=1.5)\n",
    "    ax.plot(df['step'], df['val_s1'],    label='val_s1', alpha=0.9, linewidth=1.5)\n",
    "    ax.plot(df['step'], df['val_s2'],    label='val_s2', alpha=0.9, linewidth=1.5)\n",
    "\n",
    "    ax.set_title(f'Stage {stage_idx}')\n",
    "    ax.set_xlabel('step')\n",
    "    ax.set_ylabel('loss')\n",
    "    ax.legend(fontsize=8)\n",
    "    ax.grid(True, alpha=0.3)\n",
    "\n",
    "plt.suptitle('Training Curves — All Stages', fontsize=14)\n",
    "plt.tight_layout()\n",
    "plt.savefig('logs/training_curves.png', dpi=150)\n",
    "plt.show()\n",
    "print('Saved: logs/training_curves.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. Generate sample text (inference test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from tokenizers import Tokenizer\n",
    "from model import SLM, SLMConfig\n",
    "\n",
    "CKPT      = 'checkpoints/stage2_best.pt'   # or stage0/stage1\n",
    "TOKENIZER = 'tokenizers/tokenizer_50k.json'\n",
    "PROMPT    = 'Once upon a time'\n",
    "MAX_NEW   = 100\n",
    "USE_CACHE = True   # set False to benchmark without KV cache\n",
    "\n",
    "device    = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "tokenizer = Tokenizer.from_file(TOKENIZER)\n",
    "ckpt      = torch.load(CKPT, map_location=device)\n",
    "\n",
    "model = SLM(ckpt['config']).to(device)\n",
    "model.load_state_dict(ckpt['model_state'])\n",
    "model.eval()\n",
    "\n",
    "ids    = tokenizer.encode(PROMPT).ids\n",
    "inp    = torch.tensor([ids], device=device)\n",
    "out    = model.generate(inp, max_new=MAX_NEW, temperature=0.8, top_k=50, use_cache=USE_CACHE)\n",
    "result = tokenizer.decode(out[0].tolist())\n",
    "\n",
    "print('Generated:')\n",
    "print('-' * 60)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. Benchmark: KV cache vs no cache (inference speed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import torch\n",
    "from tokenizers import Tokenizer\n",
    "from model import SLM\n",
    "\n",
    "CKPT      = 'checkpoints/stage2_best.pt'\n",
    "TOKENIZER = 'tokenizers/tokenizer_50k.json'\n",
    "MAX_NEW   = 50\n",
    "RUNS      = 3\n",
    "\n",
    "device    = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "tokenizer = Tokenizer.from_file(TOKENIZER)\n",
    "ckpt      = torch.load(CKPT, map_location=device)\n",
    "model     = SLM(ckpt['config']).to(device)\n",
    "model.load_state_dict(ckpt['model_state'])\n",
    "model.eval()\n",
    "\n",
    "ids    = tokenizer.encode('Once upon a time in a faraway land').ids\n",
    "prompt = torch.tensor([ids], device=device)\n",
    "\n",
    "results = {}\n",
    "for use_cache in (True, False):\n",
    "    times = []\n",
    "    for _ in range(RUNS):\n",
    "        torch.cuda.synchronize() if device == 'cuda' else None\n",
    "        t0 = time.time()\n",
    "        model.generate(prompt, max_new=MAX_NEW, use_cache=use_cache)\n",
    "        torch.cuda.synchronize() if device == 'cuda' else None\n",
    "        times.append(time.time() - t0)\n",
    "    results[use_cache] = sum(times) / len(times)\n",
    "\n",
    "print(f'KV cache ON  : {results[True]:.3f}s for {MAX_NEW} tokens')\n",
    "print(f'KV cache OFF : {results[False]:.3f}s for {MAX_NEW} tokens')\n",
    "print(f'Speedup      : {results[False]/results[True]:.2f}x')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nlp_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}