{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "665007e7",
   "metadata": {
    "id": "665007e7"
   },
   "source": [
    "# DiffusionGemma Sampling\n",
    "\n",
    "Example on how to load a DiffusionGemma model and run inference on it.\n",
    "\n",
    "Unlike autoregressive (AR) models that generate tokens one at a time left-to-right, DiffusionGemma generates text by iteratively denoising a block (\"canvas\") of tokens in parallel.\n",
    "\n",
    "The diffusion `ChatSampler` provides the same interface as the AR `gm.text.ChatSampler`, so you can use it as a drop-in replacement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c7378f42",
   "metadata": {
    "executionInfo": {
     "elapsed": 305,
     "status": "ok",
     "timestamp": 1781096460654,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -60
    },
    "id": "c7378f42"
   },
   "outputs": [],
   "source": [
    "!pip install -q gemma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0e7e155d",
   "metadata": {
    "executionInfo": {
     "elapsed": 318,
     "status": "ok",
     "timestamp": 1781096461241,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -60
    },
    "id": "0e7e155d"
   },
   "outputs": [],
   "source": [
    "# Common imports\n",
    "import os\n",
    "import jax\n",
    "\n",
    "# Gemma imports\n",
    "  from gemma import gm\n",
    "  from gemma import diffusion"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4a59465",
   "metadata": {
    "id": "b4a59465"
   },
   "source": [
    "## Load model and checkpoint\n",
    "\n",
    "Load the DiffusionGemma model and its parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3821c16b",
   "metadata": {
    "colab": {
     "height": 247
    },
    "executionInfo": {
     "elapsed": 57,
     "status": "error",
     "timestamp": 1781096461560,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -60
    },
    "id": "3821c16b",
    "outputId": "a6e74b78-54c6-4ce8-aac6-a94c1b76fc55"
   },
   "outputs": [],
   "source": [
    "# Initialize the DiffusionGemma model\n",
    "model = diffusion.DiffusionGemma_26B_A4B()\n",
    "\n",
    "# Load pretrained parameters\n",
    "params = gm.ckpts.load_params(diffusion.CheckpointPath.DIFFUSIONGEMMA_26B_A4B_IT, restore_concurrent_gb=16)\n",
    "print(\"Checkpoint loaded.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cb30c8c",
   "metadata": {
    "id": "5cb30c8c"
   },
   "source": [
    "## Sampling the model\n",
    "\n",
    "The easiest way to use DiffusionGemma is through `diffusion.ChatSampler`.\n",
    "\n",
    "- **`canvas_length`**: Size of the token block denoised in parallel (default: 256).\n",
    "- **`max_denoising_steps`**: Refinement steps per canvas (default: 48).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4df26535",
   "metadata": {
    "id": "4df26535"
   },
   "outputs": [],
   "source": [
    "# Initialize the tokenizer\n",
    "tokenizer = gm.text.Gemma4Tokenizer()\n",
    "\n",
    "# Create the ChatSampler with diffusion-specific parameters\n",
    "sampler = diffusion.ChatSampler(\n",
    "    model=model,\n",
    "    params=params,\n",
    "    tokenizer=tokenizer,\n",
    "    canvas_length=256,\n",
    "    max_denoising_steps=48,\n",
    "    # Settings for example colab only\n",
    "    pad_length=1024,  # Pad prompt to 1024 tokens\n",
    "    max_out_length=3072,  # Allocate the rest of the cache for the output\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2e74b98",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Compile the model\n",
    "\n",
    "DiffusionGemma runs on JAX, which traces and compiles the model on the first call. After this one-time step (~1-2 minutes), all subsequent calls are fast."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dda169d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Compiling model (one-time)...\")\n",
    "_ = sampler.chat(\"warmup\")\n",
    "print(\"Done! Model is ready.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89428118",
   "metadata": {
    "executionInfo": {
     "elapsed": 54287,
     "status": "ok",
     "timestamp": 1781031179287,
     "user": {
      "displayName": "",
      "userId": ""
     },
     "user_tz": -60
    },
    "id": "89428118",
    "outputId": "217b12cb-f503-47c9-bca4-aa6de8fef739"
   },
   "outputs": [],
   "source": [
    "prompt = \"What are diffusion LLMs?\"\n",
    "print(f\"Prompt: {prompt}\")\n",
    "print(\"Generating response...\")\n",
    "response = sampler.chat(prompt)\n",
    "print(f\"Response:\\n{response}\")"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "last_runtime": {},
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}