{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-00",
   "metadata": {},
   "source": [
    "# Load a pretrained decoder and keep training it\n",
    "\n",
    "[Notebook 05](05-train-a-language-model.ipynb) trains a language model from random weights. Most of the time you start from a model someone else already trained. In this notebook we load SmolLM2-135M, a small Llama-style model from Hugging Face, into Dew's own `CausalTransformer`, continue its training on Shakespeare for a few hundred steps, and export the result back to the Hugging Face format. At the end we load the export with the `transformers` library and check that it predicts the same next token as Dew.\n",
    "\n",
    "The notebook expects one NVIDIA GPU and takes about three minutes on a Colab L4. The first run downloads the checkpoint (about 270 MB)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-01",
   "metadata": {},
   "source": [
    "## Install\n",
    "\n",
    "`interop` adds safetensors support. `torch` is only for the last cell, where `transformers` reads the exported checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cell-02",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-09-22T20:22:06.836591Z",
     "iopub.status.busy": "2026-09-22T20:22:06.836139Z",
     "iopub.status.idle": "2026-09-22T20:22:25.769127Z",
     "shell.execute_reply": "2026-09-22T20:22:25.768533Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n"
     ]
    }
   ],
   "source": [
    "%pip install -q \"dew-ml[interop] @ git+https://github.com/AshishKumar4/dew\" \"jax[cuda12]\" torch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-03",
   "metadata": {},
   "source": [
    "## Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cell-04",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-09-22T20:22:25.772284Z",
     "iopub.status.busy": "2026-09-22T20:22:25.772033Z",
     "iopub.status.idle": "2026-09-22T20:22:25.775647Z",
     "shell.execute_reply": "2026-09-22T20:22:25.775158Z"
    }
   },
   "outputs": [],
   "source": [
    "CHECKPOINT = \"HuggingFaceTB/SmolLM2-135M\"\n",
    "SEQUENCE_LENGTH = 256\n",
    "BATCH_SIZE = 8\n",
    "STEPS = 300\n",
    "LEARNING_RATE = 3e-5\n",
    "MAX_NEW_TOKENS = 60\n",
    "PROMPT = \"ROMEO:\"\n",
    "DATA_DIR = \"data/08-shakespeare\"\n",
    "RUN_DIR = \"runs/08-continued\"\n",
    "EXPORT_DIR = \"runs/08-exported\"\n",
    "SEED = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cell-05",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-09-22T20:22:25.778074Z",
     "iopub.status.busy": "2026-09-22T20:22:25.777688Z",
     "iopub.status.idle": "2026-09-22T20:22:26.905928Z",
     "shell.execute_reply": "2026-09-22T20:22:26.905337Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[CudaDevice(id=0)]\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import os\n",
    "import urllib.request\n",
    "from pathlib import Path\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "\n",
    "print(jax.devices())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-06",
   "metadata": {},
   "source": [
    "## Loading the checkpoint\n",
    "\n",
    "`load_pretrained` takes a Hub repo id or a local folder in the Hugging Face layout. It translates the checkpoint's `config.json` into `CausalTransformer` settings and its weights into Dew's parameter tree. The weights stay in float32 and the model computes in the `dtype` we ask for. If the config has a setting that changes the computation and Dew has no equivalent, `load_pretrained` raises an error naming it instead of silently dropping it.\n",
    "\n",
    "It returns a `Pretrained` with the model on `.model`, its weights on `.variables`, and the Dew settings it derived on `.model_config`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cell-07",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-09-22T20:22:26.909029Z",
     "iopub.status.busy": "2026-09-22T20:22:26.908401Z",
     "iopub.status.idle": "2026-09-22T20:22:54.052867Z",
     "shell.execute_reply": "2026-09-22T20:22:54.052217Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.13/dist-packages/huggingface_hub/utils/_auth.py:138: UserWarning: \n",
      "Error while fetching `HF_TOKEN` secret value from your vault: 'Requesting secret HF_TOKEN timed out. Secrets can only be fetched when running from the Colab UI.'.\n",
      "  warnings.warn(f\"\\nError while fetching `HF_TOKEN` secret value from your vault: '{str(e)}'.\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "135M parameters\n",
      "{'vocab_size': 49152, 'emb_features': 576, 'num_layers': 30, 'num_heads': 9, 'num_kv_heads': 3, 'head_dim': 64}\n"
     ]
    }
   ],
   "source": [
    "from dew.interop import load_pretrained\n",
    "\n",
    "pretrained = load_pretrained(CHECKPOINT, dtype=\"bfloat16\", attention_impl=\"auto\", max_seq_len=512)\n",
    "model, variables = pretrained.model, pretrained.variables\n",
    "n_params = sum(x.size for x in jax.tree_util.tree_leaves(variables))\n",
    "print(f\"{n_params / 1e6:.0f}M parameters\")\n",
    "print({key: pretrained.model_config[key] for key in\n",
    "       (\"vocab_size\", \"emb_features\", \"num_layers\", \"num_heads\", \"num_kv_heads\", \"head_dim\")})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-08",
   "metadata": {},
   "source": [
    "The model was trained on the ids of its own tokenizer, so we must use that tokenizer too. `HFTokenizer` wraps any Hugging Face tokenizer behind the same `encode` and `decode` as the byte tokenizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cell-09",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-09-22T20:22:54.055989Z",
     "iopub.status.busy": "2026-09-22T20:22:54.055206Z",
     "iopub.status.idle": "2026-09-22T20:22:55.851592Z",
     "shell.execute_reply": "2026-09-22T20:22:55.850851Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vocabulary: 49152\n",
      "[3911, 3945, 63, 42]\n"
     ]
    }
   ],
   "source": [
    "from dew.data import HFTokenizer\n",
    "\n",
    "tokenizer = HFTokenizer(CHECKPOINT)\n",
    "print(\"vocabulary:\", tokenizer.vocab_size)\n",
    "print(tokenizer.encode(PROMPT))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-10",
   "metadata": {},
   "source": [
    "## Generating before training\n",
    "\n",
    "SmolLM2 learned from web text, code and textbooks, so it knows what a play looks like but does not write like Shakespeare. We generate greedily, taking the most likely token every time, so the output is the same on every run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cell-11",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-09-22T20:22:55.854606Z",
     "iopub.status.busy": "2026-09-22T20:22:55.853946Z",
     "iopub.status.idle": "2026-09-22T20:23:16.190164Z",
     "shell.execute_reply": "2026-09-22T20:23:16.189520Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ROMEO: I think that's a good point.\n",
      "\n",
      "JOHNSON: I think that's a good point.\n",
      "\n",
      "MARTIN: I think that's a good point.\n",
      "\n",
      "JOHNSON: I think that's a good point.\n",
      "\n",
      "MARTIN: I think\n"
     ]
    }
   ],
   "source": [
    "from dew.sampling import Sampling, generate\n",
    "\n",
    "prompt = jnp.asarray([tokenizer.encode(PROMPT)], jnp.int32)\n",
    "before = generate(model, variables, prompt, max_new_tokens=MAX_NEW_TOKENS,\n",
    "                  key=jax.random.key(0), sampling=Sampling(temperature=0.0))\n",
    "print(tokenizer.decode(before.tokens[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-12",
   "metadata": {},
   "source": [
    "## The data\n",
    "\n",
    "The token files follow the layout from notebook 05, but with SmolLM2's tokenizer. Its vocabulary has 49,152 entries, too many for one byte, so the ids are stored as `uint16`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cell-13",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-09-22T20:23:16.193191Z",
     "iopub.status.busy": "2026-09-22T20:23:16.192700Z",
     "iopub.status.idle": "2026-09-22T20:23:17.362405Z",
     "shell.execute_reply": "2026-09-22T20:23:17.361765Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[transformers] Token indices sequence length is longer than the specified maximum sequence length for this model (341094 > 8192). Running this sequence through the model will result in indexing errors\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tokenizer': 'HuggingFaceTB/SmolLM2-135M', 'vocab_size': 49152, 'dtype': 'uint16', 'train_tokens': 324040, 'val_tokens': 17054, 'eos_id': None}\n"
     ]
    }
   ],
   "source": [
    "data_dir = Path(DATA_DIR)\n",
    "data_dir.mkdir(parents=True, exist_ok=True)\n",
    "text_path = data_dir / \"input.txt\"\n",
    "urllib.request.urlretrieve(\n",
    "    \"https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt\", text_path)\n",
    "\n",
    "ids = np.asarray(tokenizer.encode(text_path.read_text(encoding=\"utf-8\")), np.uint16)\n",
    "val_len = len(ids) // 20\n",
    "ids[:val_len].tofile(data_dir / \"val.bin\")\n",
    "ids[val_len:].tofile(data_dir / \"train.bin\")\n",
    "meta = {\"tokenizer\": CHECKPOINT, \"vocab_size\": tokenizer.vocab_size, \"dtype\": \"uint16\",\n",
    "        \"train_tokens\": len(ids) - val_len, \"val_tokens\": val_len, \"eos_id\": None}\n",
    "(data_dir / \"meta.json\").write_text(json.dumps(meta, indent=2))\n",
    "print(meta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cell-14",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-09-22T20:23:17.365303Z",
     "iopub.status.busy": "2026-09-22T20:23:17.364820Z",
     "iopub.status.idle": "2026-09-22T20:23:17.371085Z",
     "shell.execute_reply": "2026-09-22T20:23:17.370553Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "training windows: 1265\n"
     ]
    }
   ],
   "source": [
    "from dew.data import Loading, TokenWindows\n",
    "\n",
    "data = TokenWindows(path=DATA_DIR, seq_len=SEQUENCE_LENGTH, val_batches=8,\n",
    "                    loading=Loading(workers=0, threads=1, read_buffer=2)).load(batch=BATCH_SIZE)\n",
    "print(\"training windows:\", data.records)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-15",
   "metadata": {},
   "source": [
    "## Continued training\n",
    "\n",
    "`LMObjective` takes the loaded weights through `pretrained=`, so training starts from them instead of from a fresh random draw. We turn off the EMA (`ema_decay=None`) because a 300-step run is too short for an average to help. The learning rate is small so the model adapts to Shakespeare without forgetting what it knew. The validation perplexity at steps 100, 200 and 300 shows the adaptation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cell-16",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-09-22T20:23:17.373765Z",
     "iopub.status.busy": "2026-09-22T20:23:17.373267Z",
     "iopub.status.idle": "2026-09-22T20:24:36.733822Z",
     "shell.execute_reply": "2026-09-22T20:24:36.732852Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training from step 0 to 300 on {'data': 1, 'expert': 1, 'fsdp': 1, 'tensor': 1, 'sequence': 1, 'stage': 1} (1 process(es))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 50: loss 3.2403\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 100: loss 3.0720\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluation val at step 100: 8 coordinated batches, 64 records, uneven_shards=False, event_key=(4286894075, 772130920): {'val/perplexity': 29.26405867480638}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 150: loss 3.1543\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 200: loss 3.1194\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluation val at step 200: 8 coordinated batches, 64 records, uneven_shards=False, event_key=(1298626594, 818800949): {'val/perplexity': 28.173023282652128}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 250: loss 3.2285\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 300: loss 3.2523\n",
      "Evaluation val at step 300: 8 coordinated batches, 64 records, uneven_shards=False, event_key=(4127360435, 4068970345): {'val/perplexity': 27.850388273194902}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Goodput: first step after 34.72 s, 34.7% of the wall time in steps\n"
     ]
    }
   ],
   "source": [
    "import optax\n",
    "from dew import Checkpoints, Trainer, metrics\n",
    "from dew.objectives.lm import LMObjective\n",
    "\n",
    "objective = LMObjective(model, SEQUENCE_LENGTH, ema_decay=None, pretrained=variables)\n",
    "trainer = Trainer(objective, optax.adamw(LEARNING_RATE), key=jax.random.key(SEED),\n",
    "                  checkpoints=Checkpoints(RUN_DIR))\n",
    "state = trainer.fit(data, steps=STEPS, log_every=50, eval_every=100, checkpoint_every=STEPS,\n",
    "                    metrics=(metrics.perplexity(),))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "cell-17",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-09-22T20:24:36.738449Z",
     "iopub.status.busy": "2026-09-22T20:24:36.737509Z",
     "iopub.status.idle": "2026-09-22T20:24:52.515173Z",
     "shell.execute_reply": "2026-09-22T20:24:52.514559Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ROMEO:\n",
      "I have heard of him, and I know him.\n",
      "\n",
      "LADY CAPULET:\n",
      "He is a nobleman, and a nobleman's son.\n",
      "\n",
      "MERCUTIO:\n",
      "He is a nobleman, and a nobleman's son.\n",
      "\n",
      "LAD\n"
     ]
    }
   ],
   "source": [
    "after = generate(model, state.params, prompt, max_new_tokens=MAX_NEW_TOKENS,\n",
    "                 key=jax.random.key(0), sampling=Sampling(temperature=0.0))\n",
    "print(tokenizer.decode(after.tokens[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-after",
   "metadata": {},
   "source": [
    "Validation perplexity went from about 29 at step 100 to about 28 at step 300, a small change from so few steps. The greedy text changed more: the speakers are now Capulets and Mercutio, and the lines have the play's layout. It still repeats itself, as greedy decoding tends to."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-18",
   "metadata": {},
   "source": [
    "## Exporting back to Hugging Face\n",
    "\n",
    "`save_pretrained_decoder` writes the trained weights in the Hugging Face layout: `config.json`, `model.safetensors`, and the tokenizer files. It runs the same name and shape translation as `load_pretrained`, backwards. Loading the export again gives back exactly the weights we trained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "cell-19",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-09-22T20:24:52.517859Z",
     "iopub.status.busy": "2026-09-22T20:24:52.517606Z",
     "iopub.status.idle": "2026-09-22T20:24:55.068883Z",
     "shell.execute_reply": "2026-09-22T20:24:55.068262Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['config.json', 'generation_config.json', 'model.safetensors', 'tokenizer.json', 'tokenizer_config.json']\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reloaded weights equal the trained ones: True\n"
     ]
    }
   ],
   "source": [
    "from dew.interop import save_pretrained_decoder\n",
    "\n",
    "save_pretrained_decoder(model, state.params, EXPORT_DIR, tokenizer=CHECKPOINT)\n",
    "print(sorted(os.listdir(EXPORT_DIR)))\n",
    "\n",
    "reloaded = load_pretrained(EXPORT_DIR, dtype=\"bfloat16\", attention_impl=\"auto\", max_seq_len=512)\n",
    "same = all(np.array_equal(np.asarray(a), np.asarray(b))\n",
    "           for a, b in zip(jax.tree_util.tree_leaves(reloaded.variables[\"params\"]),\n",
    "                           jax.tree_util.tree_leaves(state.params[\"params\"])))\n",
    "print(\"reloaded weights equal the trained ones:\", same)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-20",
   "metadata": {},
   "source": [
    "## Reading the export with `transformers`\n",
    "\n",
    "The real test of an export is whether another library reads it. We load the folder with `transformers`, run the prompt through it in float32 on the CPU, and compare its most likely next token with Dew's."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cell-21",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-09-22T20:24:55.071941Z",
     "iopub.status.busy": "2026-09-22T20:24:55.071416Z",
     "iopub.status.idle": "2026-09-22T20:25:08.332889Z",
     "shell.execute_reply": "2026-09-22T20:25:08.332265Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "transformers: 198 '\\n'\n",
      "dew:          198 '\\n'\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM\n",
    "\n",
    "hf_model = AutoModelForCausalLM.from_pretrained(EXPORT_DIR, dtype=torch.float32)\n",
    "with torch.no_grad():\n",
    "    hf_next = int(hf_model(input_ids=torch.tensor(np.asarray(prompt))).logits[0, -1].argmax())\n",
    "\n",
    "dew_next = generate(model, state.params, prompt, max_new_tokens=1,\n",
    "                    key=jax.random.key(0), sampling=Sampling(temperature=0.0))\n",
    "dew_next = int(dew_next.tokens[0, -1])\n",
    "print(\"transformers:\", hf_next, repr(tokenizer.decode([hf_next])))\n",
    "print(\"dew:         \", dew_next, repr(tokenizer.decode([dew_next])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-22",
   "metadata": {},
   "source": [
    "## Where to go next\n",
    "\n",
    "The [README model list](../README.md#models) names the decoder families Dew loads, trains, generates with and exports. The recipe flag `--pretrained` runs this notebook's flow at scale, for example `python recipes/lm/train.py data:token-windows --data.path <dir> --pretrained HuggingFaceTB/SmolLM2-135M --tokenizer HuggingFaceTB/SmolLM2-135M`."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "kernelspec": {
   "display_name": "Python 3",
   "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.13.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}