{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/michael-borck/loco-llm/blob/main/notebooks/train_math_adapter.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2Fytddn13gX3"
      },
      "source": [
        "# LocoLLM: Train a Math Adapter on Colab T4\n",
        "\n",
        "This notebook trains a QLoRA math adapter on **Qwen3-4B** using [Unsloth](https://github.com/unslothai/unsloth), then exports a merged **Q4_K_M GGUF** you can download and run locally via Ollama.\n",
        "\n",
        "**What you need:** A free Google Colab account with a T4 GPU runtime (16GB VRAM).\n",
        "\n",
        "**What you get:** A ~2.5GB GGUF file that is a complete, standalone math-specialized model.\n",
        "\n",
        "**Time:** ~20-30 minutes end to end (including install).\n",
        "\n",
        "---\n",
        "\n",
        "## How to use this notebook\n",
        "\n",
        "1. Open in Colab (click the badge above or use File > Open in Colab)\n",
        "2. Go to **Runtime > Change runtime type** and select **T4 GPU**\n",
        "3. Run all cells in order (**Runtime > Run all**)\n",
        "4. Download the GGUF file from the last cell\n",
        "5. Load it into Ollama locally: `ollama create locollm-math -f Modelfile`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LbchqnH43gX5"
      },
      "source": [
        "## Step 0: Verify GPU\n",
        "\n",
        "Make sure you have a T4 (or better) GPU assigned."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zYpgmGkg3gX6"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi --query-gpu=name,memory.total --format=csv,noheader"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BQKg4Sv73gX7"
      },
      "source": [
        "## Step 1: Install Unsloth\n",
        "\n",
        "This takes 2-3 minutes. Unsloth pulls in the right versions of transformers, peft, trl, etc."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BxTkLTG63gX7"
      },
      "outputs": [],
      "source": [
        "!pip install --upgrade --no-cache-dir unsloth unsloth_zoo 2>&1 | tail -5"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hRx_7EnE3gX7"
      },
      "source": [
        "## Step 2: Prepare Training Data (GSM8K)\n",
        "\n",
        "We download 200 examples from [GSM8K](https://huggingface.co/datasets/openai/gsm8k) and format them for Qwen3's chat template. Each example has step-by-step reasoning ending with \"The answer is N\".\n",
        "\n",
        "200 examples is enough for a proof-of-concept. For a production adapter, use 500-1000+."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eDSUhCTV3gX8"
      },
      "outputs": [],
      "source": [
        "import re\n",
        "\n",
        "import requests\n",
        "\n",
        "NUM_EXAMPLES = 200\n",
        "\n",
        "# Download from HuggingFace datasets API\n",
        "url = f\"https://datasets-server.huggingface.co/rows?dataset=openai/gsm8k&config=main&split=train&offset=0&length={NUM_EXAMPLES}\"\n",
        "resp = requests.get(url, timeout=60)\n",
        "resp.raise_for_status()\n",
        "rows = [r[\"row\"] for r in resp.json()[\"rows\"]]\n",
        "print(f\"Downloaded {len(rows)} examples from GSM8K\")\n",
        "\n",
        "\n",
        "def format_answer(answer_text: str) -> str:\n",
        "    \"\"\"Convert GSM8K answer format to clean step-by-step reasoning.\"\"\"\n",
        "    parts = answer_text.split(\"####\")\n",
        "    reasoning = re.sub(r\"<<.*?>>\", \"\", parts[0]).strip()\n",
        "    final_answer = parts[1].strip() if len(parts) > 1 else \"\"\n",
        "    return f\"{reasoning}\\nThe answer is {final_answer}\"\n",
        "\n",
        "\n",
        "# Format for Qwen3 chat template\n",
        "training_data = []\n",
        "for ex in rows:\n",
        "    training_data.append(\n",
        "        {\n",
        "            \"conversations\": [\n",
        "                {\"role\": \"user\", \"content\": ex[\"question\"].strip()},\n",
        "                {\"role\": \"assistant\", \"content\": format_answer(ex[\"answer\"])},\n",
        "            ]\n",
        "        }\n",
        "    )\n",
        "\n",
        "print(f\"Formatted {len(training_data)} training examples\")\n",
        "print(f\"\\nSample question: {training_data[0]['conversations'][0]['content'][:100]}...\")\n",
        "print(f\"Sample answer:   {training_data[0]['conversations'][1]['content'][:100]}...\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wcSPThw23gX8"
      },
      "source": [
        "## Step 3: Load Qwen3-4B with Unsloth\n",
        "\n",
        "Unsloth loads the model in 4-bit quantization. On a T4 this uses ~5GB VRAM, leaving plenty of room for training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zq0_6ab13gX9"
      },
      "outputs": [],
      "source": [
        "from unsloth import FastModel\n",
        "\n",
        "MODEL_NAME = \"unsloth/Qwen3-4B-unsloth-bnb-4bit\"\n",
        "MAX_SEQ_LENGTH = 1024\n",
        "\n",
        "model, tokenizer = FastModel.from_pretrained(\n",
        "    model_name=MODEL_NAME,\n",
        "    max_seq_length=MAX_SEQ_LENGTH,\n",
        "    load_in_4bit=True,\n",
        "    full_finetuning=False,\n",
        ")\n",
        "\n",
        "print(f\"Model loaded: {MODEL_NAME}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HaMqCME13gX9"
      },
      "source": [
        "## Step 4: Apply LoRA Adapters\n",
        "\n",
        "We apply LoRA to the attention layers (q/k/v/o projections). Rank 16 is a good middle ground for math reasoning — enough capacity to learn patterns without overfitting on 200 examples."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9LMLcQpz3gX9"
      },
      "outputs": [],
      "source": [
        "LORA_R = 16\n",
        "LORA_ALPHA = 32\n",
        "LORA_DROPOUT = 0.05\n",
        "\n",
        "model = FastModel.get_peft_model(\n",
        "    model,\n",
        "    r=LORA_R,\n",
        "    lora_alpha=LORA_ALPHA,\n",
        "    lora_dropout=LORA_DROPOUT,\n",
        "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\"],\n",
        ")\n",
        "\n",
        "trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
        "total = sum(p.numel() for p in model.parameters())\n",
        "print(f\"Trainable parameters: {trainable:,} / {total:,} ({100 * trainable / total:.2f}%)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UEjwF9pu3gX-"
      },
      "source": [
        "## Step 5: Prepare Dataset\n",
        "\n",
        "Convert our training data to a HuggingFace Dataset and apply the chat template."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c4Dph2cD3gX-"
      },
      "outputs": [],
      "source": [
        "from datasets import Dataset\n",
        "\n",
        "dataset = Dataset.from_list(training_data)\n",
        "\n",
        "\n",
        "def format_conversations(examples):\n",
        "    \"\"\"Apply Qwen3 chat template to conversations.\"\"\"\n",
        "    texts = []\n",
        "    for convos in examples[\"conversations\"]:\n",
        "        text = tokenizer.apply_chat_template(\n",
        "            convos,\n",
        "            tokenize=False,\n",
        "            add_generation_prompt=False,\n",
        "        )\n",
        "        texts.append(text)\n",
        "    return {\"text\": texts}\n",
        "\n",
        "\n",
        "dataset = dataset.map(format_conversations, batched=True)\n",
        "print(f\"Dataset ready: {len(dataset)} examples\")\n",
        "print(f\"\\nFormatted sample (first 200 chars):\\n{dataset[0]['text'][:200]}...\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hT2Fcf0V3gX-"
      },
      "source": [
        "## Step 6: Train\n",
        "\n",
        "Fine-tune for 3 epochs with SFTTrainer. On a T4, this takes ~10-15 minutes for 200 examples.\n",
        "\n",
        "Watch the training loss — it should decrease steadily. If it plateaus early, the model has learned what it can from this data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1w83qjzs3gX-"
      },
      "outputs": [],
      "source": [
        "from transformers import TrainingArguments\n",
        "from trl import SFTTrainer\n",
        "\n",
        "NUM_EPOCHS = 3\n",
        "BATCH_SIZE = 2\n",
        "GRADIENT_ACCUMULATION = 4  # effective batch size = 8\n",
        "LEARNING_RATE = 2e-4\n",
        "\n",
        "trainer = SFTTrainer(\n",
        "    model=model,\n",
        "    tokenizer=tokenizer,\n",
        "    train_dataset=dataset,\n",
        "    args=TrainingArguments(\n",
        "        output_dir=\"./checkpoints\",\n",
        "        per_device_train_batch_size=BATCH_SIZE,\n",
        "        gradient_accumulation_steps=GRADIENT_ACCUMULATION,\n",
        "        num_train_epochs=NUM_EPOCHS,\n",
        "        learning_rate=LEARNING_RATE,\n",
        "        lr_scheduler_type=\"cosine\",\n",
        "        warmup_ratio=0.1,\n",
        "        fp16=True,\n",
        "        logging_steps=5,\n",
        "        save_strategy=\"epoch\",\n",
        "        seed=42,\n",
        "        optim=\"adamw_8bit\",\n",
        "        report_to=\"none\",\n",
        "    ),\n",
        "    dataset_text_field=\"text\",\n",
        "    max_seq_length=MAX_SEQ_LENGTH,\n",
        "    packing=False,\n",
        ")\n",
        "\n",
        "print(f\"Training: {NUM_EPOCHS} epochs, effective batch size {BATCH_SIZE * GRADIENT_ACCUMULATION}\")\n",
        "trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LN8NJx9c3gX-"
      },
      "source": [
        "## Step 7: Quick Sanity Check\n",
        "\n",
        "Before exporting, test the fine-tuned model on a couple of math problems to make sure it's producing reasonable output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_oaz8GYI3gX-"
      },
      "outputs": [],
      "source": [
        "from unsloth import FastModel\n",
        "\n",
        "# Switch to inference mode\n",
        "FastModel.for_inference(model)\n",
        "\n",
        "test_questions = [\n",
        "    \"What is 15 + 27?\",\n",
        "    (\n",
        "        \"A store has 120 apples. They sell 45 in the morning and 30 in the\"\n",
        "        \" afternoon. How many apples are left?\"\n",
        "    ),\n",
        "    \"If a shirt costs $40 and is on sale for 25% off, what is the sale price?\",\n",
        "]\n",
        "\n",
        "for q in test_questions:\n",
        "    messages = [{\"role\": \"user\", \"content\": q}]\n",
        "    inputs = tokenizer.apply_chat_template(\n",
        "        messages, tokenize=True, add_generation_prompt=True, return_tensors=\"pt\"\n",
        "    ).to(model.device)\n",
        "    outputs = model.generate(inputs, max_new_tokens=256, temperature=0.0, use_cache=True)\n",
        "    response = tokenizer.decode(outputs[0][inputs.shape[-1] :], skip_special_tokens=True)\n",
        "    print(f\"Q: {q}\")\n",
        "    print(f\"A: {response}\")\n",
        "    print(\"-\" * 60)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5ZVsPFdG3gX_"
      },
      "source": [
        "## Step 8: Export Merged GGUF\n",
        "\n",
        "Merge the LoRA weights into the base model and export as a Q4_K_M GGUF. This creates a standalone model file (~2.5GB) that Ollama can load directly.\n",
        "\n",
        "**Why merge?** Ollama doesn't support Qwen3 LoRA adapters via the `ADAPTER` directive (only Llama/Mistral/Gemma). Merging produces a complete, self-contained model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yNF0qt4M3gX_"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "OUTPUT_DIR = \"./gguf_output\"\n",
        "QUANTIZATION_METHOD = \"q4_k_m\"\n",
        "\n",
        "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
        "\n",
        "print(f\"Exporting merged GGUF ({QUANTIZATION_METHOD})...\")\n",
        "print(\"This takes a few minutes — merging weights and quantizing.\")\n",
        "\n",
        "model.save_pretrained_gguf(\n",
        "    OUTPUT_DIR,\n",
        "    tokenizer,\n",
        "    quantization_method=QUANTIZATION_METHOD,\n",
        ")\n",
        "\n",
        "# Show the exported file\n",
        "for f in os.listdir(OUTPUT_DIR):\n",
        "    if f.endswith(\".gguf\"):\n",
        "        size_mb = os.path.getsize(os.path.join(OUTPUT_DIR, f)) / (1024 * 1024)\n",
        "        print(f\"\\nExported: {f} ({size_mb:.0f} MB)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n3eLWQ103gX_"
      },
      "source": [
        "## Step 9: Save to Google Drive (Optional)\n",
        "\n",
        "Save the GGUF to your Google Drive so it persists after the Colab session ends. You can also download it directly from the next cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LYKjJOdo3gX_"
      },
      "outputs": [],
      "source": [
        "import shutil\n",
        "\n",
        "from google.colab import drive\n",
        "\n",
        "drive.mount(\"/content/drive\")\n",
        "\n",
        "DRIVE_DIR = \"/content/drive/MyDrive/locollm-adapters/math\"\n",
        "os.makedirs(DRIVE_DIR, exist_ok=True)\n",
        "\n",
        "for f in os.listdir(OUTPUT_DIR):\n",
        "    if f.endswith(\".gguf\"):\n",
        "        src = os.path.join(OUTPUT_DIR, f)\n",
        "        dst = os.path.join(DRIVE_DIR, f)\n",
        "        print(f\"Copying {f} to Google Drive...\")\n",
        "        shutil.copy2(src, dst)\n",
        "        print(f\"Saved to: {dst}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cn7VLUeS3gX_"
      },
      "source": [
        "## Step 10: Download the GGUF\n",
        "\n",
        "Download the GGUF file to your local machine. Then load it into Ollama:\n",
        "\n",
        "```bash\n",
        "# On your local machine:\n",
        "cd loco-llm\n",
        "mkdir -p adapters/math/gguf\n",
        "mv ~/Downloads/unsloth.Q4_K_M.gguf adapters/math/gguf/\n",
        "\n",
        "# Create Ollama model\n",
        "echo 'FROM ./adapters/math/gguf/unsloth.Q4_K_M.gguf' > adapters/math/Modelfile\n",
        "ollama create locollm-math -f adapters/math/Modelfile\n",
        "\n",
        "# Test it\n",
        "ollama run locollm-math \"What is 15 + 27?\"\n",
        "\n",
        "# Or use the LocoLLM CLI\n",
        "uv run loco setup\n",
        "uv run loco eval math\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gXnOULvA3gX_"
      },
      "outputs": [],
      "source": [
        "from google.colab import files\n",
        "\n",
        "for f in os.listdir(OUTPUT_DIR):\n",
        "    if f.endswith(\".gguf\"):\n",
        "        print(f\"Downloading {f}...\")\n",
        "        files.download(os.path.join(OUTPUT_DIR, f))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y8bL9Ijk3gX_"
      },
      "source": [
        "---\n",
        "\n",
        "## Training Summary\n",
        "\n",
        "| Parameter | Value |\n",
        "|-----------|-------|\n",
        "| Base model | Qwen3-4B (4-bit via Unsloth) |\n",
        "| Method | QLoRA |\n",
        "| LoRA rank | 16 |\n",
        "| LoRA alpha | 32 |\n",
        "| Dataset | GSM8K (200 examples) |\n",
        "| Epochs | 3 |\n",
        "| Effective batch size | 8 |\n",
        "| Learning rate | 2e-4 |\n",
        "| Export format | Q4_K_M GGUF (~2.5GB) |\n",
        "| Hardware | Colab T4 (16GB VRAM) |\n",
        "\n",
        "## What Next?\n",
        "\n",
        "- **Evaluate**: Run `uv run loco eval math` to compare the adapter against the base model\n",
        "- **More data**: Try 500 or 1000 GSM8K examples (change `NUM_EXAMPLES` in Step 2)\n",
        "- **Different domain**: Fork this notebook and swap GSM8K for your own dataset\n",
        "- **Iterate**: See [adapter-guide.md](https://github.com/michael-borck/loco-llm/blob/main/docs/adapter-guide.md) for the full development cycle"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}