{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vashtra-0.6B: build the corpus, train the model, ship the demo\n",
    "\n",
    "Two-stage domain adaptation of **Qwen3-0.6B** into a machine-learning specialist.\n",
    "\n",
    "| Stage | Data | What it teaches |\n",
    "|---|---|---|\n",
    "| A, continued pretraining | full-text arXiv ML papers + 117k cs.LG abstracts | ML vocabulary, concepts, how arguments are made |\n",
    "| B, supervised fine-tuning | ML Q&A + general chat | answering, instruction following |\n",
    "\n",
    "**Runtime \u2192 Change runtime type \u2192 T4 GPU**, then **Runtime \u2192 Run all**.\n",
    "Everything pushes to the Hub automatically. Budget ~2 to 3 hours on the free tier.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Confirm you have a GPU\n",
    "If this cell fails you are on CPU. Switch the runtime first, or nothing below will finish."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!nvidia-smi -L\n",
    "import torch\n",
    "assert torch.cuda.is_available(), \"No GPU. Runtime -> Change runtime type -> T4 GPU\"\n",
    "print(torch.cuda.get_device_name(0),\n",
    "      f\"{torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Install pinned dependencies\n",
    "Versions are pinned to what the pipeline was verified against."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q -U \"transformers==5.9.0\" \"datasets==4.8.5\" accelerate bitsandbytes huggingface_hub\n",
    "import transformers, datasets\n",
    "print(\"transformers\", transformers.__version__, \"| datasets\", datasets.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Sign in to Hugging Face\n",
    "\n",
    "Paste a **write** token from https://huggingface.co/settings/tokens"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login, whoami\n",
    "notebook_login()\n",
    "print(\"signed in as:\", whoami()[\"name\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Fetch the pipeline scripts"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!wget -q -O build_corpus.py  https://huggingface.co/vigneshwar234/Vashtra-0.6B/resolve/main/build_corpus.py\n",
    "!wget -q -O train_vashtra.py https://huggingface.co/vigneshwar234/Vashtra-0.6B/resolve/main/train_vashtra.py\n",
    "!wget -q -O export_onnx.py   https://huggingface.co/vigneshwar234/Vashtra-0.6B/resolve/main/export_onnx.py\n",
    "!ls -la *.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Build the corpus\n",
    "\n",
    "Streams full-text arXiv papers, cs.LG abstracts, ML/stats/CS StackExchange, and\n",
    "the smoltalk2 mixture into `cpt.jsonl` + `sft.jsonl`, then pushes them to\n",
    "`vigneshwar234/vashtra-ml-corpus`. Takes ~25 to 40 minutes, mostly the\n",
    "StackExchange and full-text scans.\n",
    "\n",
    "Skip this cell if the corpus is already built. Training pulls it from the Hub."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!python build_corpus.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. Train\n",
    "\n",
    "Stage A (continued pretraining) then Stage B (SFT), then the model is pushed to\n",
    "`vigneshwar234/Vashtra-0.6B`.\n",
    "\n",
    "The corpus is about 90M tokens, which is more than a free Colab session will get\n",
    "through. Stage A therefore takes a token budget, defaulting to 30M, which is\n",
    "roughly two hours on a T4. Checkpoints are pushed to the Hub as it goes.\n",
    "\n",
    "If the session drops, open a new one and run the same cell with `--resume`. It\n",
    "pulls the last checkpoint back down and carries on instead of starting over.\n",
    "To train on more of the corpus, raise `--cpt_tokens` (or pass `0` for all of it)\n",
    "and expect to resume across two or three sessions."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!python train_vashtra.py --stage all --cpt_tokens 30000000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the session above was cut off, run this instead of starting over."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# !python train_vashtra.py --stage all --cpt_tokens 30000000 --resume"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7. Smoke-test what you trained\n",
    "\n",
    "Sanity check, not an evaluation. You are looking for coherent, on-topic ML answers,\n",
    "and for the general question to still work, which is what the anchor data buys."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "import torch\n",
    "\n",
    "m = AutoModelForCausalLM.from_pretrained(\"vigneshwar234/Vashtra-0.6B\",\n",
    "                                         dtype=torch.bfloat16).cuda().eval()\n",
    "t = AutoTokenizer.from_pretrained(\"vigneshwar234/Vashtra-0.6B\")\n",
    "\n",
    "for q in [\"Explain why batch normalization helps training.\",\n",
    "          \"What is the difference between LoRA and full fine-tuning?\",\n",
    "          \"When would I use focal loss instead of cross-entropy?\",\n",
    "          \"Write me a two-line poem about the sea.\"]:   # <- anchor still alive?\n",
    "    ids = t.apply_chat_template([{\"role\":\"user\",\"content\":q}],\n",
    "                                add_generation_prompt=True, return_tensors=\"pt\").cuda()\n",
    "    out = m.generate(ids, max_new_tokens=220, do_sample=False)\n",
    "    print(\"Q:\", q)\n",
    "    print(\"A:\", t.decode(out[0][ids.shape[1]:], skip_special_tokens=True))\n",
    "    print(\"-\"*72)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 8. Measure it against the base model\n",
    "\n",
    "Compare perplexity on held-out arXiv ML text. If stage A worked, Vashtra should be\n",
    "meaningfully lower than Qwen3-0.6B here. Put the real number in the model card.\n",
    "This is the claim that makes the model credible."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch, math\n",
    "from datasets import load_dataset\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "tok = AutoTokenizer.from_pretrained(\"Qwen/Qwen3-0.6B\")\n",
    "held = load_dataset(\"vigneshwar234/vashtra-ml-corpus\", data_files=\"cpt.jsonl\",\n",
    "                    split=\"train\").select(range(300))\n",
    "\n",
    "def perplexity(name):\n",
    "    mdl = AutoModelForCausalLM.from_pretrained(name, dtype=torch.bfloat16).cuda().eval()\n",
    "    nll, ntok = 0.0, 0\n",
    "    with torch.no_grad():\n",
    "        for r in held:\n",
    "            ids = tok(r[\"text\"], return_tensors=\"pt\",\n",
    "                      truncation=True, max_length=1024).input_ids.cuda()\n",
    "            if ids.shape[1] < 32: continue\n",
    "            loss = mdl(ids, labels=ids).loss\n",
    "            nll += loss.item() * (ids.shape[1]-1); ntok += ids.shape[1]-1\n",
    "    del mdl; torch.cuda.empty_cache()\n",
    "    return math.exp(nll/ntok)\n",
    "\n",
    "base = perplexity(\"Qwen/Qwen3-0.6B\")\n",
    "ours = perplexity(\"vigneshwar234/Vashtra-0.6B\")\n",
    "print(f\"arXiv cs.LG held-out perplexity\")\n",
    "print(f\"  Qwen3-0.6B    {base:.2f}\")\n",
    "print(f\"  Vashtra-0.6B  {ours:.2f}   ({(base-ours)/base*100:+.1f}%)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 9. Export to ONNX for the in-browser demo\n",
    "\n",
    "Produces `vigneshwar234/Vashtra-0.6B-ONNX`, which the static WebGPU Space loads\n",
    "directly, so the demo costs nothing to host and runs on the visitor's own GPU."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q \"optimum[onnxruntime]>=1.20\" onnx onnxruntime\n",
    "!python export_onnx.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Done\n",
    "\n",
    "| | |\n",
    "|---|---|\n",
    "| Model | https://huggingface.co/vigneshwar234/Vashtra-0.6B |\n",
    "| Corpus | https://huggingface.co/datasets/vigneshwar234/vashtra-ml-corpus |\n",
    "| Live demo | https://huggingface.co/spaces/vigneshwar234/Vashtra |\n",
    "\n",
    "**Last step:** put the real perplexity number from cell 8 into the model card, and\n",
    "replace the \"weights are training\" banner. Measured results are what make people\n",
    "trust a small model, so ship them.\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [],
   "gpuType": "T4",
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}