colabs/sampling_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SC77q_zBESaM"
   },
   "source": [
    "Copyright 2024 DeepMind Technologies Limited.\n",
    "\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at\n",
    "\n",
    "http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TpESp4p5ESaM"
   },
   "source": [
    "# Getting Started with Gemma Sampling: A Step-by-Step Guide\n",
    "\n",
    "You will find in this colab a detailed tutorial explaining how to load a Gemma checkpoint and sample from it.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2lFHiDvvESaM"
   },
   "source": [
    "## Setup\n",
    "\n",
    "Please follow installation instructions at https://github.com/google-deepmind/gemma/blob/main/README.md."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "cellView": "form",
    "id": "O-sxcasvESaM"
   },
   "outputs": [],
   "source": [
    "# @title Download the checkpoints\n",
    "# Download the Flax checkpoints from https://www.kaggle.com/models/google/gemma\n",
    "# and put the local paths below.\n",
    "\n",
    "ckpt_path = ''\n",
    "vocab_path = ''\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "cellView": "form",
    "id": "-jpTUa1YESaM"
   },
   "outputs": [],
   "source": [
    "# @title Python imports\n",
    "from gemma import params as params_lib\n",
    "from gemma import sampler as sampler_lib\n",
    "from gemma import transformer as transformer_lib\n",
    "import sentencepiece as spm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4fDQsC87ESaN"
   },
   "source": [
    "## Start Generating with Your Model\n",
    "\n",
    "Load and prepare your LLM's checkpoint for use with Flax."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "cellView": "form",
    "id": "57nMYQ4HESaN"
   },
   "outputs": [],
   "source": [
    "# Load parameters\n",
    "params = params_lib.load_and_format_params(ckpt_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NWJ3UvHXESaN"
   },
   "source": [
    "Load your tokenizer, which we'll construct using the [SentencePiece](https://github.com/google/sentencepiece) library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "cellView": "form",
    "id": "khXrjEF0ESaN"
   },
   "outputs": [],
   "source": [
    "vocab = spm.SentencePieceProcessor()\n",
    "vocab.Load(vocab_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tCRtZMg0ESaN"
   },
   "source": [
    "Use the `transformer_lib.TransformerConfig.from_params` function to automatically load the correct configuration from a checkpoint. Note that the vocabulary size is smaller than the number of input embeddings due to unused tokens in this release."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "cellView": "form",
    "id": "7InOzQtcESaN"
   },
   "outputs": [],
   "source": [
    "transformer_config=transformer_lib.TransformerConfig.from_params(\n",
    "    params,\n",
    "    cache_size=1024  # Number of time steps in the transformer's cache\n",
    ")\n",
    "transformer = transformer_lib.Transformer(transformer_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KaU-X3_jESaN"
   },
   "source": [
    "Finally, build a sampler on top of your model and your tokenizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "cellView": "form",
    "id": "bdstASGrESaN"
   },
   "outputs": [],
   "source": [
    "# Create a sampler with the right param shapes.\n",
    "sampler = sampler_lib.Sampler(\n",
    "    transformer=transformer,\n",
    "    vocab=vocab,\n",
    "    params=params['transformer'],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C1fLns-_ESaN"
   },
   "source": [
    "You're ready to start sampling ! This sampler uses just-in-time compilation, so changing the input shape triggers recompilation, which can slow things down. For the fastest and most efficient results, keep your batch size consistent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "cellView": "form",
    "id": "qA0BhNQvESaN"
   },
   "outputs": [],
   "source": [
    "input_batch = [\n",
    "    \"\\n# Python program for implementation of Bubble Sort\\n\\ndef bubbleSort(arr):\",\n",
    "    \"What are the planets of the solar system?\",\n",
    "  ]\n",
    "\n",
    "out_data = sampler(\n",
    "    input_strings=input_batch,\n",
    "    total_generation_steps=300,  # number of steps performed when generating\n",
    "  )\n",
    "\n",
    "for input_string, out_string in zip(input_batch, out_data.text):\n",
    "  print(f\"Prompt:\\n{input_string}\\nOutput:\\n{out_string}\")\n",
    "  print()\n",
    "  print(10*'#')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tqbJ1SUcESaN"
   },
   "source": [
    "You should get an implementation of bubble sort and a description of the solar system.\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "private_outputs": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}