{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train YCode-LM v5 (~385M parameters) on free Kaggle GPUs\n",
    "\n",
    "This notebook grows **YCode-LM v4** (100M parameters, trained from scratch) into **YCode-LM v5**\n",
    "(34 layers \u00d7 1024 wide, about 385M parameters). It then continues pretraining, instruction-tunes,\n",
    "evaluates, and **publishes the model to your Hugging Face account**.\n",
    "\n",
    "A Kaggle session lasts up to 12 hours, so training is split across sessions. Every ~90 minutes the\n",
    "checkpoint is uploaded to a *private* Hugging Face repo. **Just run this notebook again** each time a\n",
    "session ends; it resumes automatically and publishes the finished model at the end.\n",
    "\n",
    "## One-time setup\n",
    "\n",
    "1. **Settings \u2192 Accelerator \u2192 GPU T4 x2**.\n",
    "2. **Settings \u2192 Internet \u2192 On**. Kaggle requires a verified phone number for this.\n",
    "3. Create a Hugging Face token with **Write** access at https://huggingface.co/settings/tokens.\n",
    "   In Kaggle, open **Add-ons \u2192 Secrets \u2192 Add a new secret**, name it `HF_TOKEN` and paste the token.\n",
    "4. Set `HF_USER` below to your Hugging Face username.\n",
    "\n",
    "## Each session\n",
    "\n",
    "Click **Save Version \u2192 Save & Run All (Commit)** so it runs in the background for up to 12 hours,\n",
    "even with your laptop closed. Check the log at the end to see the progress. When it says\n",
    "**\"All phases complete\"**, your model is live at `https://huggingface.co/<HF_USER>/ycode-lm-v5`.\n",
    "\n",
    "Expect about 3\u20135 sessions with the default settings. Kaggle's free quota is about 30 GPU-hours per week."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# ---- Settings -------------------------------------------------------------\n",
    "HF_USER = \"Yazhiniyan\"          # your Hugging Face username\n",
    "MODEL_REPO = f\"{HF_USER}/ycode-lm-v5\"\n",
    "PUBLIC = True                   # False keeps the finished model private\n",
    "PRESET = \"v5-385m\"              # 34 x 1024, ~385M parameters\n",
    "PRETRAIN_STEPS = 30000          # x 16k tokens/step (2 GPUs x batch 8 x 1024) \u2248 490M tokens\n",
    "SFT_STEPS = 1000\n",
    "SESSION_HOURS = 11.0            # stop starting new segments before Kaggle's 12 h limit\n",
    "BRANCH = \"claude/affectionate-rubin-7rl1am\"   # switch to \"main\" once the pull request is merged"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import os, subprocess\n",
    "from kaggle_secrets import UserSecretsClient\n",
    "\n",
    "os.environ[\"HF_TOKEN\"] = UserSecretsClient().get_secret(\"HF_TOKEN\")\n",
    "if not os.path.isdir(\"Code-agent\"):\n",
    "    subprocess.run([\"git\", \"clone\", \"--depth\", \"1\", \"-b\", BRANCH,\n",
    "                    \"https://github.com/yazhiniyanMS/Code-agent.git\"], check=True)\n",
    "os.chdir(\"Code-agent\")\n",
    "subprocess.run([\"pip\", \"install\", \"-q\", \"-e\", \".[local,hub]\"], check=True)\n",
    "\n",
    "import torch\n",
    "print(\"GPUs:\", torch.cuda.device_count(), [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])\n",
    "assert torch.cuda.is_available(), \"Turn on a GPU accelerator in Settings.\""
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Grows v4 -> v5 on the first run; on later runs it resumes from the Hugging Face checkpoint.\n",
    "!python -m ycode.lm.cli autotrain \\\n",
    "    --base-model models/ycode-lm-v4 --preset {PRESET} \\\n",
    "    --hf-repo {MODEL_REPO} {\"--public\" if PUBLIC else \"\"} \\\n",
    "    --pretrain-steps {PRETRAIN_STEPS} --sft-steps {SFT_STEPS} \\\n",
    "    --batch-size 8 --hours {SESSION_HOURS} --segment-minutes 90 --max-mb 1000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## After it finishes\n",
    "\n",
    "Try the model here, or locally after downloading it from Hugging Face:\n",
    "\n",
    "```bash\n",
    "huggingface-cli download <HF_USER>/ycode-lm-v5 --local-dir models/ycode-lm-v5\n",
    "ycode --local models/ycode-lm-v5\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import os\n",
    "if os.path.exists(\"ycode-autotrain/release/model.pt\"):\n",
    "    from ycode.lm.generate import LocalLM\n",
    "    lm = LocalLM(\"ycode-autotrain/release\")\n",
    "    print(lm.chat([(\"user\", \"Write a Python function `is_even(n)` that returns True if n is even.\")],\n",
    "                  max_new_tokens=150, temperature=0.2))\n",
    "else:\n",
    "    print(\"Not finished yet: run the notebook again in a new session to continue training.\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "kaggle": {
   "accelerator": "nvidiaTeslaT4",
   "isInternetEnabled": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}