guides/ipynb/keras_hub/getting_started.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "# Getting Started with KerasHub\n",
    "\n",
    "**Author:** [Matthew Watson](https://github.com/mattdangerw/), [Jonathan Bischof](https://github.com/jbischof)<br>\n",
    "**Date created:** 2022/12/15<br>\n",
    "**Last modified:** 2024/10/17<br>\n",
    "**Description:** An introduction to the KerasHub API."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "**KerasHub** is a pretrained modeling library that aims to be simple, flexible, and fast.\n",
    "The library provides [Keras 3](https://keras.io/keras_3/) implementations of popular\n",
    "model architectures, paired with a collection of pretrained checkpoints available on\n",
    "[Kaggle](https://www.kaggle.com/organizations/keras/models). Models can be used for both\n",
    "training and inference on any of the TensorFlow, Jax, and Torch backends.\n",
    "\n",
    "KerasHub is an extension of the core Keras API; KerasHub components are provided as\n",
    "`keras.Layer`s and `keras.Model`s. If you are familiar with Keras, congratulations! You\n",
    "already understand most of KerasHub.\n",
    "\n",
    "This guide is meant to be an accessible introduction to the entire library. We will start\n",
    "by using high-level APIs to classify images and generate text, then progressively show\n",
    "deeper customization of models and training. Throughout the guide, we use Professor Keras,\n",
    "the official Keras mascot, as a visual reference for the complexity of the material:\n",
    "\n",
    "![](https://storage.googleapis.com/keras-nlp/getting_started_guide/prof_keras_evolution.png)\n",
    "\n",
    "As always, we'll keep our Keras guides focused on real-world code examples. You can play\n",
    "with the code here at any time by clicking the Colab link at the top of the guide."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Installation and Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "To begin, let's install keras-hub. The library is available on PyPI, so we can simply\n",
    "install it with pip."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade --quiet keras-hub keras"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Keras 3 was built to work on top of TensorFlow, Jax, and Torch backends. You should\n",
    "specify the backend first thing when writing Keras code, before any library imports. We\n",
    "will use the Jax backend for this guide, but you can use `torch` or `tensorflow` without\n",
    "changing a single line in the rest of this guide. That's the power of Keras 3!\n",
    "\n",
    "We will also set `XLA_PYTHON_CLIENT_MEM_FRACTION`, which frees up the whole GPU for\n",
    "Jax to use from the start."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"KERAS_BACKEND\"] = \"jax\"  # or \"tensorflow\" or \"torch\"\n",
    "os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"] = \"1.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Lastly, we need to do some extra setup to access the models used in this guide. Many\n",
    "popular open LLMs, such as Gemma from Google and Llama from Meta, require accepting\n",
    "a community license before accessing the model weights. We will be using Gemma in this\n",
    "guide, so we can follow the following steps:\n",
    "\n",
    "1. Go to the [Gemma 2](https://www.kaggle.com/models/keras/gemma2) model page, and accept\n",
    "   the license at the banner at the top.\n",
    "2. Generate an Kaggle API key by going to [Kaggle settings](https://www.kaggle.com/settings)\n",
    "   and clicking \"Create New Token\" button under the \"API\" section.\n",
    "3. Inside your colab notebook, click on the key icon on the left hand toolbar. Add two\n",
    "   secrets: `KAGGLE_USERNAME` with your username, and `KAGGLE_KEY` with the API key you just\n",
    "   created. Make these secrets visible to the notebook you are running."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## API Quickstart\n",
    "\n",
    "Before we begin, let's take a look at the key classes we will use in the KerasHub library.\n",
    "\n",
    "* **Task**: e.g., `keras_hub.models.CausalLM`, `keras_hub.models.ImageClassifier`, and\n",
    "`keras_hub.models.TextClassifier`.\n",
    "  * **What it does**: A task maps from raw image, audio, and text inputs to model\n",
    "    predictions.\n",
    "  * **Why it's important**: A task is the highest-level entry point to the KerasHub API. It\n",
    "    encapsulates both preprocessing and modeling into a single, easy-to-use class. Tasks can\n",
    "    be used both for fine-tuning and inference.\n",
    "  * **Has a**: `backbone` and `preprocessor`.\n",
    "  * **Inherits from**: `keras.Model`.\n",
    "* **Backbone**: `keras_hub.models.Backbone`.\n",
    "  * **What it does**: Maps preprocessed tensor inputs to the latent space of the model.\n",
    "  * **Why it's important**: The backbone encapsulates the architecture and parameters of a\n",
    "    pretrained models in a way that is unspecialized to any particular task. A backbone can\n",
    "    be combined with arbitrary preprocessing and \"head\" layers mapping dense features to\n",
    "    predictions to accomplish any ML task.\n",
    "  * **Inherits from**: `keras.Model`.\n",
    "* **Preprocessor**: e.g.,`keras_hub.models.CausalLMPreprocessor`,\n",
    "  `keras_hub.models.ImageClassifierPreprocessor`, and\n",
    "  `keras_hub.models.TextClassifierPreprocessor`.\n",
    "  * **What it does**: A preprocessor maps from raw image, audio and text inputs to\n",
    "    preprocessed tensor inputs.\n",
    "  * **Why it's important**: A preprocessing layer encapsulates all tasks specific\n",
    "    preprocessing, e.g. image resizing and text tokenization, in a way that can be used\n",
    "    standalone to precompute preprocessed inputs. Note that if you are using a high-level\n",
    "    task class, this preprocessing is already baked in by default.\n",
    "  * **Has a**: `tokenizer`, `audio_converter`, and/or `image_converter`.\n",
    "  * **Inherits from**: `keras.layers.Layer`.\n",
    "* **Tokenizer**: `keras_hub.tokenizers.Tokenizer`.\n",
    "  * **What it does**: Converts strings to sequences of token ids.\n",
    "  * **Why it's important**: The raw bytes of a string are an inefficient representation of\n",
    "    text input, so we first map string inputs to integer token ids. This class encapsulated\n",
    "    the mapping of strings to ints and the reverse (via the `detokenize()` method).\n",
    "  * **Inherits from**: `keras.layers.Layer`.\n",
    "* **ImageConverter**: `keras_hub.layers.ImageConverter`.\n",
    "  * **What it does**: Resizes and rescales image input.\n",
    "  * **Why it's important**: Image models often need to normalize image inputs to a specific\n",
    "    range, or resizing inputs to a specific size. This class encapsulates the image-specific\n",
    "    preprocessing.\n",
    "  * **Inherits from**: `keras.layers.Layer`.\n",
    "* **AudioConveter**: `keras_hub.layers.AudioConveter`.\n",
    "  * **What it does**: Converts raw audio to model ready input.\n",
    "  * **Why it's important**: Audio models often need to preprocess raw audio input before\n",
    "    passing it to a model, e.g. by computing a spectrogram of the audio signal. This class\n",
    "    encapsulates the image specific preprocessing in an easy to use layer.\n",
    "  * **Inherits from**: `keras.layers.Layer`.\n",
    "\n",
    "All of the classes listed here have a `from_preset()` constructor, which will instantiate\n",
    "the component with weights and state for the given pre-trained model identifier. E.g.\n",
    "`keras_hub.tokenizers.Tokenizer.from_preset(\"gemma2_2b_en\")` will create a layer that\n",
    "tokenizes text using a Gemma2 tokenizer vocabulary.\n",
    "\n",
    "The figure below shows how all these core classes interact. Arrow indicate composition\n",
    "not inheritance (e.g., a task *has a* backbone).\n",
    "\n",
    "![png](/img/guides/getting_started/class-diagram.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Classify an image\n",
    "\n",
    "![](https://storage.googleapis.com/keras-nlp/getting_started_guide/prof_keras_beginner.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Enough setup! Let's have some fun with pre-trained models. Let's load a test image of a\n",
    "California Quail and classify it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import keras\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "image_url = \"https://upload.wikimedia.org/wikipedia/commons/a/aa/California_quail.jpg\"\n",
    "image_path = keras.utils.get_file(origin=image_url)\n",
    "image = keras.utils.load_img(image_path)\n",
    "plt.imshow(image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "We can use a ResNet vision model trained on the ImageNet-1k database. This model will\n",
    "give each input sample and output label from `[0, 1000)`, where each label corresponds to\n",
    "some real word entity, like a \"milk can\" or a \"porcupine.\" The dataset actually has a\n",
    "specific label for quail, at index 85. Let's download the model and predict a label."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import keras_hub\n",
    "\n",
    "image_classifier = keras_hub.models.ImageClassifier.from_preset(\n",
    "    \"resnet_50_imagenet\",\n",
    "    activation=\"softmax\",\n",
    ")\n",
    "batch = np.array([image])\n",
    "image_classifier.preprocessor.image_size = (224, 224)\n",
    "preds = image_classifier.predict(batch)\n",
    "preds.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "These ImageNet labels aren't a particularly \"human readable,\" so we can use a built-in\n",
    "utility function to decode the predictions to a set of class names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "keras_hub.utils.decode_imagenet_predictions(preds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Looking good! The model weights successfully downloaded, and we predicted the\n",
    "correct classification label for our quail image with near certainty.\n",
    "\n",
    "This was our first example of the high-level **task** API mentioned in the API quickstart\n",
    "above. An `keras_hub.models.ImageClassifier` is a task for classifying images, and can be\n",
    "used with a number of different model architectures (ResNet, VGG, MobileNet, etc). You\n",
    "can view the full list of models shipped directly by the Keras team on\n",
    "[Kaggle](https://www.kaggle.com/organizations/keras/models).\n",
    "\n",
    "A task is just a subclass of `keras.Model` \u2014 you can use `fit()`, `compile()`, and\n",
    "`save()` on our `classifier` object same as any other model. But tasks come with a few\n",
    "extras provided by the KerasHub library. The first and most important is `from_preset()`,\n",
    "a special constructor you will see on many classes in KerasHub.\n",
    "\n",
    "A **preset** is a directory of model state. It defines both the architecture we should\n",
    "load and the pretrained weights that go with it. `from_preset()` allows us to load\n",
    "**preset** directories from a number of different locations:\n",
    "\n",
    "- A local directory.\n",
    "- The Kaggle Model hub.\n",
    "- The HuggingFace model hub.\n",
    "\n",
    "You can take a look at the `keras_hub.models.ImageClassifier.from_preset` docs to better\n",
    "understand all the options when constructing a Keras model from a preset.\n",
    "\n",
    "All tasks use two main sub-objects. A `keras_hub.models.Backbone` and a\n",
    "`keras_hub.layers.Preprocessor`. You might be familiar already with the term **backbone**\n",
    "from computer vision, where it is often used to describe a feature extractor network that\n",
    "maps images to a latent space. A KerasHub backbone is this concept generalized, we use it\n",
    "to refer to any pretrained model without a task-specific head. That is, a KerasHub\n",
    "backbone maps raw images, audio and text (or a combination of these inputs) to a\n",
    "pretrained model's latent space. We can then map this latent space to any number of task\n",
    "specific outputs, depending on what we are trying to do with the model.\n",
    "\n",
    "A **preprocessor** is just a Keras layer that does all the preprocessing for a specific\n",
    "task. In our case, preprocessing with will resize our input image and rescale it to the\n",
    "range `[0, 1]` using some ImageNet specific mean and variance data. Let's call our\n",
    "task's preprocessor and backbone in succession to see what happens to our input shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "print(\"Raw input shape:\", batch.shape)\n",
    "resized_batch = image_classifier.preprocessor(batch)\n",
    "print(\"Preprocessed input shape:\", resized_batch.shape)\n",
    "hidden_states = image_classifier.backbone(resized_batch)\n",
    "print(\"Latent space shape:\", hidden_states.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Our raw image is rescaled to `(224, 224)` during preprocessing and finally\n",
    "downscaled to a `(7, 7)` image of 2048 feature vectors \u2014 the latent space of the\n",
    "ResNet model. Note that ResNet can actually handle images of arbitrary sizes,\n",
    "though performance will eventually fall off if your image is very different\n",
    "sized than the pretrained data. If you'd like to disable the resizing in the\n",
    "preprocessing layer, you can run `image_classifier.preprocessor.image_size = None`.\n",
    "\n",
    "If you are ever wondering the exact structure of the task you loaded, you can\n",
    "use `model.summary()` same as any Keras model. The model summary for tasks will\n",
    "included extra information on model preprocessing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "image_classifier.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Generate text with an LLM\n",
    "\n",
    "![](https://storage.googleapis.com/keras-nlp/getting_started_guide/prof_keras_intermediate.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Next up, let's try working with and generating text. The task we can use when generating\n",
    "text is `keras_hub.models.CausalLM` (where LM is short for **L**anguage **M**odel). Let's\n",
    "download the 2 billion parameter Gemma 2 model and try it out.\n",
    "\n",
    "Since this is about 100x larger model than the ResNet model we just downloaded, we need to be\n",
    "a little more careful about our GPU memory usage. We can use a half-precision type to\n",
    "load each parameter of our ~2.5 billion as a two-byte float instead of four. To do this\n",
    "we can pass `dtype` to the `from_preset()` constructor. `from_preset()` will forward any\n",
    "kwargs to the main constructor for the class, so you can pass kwargs that work on all\n",
    "Keras layers like `dtype`, `trainable`, and `name`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "causal_lm = keras_hub.models.CausalLM.from_preset(\n",
    "    \"gemma2_instruct_2b_en\",\n",
    "    dtype=\"bfloat16\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "The model we just loaded was an instruction-tuned version of Gemma, which means the model\n",
    "was further fine-tuned for chat. We can take advantage of these capabilities as long as\n",
    "we stick to the particular template for text used when training the model. These special\n",
    "tokens vary per model and can be hard to track, the [Kaggle model\n",
    "page](https://www.kaggle.com/models/keras/gemma2/) will contain details such as this.\n",
    "\n",
    "`CausalLM` come with an extra function called `generate()` which can be used generate\n",
    "predict tokens in a loop and decode them as a string."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "template = \"<start_of_turn>user\\n{question}<end_of_turn>\\n<start_of_turn>model\"\n",
    "\n",
    "question = \"\"\"Write a python program to generate the first 1000 prime numbers.\n",
    "Just show the actual code.\"\"\"\n",
    "print(causal_lm.generate(template.format(question=question), max_length=512))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Note that on the Jax and TensorFlow backends, this `generate()` function is compiled, so\n",
    "the second time you call for the same `max_length`, it will actually be much faster.\n",
    "KerasHub will use Jax an TensorFlow to compute an optimized version of the generation\n",
    "computational graph that can be reused."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "question = \"Share a very simple brownie recipe.\"\n",
    "print(causal_lm.generate(template.format(question=question), max_length=512))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "As with our image classifier, we can use model summary to see the details of our task\n",
    "setup, including preprocessing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "causal_lm.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Our text preprocessing includes a tokenizer, which is how all KerasHub models handle\n",
    "input text. Let's try using it directly to get a better sense of how it works. All\n",
    "tokenizers include `tokenize()` and `detokenize()` methods, to map strings to integer\n",
    "sequences and integer sequences to strings. Directly calling the layer with\n",
    "`tokenizer(inputs)` is equivalent to calling `tokenizer.tokenize(inputs)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "tokenizer = causal_lm.preprocessor.tokenizer\n",
    "tokens_ids = tokenizer.tokenize(\"The quick brown fox jumps over the lazy dog.\")\n",
    "print(tokens_ids)\n",
    "string = tokenizer.detokenize(tokens_ids)\n",
    "print(string)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "The `generate()` function for `CausalLM` models involved a sampling step. The Gemma model\n",
    "will be called once for each token we want to generate, and return a probability\n",
    "distribution over all tokens. This distribution is then sampled to choose the next token\n",
    "in the sequence.\n",
    "\n",
    "For Gemma models, we default to greedy sampling, meaning we simply pick the most likely\n",
    "output from the model at each step. But we can actually control this process with an\n",
    "extra `sampler` argument to the standard `compile` function on all Keras models. Let's\n",
    "try it out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "causal_lm.compile(\n",
    "    sampler=keras_hub.samplers.TopKSampler(k=10, temperature=2.0),\n",
    ")\n",
    "\n",
    "question = \"Share a very simple brownie recipe.\"\n",
    "print(causal_lm.generate(template.format(question=question), max_length=512))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Here we used a Top-K sampler, meaning we will randomly sample the partial distribution formed\n",
    "by looking at just the top 10 predicted tokens at each time step. We also pass a `temperature` of 2,\n",
    "which flattens our predicted distribution before we sample.\n",
    "\n",
    "The net effect is that we will explore our model's distribution much more broadly each\n",
    "time we generate output. Generation will now be a random process, each time we re-run\n",
    "generate we will get a different result. We can note that the results feel \"looser\" than\n",
    "greedy search \u2014 more minor mistakes and a less consistent tone.\n",
    "\n",
    "You can look at all the samplers Keras supports at [keras_hub.samplers](https://keras.io/api/keras_hub/samplers/).\n",
    "\n",
    "Let's free up the memory from our large Gemma model before we jump to the next section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "del causal_lm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Fine-tune and publish an image classifier\n",
    "\n",
    "![](https://storage.googleapis.com/keras-nlp/getting_started_guide/prof_keras_advanced.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Now that we've tried running inference for both images and text, let's try running\n",
    "training. We will take our ResNet image classifier from earlier and fine-tune it on\n",
    "simple cats vs dogs dataset. We can start by downloading and extracting the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import pathlib\n",
    "\n",
    "extract_dir = keras.utils.get_file(\n",
    "    \"cats_vs_dogs\",\n",
    "    \"https://download.microsoft.com/download/3/E/1/3E1C3F21-ECDB-4869-8368-6DEBA77B919F/kagglecatsanddogs_5340.zip\",\n",
    "    extract=True,\n",
    ")\n",
    "data_dir = pathlib.Path(extract_dir) / \"PetImages\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "When working with lots of real-world image data, corrupted images are a common occurrence.\n",
    "Let's filter out badly-encoded images that do not feature the string \"JFIF\" in their\n",
    "header."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "num_skipped = 0\n",
    "\n",
    "for path in data_dir.rglob(\"*.jpg\"):\n",
    "    with open(path, \"rb\") as file:\n",
    "        is_jfif = b\"JFIF\" in file.peek(10)\n",
    "    if not is_jfif:\n",
    "        num_skipped += 1\n",
    "        os.remove(path)\n",
    "\n",
    "print(f\"Deleted {num_skipped} images.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "We can load the dataset with `keras.utils.image_dataset_from_directory`. One important\n",
    "thing to note here is that the `train_ds` and `val_ds` will both be returned as\n",
    "`tf.data.Dataset` objects, including on the `torch` and `jax` backends.\n",
    "\n",
    "KerasHub will use [tf.data](https://www.tensorflow.org/guide/data) as the default API for\n",
    "running multi-threaded preprocessing on the CPU. `tf.data` is a powerful API for training\n",
    "input pipelines that can scale up to complex, multi-host training jobs easily. Using it\n",
    "does not restrict your choice of backend, a `tf.data.Dataset` can be as an iterator of\n",
    "regular numpy data and passed to `fit()` on any Keras backend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "train_ds, val_ds = keras.utils.image_dataset_from_directory(\n",
    "    data_dir,\n",
    "    validation_split=0.2,\n",
    "    subset=\"both\",\n",
    "    seed=1337,\n",
    "    image_size=(256, 256),\n",
    "    batch_size=32,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "At its simplest, training our classifier could consist of simply calling `fit()` on our\n",
    "model with our dataset. But to make this example a little more interesting, let's show\n",
    "how to customize preprocessing within a task.\n",
    "\n",
    "In the first example, we saw how, by default, the preprocessing for our ResNet model resized\n",
    "and rescaled our input. This preprocessing can be customized when we create our model. We\n",
    "can use Keras' image preprocessing layers to create a `keras.layers.Pipeline` that will\n",
    "rescale, randomly flip, and randomly rotate our input images. These random image\n",
    "augmentations will allow our smaller dataset to function as a larger, more varied one.\n",
    "Let's try it out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "preprocessor = keras.layers.Pipeline(\n",
    "    [\n",
    "        keras.layers.Rescaling(1.0 / 255),\n",
    "        keras.layers.RandomFlip(\"horizontal\"),\n",
    "        keras.layers.RandomRotation(0.2),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Now that we have created a new layer for preprocessing, we can simply pass it to the\n",
    "`ImageClassifier` during the `from_preset()` constructor. We can also pass\n",
    "`num_classes=2` to match our two labels for \"cat\" and \"dog.\" When `num_classes` is\n",
    "specified like this, our head weights for the model will be randomly initialized\n",
    "instead of containing the weights for our 1000 class image classification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "image_classifier = keras_hub.models.ImageClassifier.from_preset(\n",
    "    \"resnet_50_imagenet\",\n",
    "    activation=\"softmax\",\n",
    "    num_classes=2,\n",
    "    preprocessor=preprocessor,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Note that if you want to preprocess your input data outside of Keras, you can simply\n",
    "pass `preprocessor=None` to the task `from_preset()` call. In this case, KerasHub will\n",
    "apply no preprocessing at all, and you are free to preprocess your data with any library\n",
    "or workflow before passing your data to `fit()`.\n",
    "\n",
    "Next, we can compile our model for fine-tuning. A KerasHub task is just a regular\n",
    "`keras.Model` with some extra functionality, so we can `compile()` as normal for a\n",
    "classification task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "image_classifier.compile(\n",
    "    optimizer=keras.optimizers.Adam(1e-4),\n",
    "    loss=\"sparse_categorical_crossentropy\",\n",
    "    metrics=[\"accuracy\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "With that, we can simply run `fit()`. The image classifier will automatically apply our\n",
    "preprocessing to each batch when training the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "image_classifier.fit(\n",
    "    train_ds,\n",
    "    validation_data=val_ds,\n",
    "    epochs=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "After three epochs of data, we achieve 99% accuracy on our cats vs dogs\n",
    "validation dataset. This is unsurprising, given that the ImageNet pretrained weights we began\n",
    "with could already classify some breeds of cats and dogs individually.\n",
    "\n",
    "Now that we have a fine-tuned model let's try saving it. You can create a new saved preset with a\n",
    "fine-tuned model for any task simply by running `task.save_to_preset()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "image_classifier.save_to_preset(\"cats_vs_dogs\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "One of the most powerful features of KerasHub is the ability upload models to Kaggle or\n",
    "Huggingface models hub and share them with others. `keras_hub.upload_preset` allows you\n",
    "to upload a saved preset.\n",
    "\n",
    "In this case, we will upload to Kaggle. We have already authenticated with Kaggle to,\n",
    "download the Gemma model earlier. Running the following cell well upload a new model\n",
    "to Kaggle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "from google.colab import userdata\n",
    "\n",
    "username = userdata.get(\"KAGGLE_USERNAME\")\n",
    "keras_hub.upload_preset(\n",
    "    f\"kaggle://{username}/resnet/keras/cats_vs_dogs\",\n",
    "    \"cats_vs_dogs\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Let's take a look at a test image from our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "image = keras.utils.load_img(data_dir / \"Cat\" / \"6779.jpg\")\n",
    "plt.imshow(image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "If we wait for a few minutes for our model upload to finish processing on the Kaggle\n",
    "side, we can go ahead and download the model we just created and use it to classify this\n",
    "test image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "image_classifier = keras_hub.models.ImageClassifier.from_preset(\n",
    "    f\"kaggle://{username}/resnet/keras/cats_vs_dogs\",\n",
    ")\n",
    "print(image_classifier.predict(np.array([image])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Congratulations on uploading your first model with KerasHub! If you want to share your\n",
    "work with others, you can go to the model link printed out when we uploaded the model, and\n",
    "turn the model public in settings.\n",
    "\n",
    "Let's delete this model to free up memory before we move on to our final example for this\n",
    "guide."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "del image_classifier"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Building a custom text classifier\n",
    "\n",
    "![](https://storage.googleapis.com/keras-nlp/getting_started_guide/prof_keras_expert.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "As a final example for this getting started guide, let's take a look at how we can build\n",
    "custom models from lower-level Keras and KerasHub components. We will build a text\n",
    "classifier to classify movie reviews in the IMDb dataset as either positive or negative.\n",
    "\n",
    "Let's download the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "extract_dir = keras.utils.get_file(\n",
    "    \"imdb_reviews\",\n",
    "    origin=\"https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz\",\n",
    "    extract=True,\n",
    ")\n",
    "data_dir = pathlib.Path(extract_dir) / \"aclImdb\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "The IMDb dataset contrains a large amount of unlabeled movie reviews. We don't need those\n",
    "here, we can simply delete them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import shutil\n",
    "\n",
    "shutil.rmtree(data_dir / \"train\" / \"unsup\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Next up, we can load our data with `keras.utils.text_dataset_from_directory`. As with our\n",
    "image dataset creation above, the returned datasets will be `tf.data.Dataset` objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "raw_train_ds = keras.utils.text_dataset_from_directory(\n",
    "    data_dir / \"train\",\n",
    "    batch_size=2,\n",
    ")\n",
    "raw_val_ds = keras.utils.text_dataset_from_directory(\n",
    "    data_dir / \"test\",\n",
    "    batch_size=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "KerasHub is designed to be a layered API. At the top-most level, tasks aim to make it\n",
    "easy to quickly tackle a problem. We could keep using the task API here, and create a\n",
    "`keras_hub.models.TextClassifer` for a text classification model like BERT, and fine-tune\n",
    "it in 10 or so lines of code.\n",
    "\n",
    "Instead, to make our final example a little more interesting, let's show how we can use\n",
    "lower-level API components to do something that isn't directly baked in to the library.\n",
    "We will take the Gemma 2 model we used earlier, which is usually used for generating text,\n",
    "and modify it to output classification predictions.\n",
    "\n",
    "A common approach for classifying with a generative model would keep using it in a generative\n",
    "context, by prompting it with the review and a question (`\"Is this review positive or negative?\"`).\n",
    "But making an actual classifier is more useful if you want an actual probability score associated\n",
    "with your labels.\n",
    "\n",
    "Instead of loading the Gemma 2 model through the `CausalLM` task, we can load two\n",
    "lower-level components: a **backbone** and a **tokenizer**. Much like the task classes we have\n",
    "used so far, `keras_hub.models.Backbone` and `keras_hub.tokenizers.Tokenizer` both have a\n",
    "`from_preset()` constructor for loading pretrained models. If you are running this code,\n",
    "you will note you don't have to wait for a download as we use the model a second time,\n",
    "the weights files are cached locally the first time we use the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "tokenizer = keras_hub.tokenizers.Tokenizer.from_preset(\n",
    "    \"gemma2_instruct_2b_en\",\n",
    ")\n",
    "backbone = keras_hub.models.Backbone.from_preset(\n",
    "    \"gemma2_instruct_2b_en\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "We saw what the tokenizer does in the second example of this guide. We can use it to map\n",
    "from string inputs to token ids in a way that matches the pretrained weights of the Gemma\n",
    "model.\n",
    "\n",
    "The backbone will map from a sequence of token ids to a sequence of embedded tokens in\n",
    "the latent space of the model. We can use this rich representation to build a classifier.\n",
    "\n",
    "Let's start by defining a custom preprocessing routine. `keras_hub.layers` contains a\n",
    "collection of modeling and preprocessing layers, included some layers for token\n",
    "preprocessing. We can use `keras_hub.layers.StartEndPacker`, which will append a special\n",
    "start token to the beginning of each review, a special end token to the end, and finally\n",
    "truncate or pad each review to a fixed length.\n",
    "\n",
    "If we combine this with our `tokenizer`, we can build a preprocessing function that will\n",
    "output batches of token ids with shape `(batch_size, sequence_length)`. We should also\n",
    "output a padding mask that marks which tokens are padding tokens, so we can later exclude\n",
    "these positions from our Transformer's attention computation. Most Transformer backbones\n",
    "in KerasNLP take in a `\"padding_mask\"` input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "packer = keras_hub.layers.StartEndPacker(\n",
    "    start_value=tokenizer.start_token_id,\n",
    "    end_value=tokenizer.end_token_id,\n",
    "    pad_value=tokenizer.pad_token_id,\n",
    "    sequence_length=None,\n",
    ")\n",
    "\n",
    "\n",
    "def preprocess(x, y=None, sequence_length=256):\n",
    "    x = tokenizer(x)\n",
    "    x = packer(x, sequence_length=sequence_length)\n",
    "    x = {\n",
    "        \"token_ids\": x,\n",
    "        \"padding_mask\": x != tokenizer.pad_token_id,\n",
    "    }\n",
    "    return keras.utils.pack_x_y_sample_weight(x, y)\n",
    ""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "With our preprocessing defined, we can simply use `tf.data.Dataset.map` to apply our\n",
    "preprocessing to our input data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "train_ds = raw_train_ds.map(preprocess, num_parallel_calls=16)\n",
    "val_ds = raw_val_ds.map(preprocess, num_parallel_calls=16)\n",
    "next(iter(train_ds))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Running fine-tuning on a 2.5 billion parameter model is quite expensive compared to the\n",
    "image classifier we trained earlier, for the simple reason that this model is 100x the\n",
    "size of ResNet! To speed things up a bit, let's reduce the size of our training data to a\n",
    "tenth of the original size. Of course, this is leaving some performance on the table\n",
    "compared to full training, but it will keep things running quickly for our guide."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "train_ds = train_ds.take(1000)\n",
    "val_ds = val_ds.take(1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Next, we need to attach a classification head to our backbone model. In general, text\n",
    "transformer backbones will output a tensor with shape\n",
    "`(batch_size, sequence_length, hidden_dim)`. The main thing we will need to\n",
    "classify with this input is to pool on the sequence dimension so we have a single\n",
    "feature vector per input example.\n",
    "\n",
    "Since the Gemma model is a generative model, information only passed from left to right\n",
    "in the sequence. The only token representation that can \"see\" the entire movie review\n",
    "input is the final token in each review. We can write a simple pooling layer to do this \u2014\n",
    "we will simply grab the last non-padding position of each input sequence. There's no special\n",
    "process to writing a layer like this, we can use Keras and `keras.ops` normally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "from keras import ops\n",
    "\n",
    "\n",
    "class LastTokenPooler(keras.layers.Layer):\n",
    "    def call(self, inputs, padding_mask):\n",
    "        end_positions = ops.sum(padding_mask, axis=1, keepdims=True) - 1\n",
    "        end_positions = ops.cast(end_positions, \"int\")[:, :, None]\n",
    "        outputs = ops.take_along_axis(inputs, end_positions, axis=1)\n",
    "        return ops.squeeze(outputs, axis=1)\n",
    ""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "With this pooling layer, we are ready to write our Gemma classifier. All task and backbone\n",
    "models in KerasHub are [functional](https://keras.io/guides/functional_api/) models, so\n",
    "we can easily manipulate the model structure. We will call our backbone on our inputs, add\n",
    "our new pooling layer, and finally add a small feedforward network with a `\"relu\"` activation\n",
    "in the middle. Let's try it out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "inputs = backbone.input\n",
    "x = backbone(inputs)\n",
    "x = LastTokenPooler(\n",
    "    name=\"pooler\",\n",
    ")(x, inputs[\"padding_mask\"])\n",
    "x = keras.layers.Dense(\n",
    "    2048,\n",
    "    activation=\"relu\",\n",
    "    name=\"pooled_dense\",\n",
    ")(x)\n",
    "x = keras.layers.Dropout(\n",
    "    0.1,\n",
    "    name=\"output_dropout\",\n",
    ")(x)\n",
    "outputs = keras.layers.Dense(\n",
    "    2,\n",
    "    activation=\"softmax\",\n",
    "    name=\"output_dense\",\n",
    ")(x)\n",
    "text_classifier = keras.Model(inputs, outputs)\n",
    "text_classifier.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "Before we train, there is one last trick we should employ to make this code run on free\n",
    "tier colab GPUs. We can see from our model summary our model takes up almost 10 gigabytes\n",
    "of space. An optimizer will need to make multiple copies of each parameter during\n",
    "training, taking the total space of our model during training close to 30 or 40\n",
    "gigabytes.\n",
    "\n",
    "This would OOM many GPUs. A useful trick we can employ is to enable LoRA on our\n",
    "backbone. LoRA is an approach which freezes the entire model, and only trains a low-parameter\n",
    "decomposition of large weight matrices. You can read more about LoRA in this [Keras\n",
    "example](https://keras.io/examples/nlp/parameter_efficient_finetuning_of_gpt2_with_lora/).\n",
    "Let's try enabling it and re-printing our summary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "backbone.enable_lora(4)\n",
    "text_classifier.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "After enabling LoRA, our model goes from 10GB of traininable parameters to just 20MB.\n",
    "That means the space used by optimizer variables will no longer be a concern.\n",
    "\n",
    "With all that set up, we can compile and train our model as normal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "text_classifier.compile(\n",
    "    optimizer=keras.optimizers.Adam(5e-5),\n",
    "    loss=\"sparse_categorical_crossentropy\",\n",
    "    metrics=[\"accuracy\"],\n",
    ")\n",
    "text_classifier.fit(\n",
    "    train_ds,\n",
    "    validation_data=val_ds,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "We are able to achieve over ~93% accuracy on the movie review sentiment\n",
    "classification problem. This is not bad, given that we only used a 10th of our\n",
    "original dataset to train.\n",
    "\n",
    "Taken together, the `backbone` and `tokenizer` we created in this example\n",
    "allowed us access the full power of pretrained Gemma checkpoints, without\n",
    "restricting what we could do with them. This is a central aim of the KerasHub\n",
    "API. Simple workflows should be easy, and as you go deeper, you gain access to a\n",
    "deeply customizable set of building blocks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Going further\n",
    "\n",
    "This is just scratching the surface of what you can do with the KerasHub.\n",
    "\n",
    "This guide shows a few of the high-level tasks that we ship with the KerasHub library,\n",
    "but there are many tasks we did not cover here. Try [generating images with Stable\n",
    "Diffusion](https://keras.io/guides/keras_hub/stable_diffusion_3_in_keras_hub/), for\n",
    "example.\n",
    "\n",
    "The most significant advantage of KerasHub is it gives you the flexibility to combine pre-trained\n",
    "building blocks with the full power of Keras 3. You can train large LLMs on TPUs with model\n",
    "parallelism with the [keras.distribution](https://keras.io/guides/distribution/) API. You can\n",
    "quantize models with Keras' [quatize\n",
    "method](https://keras.io/examples/keras_recipes/float8_training_and_inference_with_transfo\n",
    "rmer/). You can write custom training loops and even mix in direct Jax, Torch, or\n",
    "Tensorflow calls.\n",
    "\n",
    "See [keras.io/keras_hub](https://keras.io/keras_hub/) for a full list of guides and\n",
    "examples to continue digging into the library."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "getting_started",
   "private_outputs": false,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}