text-bison-notebook.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# 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\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fsv4jGuU89rX"
      },
      "source": [
        "# PaLM 2 Text Bison Reinforcement Learning from Human Feedback"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "overview:automl"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This project showcases use of Reinforcement Learning from Human Feedback (RLHF) to tune the [PaLM 2 Text Bison][bison] model from Vertex AI Model Garden. Learn more about [RLHF tuning][rlhf] on Vertex AI.\n",
        "\n",
        "[vertex]: https://cloud.google.com/vertex-ai\n",
        "[ar]: https://cloud.google.com/artifact-registry\n",
        "[cb]: https://cloud.google.com/build\n",
        "[cd]: https://cloud.google.com/deploy\n",
        "[cs]: https://cloud.google.com/storage\n",
        "[pytorch]: https://pytorch.org/\n",
        "[tune]: https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-models\n",
        "[bison]: https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text\n",
        "[rlhf]: https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-text-models-rlhf "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "objective:pipelines,automl"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you will learn to tune and deploy a large language model model. To facilitate model development, deployment, and management, the project leverages various services including [Vertex AI][vertex], [Artifact Registry][ar], [Cloud Build][cb], [Cloud Deploy][cd] and [Cloud Storage][cs].\n",
        "  \n",
        "\n",
        "[vertex]: https://cloud.google.com/vertex-ai\n",
        "[ar]: https://cloud.google.com/artifact-registry\n",
        "[cb]: https://cloud.google.com/build\n",
        "[cd]: https://cloud.google.com/deploy\n",
        "[cs]: https://cloud.google.com/storage"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "costs"
      },
      "source": [
        "### Costs\n",
        "\n",
        "Learn about pricing for [Vertex AI](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage](https://cloud.google.com/storage/pricing), [Cloud Build](https://cloud.google.com/build/pricing), [Cloud Deploy](https://cloud.google.com/deploy/pricing) and [Artifact Registry](https://cloud.google.com/artifact-registry/pricing). Use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Setup\n",
        "\n",
        "Before you begin, use the [setup](/docs/SETUP.md) guide to configure your project, service account, permissions, storage and more before executing this notebook."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_aip:mbsdk"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the following packages required to execute this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F9dJ5Of-dORl"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet -r requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "restart"
      },
      "source": [
        "### Colab Only"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HKDBjmNr9T6t"
      },
      "source": [
        "Automatically restart kernel after installs so that your environment can access the new packages. Uncomment the cell below to restart the kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D-ZBOjErv5mM"
      },
      "outputs": [],
      "source": [
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "before_you_begin:nogpu"
      },
      "source": [
        "## Configure variables\n",
        "\n",
        "### Set Project ID"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "set_project_id"
      },
      "outputs": [],
      "source": [
        "import google.auth\n",
        "\n",
        "_, PROJECT_ID = google.auth.default()\n",
        "print(\"Project ID: \", PROJECT_ID)\n",
        "\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "### Set Region\n",
        "\n",
        "Set the `REGION` variable; it defaults to `\"us-central1\"`. For now, the only other supported region is `europe-west4`. Learn more about Vertex AI [regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2dw8q9fdQEH5"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\" # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### (Optional) Update Configuration\n",
        "\n",
        "If needed, you can update the default configuration for [tuner](/pkg/tuner/metadata.py) and [predictor](/pkg/tuner/metadata.py) packages."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gcp_authenticate"
      },
      "source": [
        "## Authenticate Google Cloud Account\n",
        "\n",
        "Depending on your Jupyter environment, you may need to authenticate manually. Follow the relevant instructions below.\n",
        "\n",
        "**1. Vertex AI Workbench**\n",
        "\n",
        "Do nothing as you are already authenticated."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0367eac06a10"
      },
      "source": [
        "**3. Vertex AI Colab**\n",
        "\n",
        "If you are using Colab, uncomment and run:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "21ad4dbb4a61"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n",
        "**2. Local**\n",
        "\n",
        "If you are running the notebook locally, uncomment and run:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ce6043da7b33"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "setup_vars"
      },
      "source": [
        "## Import Libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ec798de2db0d"
      },
      "outputs": [],
      "source": [
        "%load_ext autoreload\n",
        "%autoreload 2\n",
        "\n",
        "import google.cloud.aiplatform as aiplatform\n",
        "from kfp import compiler\n",
        "\n",
        "from pkg.tuner import metadata, parameters, registry, rlhf, steps\n",
        "from pkg.predictor import online"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk"
      },
      "source": [
        "## Initialize Vertex AI SDK\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rUP-o_VB4JEY"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=metadata.BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aef4f59195ad"
      },
      "source": [
        "## Move Data to Cloud Storage Bucket\n",
        "\n",
        "RLHF tuning uses three datasets:\n",
        "\n",
        "* Prompt [dataset](/data/prompt/shard-00000-of-00002.jsonl) with unlabelled prompts.\n",
        "* Human preference [dataset](/data/preference/shard-00000-of-00002.jsonl) with prompts labelled with preferences from humans.\n",
        "* Evaluation [dataset](/data/evaluation/shard-00000-of-00002.jsonl) with unlabeled prompts for prediction after the model is tuned. This is optional.\n",
        "  \n",
        "Learn more about data for RLFH in Vertex AI [documentation]. \n",
        "\n",
        "[documentation]: https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-text-models-rlhf#prepare_rlhf_tuning_datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "! gsutil -m cp -r data {metadata.BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A5hpItbK_KOc"
      },
      "source": [
        "## Compile the RLHF Pipeline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dH7V802zA4GW"
      },
      "source": [
        "Compile the pipeline into a YAML file that will be submitted to Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SkGgfDgr_20j"
      },
      "outputs": [],
      "source": [
        "compiler.Compiler().compile(\n",
        "    pipeline_func=rlhf.tune, package_path=metadata.COMPILED_PIPELINE_PATH\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Upload the compiled pipeline into Vertex AI Pipeline Registry which is backed by Artifact Registry."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "PIPELINE_TEMPLATE_URI = registry.upload(\n",
        "  project_id=PROJECT_ID,\n",
        "  region=REGION,\n",
        "  pipeline_registry=metadata.PIPELINE_REGISTRY,\n",
        "  compiled_pipeline_path=metadata.COMPILED_PIPELINE_PATH\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Calculate Training Steps\n",
        "\n",
        "Choose a suitable value for the number of reward model and reinforcing learning train steps to avoid overfitting. This depends on the size of the datasets. If needed, configure the values in [metadata.py](/pkg/trainer/metadata.py)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "REWARD_MODEL_TRAIN_STEPS = steps.get_reward_model_train_steps()\n",
        "\n",
        "REINFORCEMENT_LEARNING_TRAIN_STEPS = steps.get_reinforcement_learning_train_steps()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3o6EaTXJ4JEY"
      },
      "source": [
        "## Construct the Pipeline Job and Run on Vertex AI\n",
        "\n",
        "Define a pipeline job via the pipeline template compiled in the previous step.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "df55e79c"
      },
      "outputs": [],
      "source": [
        "job = aiplatform.PipelineJob(\n",
        "    display_name=metadata.MODEL_DISPLAY_NAME,\n",
        "    pipeline_root=metadata.PIPELINE_ROOT,\n",
        "    template_path=PIPELINE_TEMPLATE_URI,\n",
        "    parameter_values=parameters.get_values(\n",
        "        preference_dataset=metadata.PREFERENCE_DATASET,\n",
        "        prompt_dataset=metadata.PROMPT_DATASET,\n",
        "        eval_dataset=metadata.EVALUATION_DATASET,\n",
        "        reward_model_train_steps=REWARD_MODEL_TRAIN_STEPS,\n",
        "        reinforcement_learning_train_steps=REINFORCEMENT_LEARNING_TRAIN_STEPS\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lz99NXISmL0I"
      },
      "source": [
        "Run the pipeline on Vertex AI. Note that this will take about 2 hours."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CtAF_q4I4JEZ"
      },
      "outputs": [],
      "source": [
        "job.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1jSGT6VUo3FT"
      },
      "source": [
        "## Get Predictions\n",
        "\n",
        "Make predictions using either [Online Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-online-predictions) or [Batch Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-batch-predictions).\n",
        "\n",
        "First, get the endpoint resource name."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zFZgJ4uxrb0p"
      },
      "outputs": [],
      "source": [
        "stdout = ! gcloud ai models list --project=$PROJECT_ID --region=$REGION --filter=\"DISPLAY_NAME: $metadata.MODEL_DISPLAY_NAME\" --sort-by=~creationTimestamp --limit=1 --format=\"flattened(deployedModels[0].endpoint)\" 2>/dev/null\n",
        "ENDPOINT_RESOURCE = stdout[1].split()[1]\n",
        "print(\"Endpoint resource:\", ENDPOINT_RESOURCE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Online Predictions "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Invoke the endpoint where the model is deployed to make predictions for some test instances."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "PROMPT = \"How can I use VR in fitness?\"\n",
        "\n",
        "online.send(\n",
        "    resource=ENDPOINT_RESOURCE,\n",
        "    prompt=PROMPT\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Batch Predictions\n",
        "\n",
        "Coming soon..."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "text-bison-notebook.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.11.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}