{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NoTokenLM-MicroGen — Use Notebook\n",
    "\n",
    "This is the notebook to use if you just want to **run one of the MicroGen models**. Every model in the series ships as a real `transformers`-compatible checkpoint, so loading one is a single `from_pretrained(..., trust_remote_code=True)` call — no manual architecture code, no private repos.\n",
    "\n",
    "Because this single repository hosts six different checkpoints side by side (one per subfolder), Hugging Face's default \"Use this model\" widget doesn't represent it correctly — this notebook is the intended entry point instead.\n",
    "\n",
    "**Works on Google Colab and Kaggle** — the setup cell below detects which platform you're on and installs accordingly.\n",
    "\n",
    "## What this notebook does\n",
    "1. Lets you pick a model from a dropdown\n",
    "2. Loads it directly from the Hub with `AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True)`\n",
    "3. Runs text generation on your own prompt\n",
    "\n",
    "No token is required (the repo is public) — a GPU (T4 is enough) will make generation faster, but these models are small enough to run fine on CPU too.\n",
    "\n",
    "> These models have no tokenizer — input and output are raw UTF-8 bytes. `model.generate_bytes(prompt, ...)` handles that for you."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "**If you hit a `ModuleNotFoundError: Could not import module 'PreTrainedModel'` / `operator torchvision::nms does not exist` error:** this comes from Colab/Kaggle's pre-installed `torch` and `torchvision` going out of sync after `pip install -U torch` upgrades one but not the other. The cell below pins compatible versions of both together to avoid that. If you already hit the error before running this cell, use **Runtime → Restart runtime** (Colab) or **Run → Restart session** (Kaggle) after this cell finishes, then re-run from the top."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os\n",
    "\n",
    "IN_COLAB = 'google.colab' in sys.modules\n",
    "IN_KAGGLE = 'kaggle_secrets' in sys.modules or 'KAGGLE_KERNEL_RUN_TYPE' in os.environ\n",
    "\n",
    "if IN_COLAB:\n",
    "    print('Platform: Google Colab')\n",
    "elif IN_KAGGLE:\n",
    "    print('Platform: Kaggle')\n",
    "else:\n",
    "    print('Platform: local / other Jupyter')\n",
    "\n",
    "# Pin torch + torchvision TOGETHER (same release) to avoid the ABI mismatch\n",
    "# that causes \"operator torchvision::nms does not exist\". Upgrading only\n",
    "# torch (as a plain `pip install -U torch` would) is what breaks this.\n",
    "!pip install -q -U \"torch==2.5.1\" \"torchvision==0.20.1\" --index-url https://download.pytorch.org/whl/cu121\n",
    "!pip install -q -U transformers huggingface_hub\n",
    "\n",
    "print('\\nInstall complete.')\n",
    "print('If this is the FIRST time you are running this in this session and you')\n",
    "print('previously hit the torchvision::nms error, please restart the runtime now')\n",
    "print('(Runtime > Restart runtime on Colab, Run > Restart session on Kaggle)')\n",
    "print('and then re-run all cells from the top.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(f'Device: {device}')\n",
    "if device == 'cuda':\n",
    "    print(torch.cuda.get_device_name(0))\n",
    "print(f'torch: {torch.__version__}')\n",
    "import torchvision\n",
    "print(f'torchvision: {torchvision.__version__}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model registry\n",
    "\n",
    "Used only to show languages/context length and to offer default prompts — not needed for loading, that part is fully generic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PUBLIC_REPO = \"omurberaisik/NoTokenLM-MicroGen\"\n",
    "\n",
    "MODEL_INFO = {\n",
    "    \"2.5\": dict(languages=[\"en\"], prompts=[\"The \", \"I think that \", \"Once upon a time \"]),\n",
    "    \"2.6\": dict(languages=[\"en\"], prompts=[\"The \", \"I think that \", \"Once upon a time \"]),\n",
    "    \"2.7\": dict(languages=[\"en\"], prompts=[\"The \", \"I think that \", \"Once upon a time \"]),\n",
    "    \"3.5\": dict(languages=[\"en\"], prompts=[\"The \", \"I think that \", \"Once upon a time \"]),\n",
    "    \"3.6\": dict(languages=[\"en\"], prompts=[\"The \", \"I think that \", \"Once upon a time \"]),\n",
    "    \"3.7\": dict(languages=[\"en\", \"es\", \"id\", \"it\"],\n",
    "                prompts=[(\"en\", \"The \"), (\"en\", \"I think that \"),\n",
    "                         (\"es\", \"El \"), (\"id\", \"Saya \"), (\"it\", \"Il \")]),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pick a model and load it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SELECTED_VERSION = \"3.6\"  #@param [\"2.5\", \"2.6\", \"2.7\", \"3.5\", \"3.6\", \"3.7\"]\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    PUBLIC_REPO,\n",
    "    subfolder=SELECTED_VERSION,\n",
    "    trust_remote_code=True,\n",
    ")\n",
    "model.to(device)\n",
    "model.eval()\n",
    "\n",
    "n_params = sum(p.numel() for p in model.parameters())\n",
    "info = MODEL_INFO[SELECTED_VERSION]\n",
    "print(f\"Loaded NoTokenLM-MicroGen-{SELECTED_VERSION}: {n_params:,} parameters\")\n",
    "print(f\"Languages: {info['languages']}\")\n",
    "print(f\"Context length: {model.config.max_len} bytes\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate text\n",
    "\n",
    "Edit `MY_PROMPT` and re-run. `temperature` controls randomness — lower is more predictable, higher is more chaotic; these models were evaluated at 0.5."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MY_PROMPT = \"The \"  #@param {type:\"string\"}\n",
    "TEMPERATURE = 0.5  #@param {type:\"number\"}\n",
    "N_NEW_BYTES = 200  #@param {type:\"integer\"}\n",
    "\n",
    "out = model.generate_bytes(MY_PROMPT, n_new_bytes=N_NEW_BYTES, temperature=TEMPERATURE)\n",
    "print(f\"Model:  NoTokenLM-MicroGen-{SELECTED_VERSION} ({n_params:,} params)\")\n",
    "print(f\"Prompt: {MY_PROMPT!r}\")\n",
    "print(f\"Output: {out!r}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Try the model's default prompt set\n",
    "\n",
    "Each model has a small built-in prompt set (the same ones used for the README examples). Run this to see how the currently loaded model does on all of them at once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for item in info[\"prompts\"]:\n",
    "    if isinstance(item, tuple):\n",
    "        lang, prompt = item\n",
    "        tag = f\"[{lang}] \"\n",
    "    else:\n",
    "        prompt = item\n",
    "        tag = \"\"\n",
    "    out = model.generate_bytes(prompt, n_new_bytes=200, temperature=0.5)\n",
    "    print(f\"{tag}{prompt!r} -> {out!r}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "### A note on what to expect\n",
    "\n",
    "These are small, from-scratch, byte-level models with no instruction tuning. From 2.7 onward they tend to produce grammatically plausible continuations, but they are not chat assistants and will not reliably stay on topic for long. See the [model card](https://huggingface.co/omurberaisik/NoTokenLM-MicroGen) for a full, honest rundown of each version's limitations."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU"
 },
 "nbformat": 4,
 "nbformat_minor": 0
}