notebooks/official/prediction/llm_streaming_prediction.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "copyright"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 Google LLC\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",
        "#     https://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": "title:generic,gcp"
      },
      "source": [
        "# Vertex AI LLM and streaming prediction\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official//prediction/llm_streaming_prediction.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fprediction%2Fllm_streaming_prediction.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>    \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/prediction/llm_streaming_prediction.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/prediction/llm_streaming_prediction.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "overview:mlops"
      },
      "source": [
        "## Overview\n",
        "\n",
        "\n",
        "This tutorial demonstrates how to use Vertex AI LLM for making streaming predictions on large language models.\n",
        "\n",
        "Learn more about [Vertex AI Language Models](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.language_models.TextGenerationModel#vertexai_language_models_TextGenerationModel_predict_streaming)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "objective:mlops,stage2,get_started_automl_training"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to use Vertex AI LLM to download pretrained LLM model, make predictions and finetuning the model.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
        "- `Vertex AI LLM`\n",
        "- `Vertex AI Prediction`\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Load a pretrained text generation model.\n",
        "- Make a non-streaming prediction\n",
        "- Load a pretrained text generation model, which supports streaming.\n",
        "- Make a streaming prediction\n",
        "- Load a pretrained chat model.\n",
        "- Do a local interactive chat session.\n",
        "- Do a batch prediction with a text generation model.\n",
        "- Do a batch prediction with a text embedding model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dataset:flowers,icn"
      },
      "source": [
        "### Model\n",
        "\n",
        "The pre-trained models used for this tutorial are from Vertex AI LLM repository. The models used are:\n",
        "\n",
        "- text-bison\n",
        "- chat-bison\n",
        "- text-embedding-gecko"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fb3451ce8e47"
      },
      "source": [
        "### Costs\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "- Vertex AI\n",
        "- Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4b331e2fd155"
      },
      "source": [
        "## Get Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a2c2cb2109a0"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "24166b67eab5"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                                pandas"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ff555b32bab8"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ee775571c2b5"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "92e68cfc3a90"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "46604f70e831"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4f872cd812d0"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6ddefefd807a"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bucket:mbsdk"
      },
      "source": [
        "#### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bucket"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "autoset_bucket"
      },
      "source": [
        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "91c46850b49b"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b7c2a8a91840"
      },
      "source": [
        "#### Set Google Cloud project"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d72078ee9da2"
      },
      "outputs": [],
      "source": [
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "setup_vars"
      },
      "source": [
        "## Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8696271b314f"
      },
      "outputs": [],
      "source": [
        "from vertexai.preview.language_models import TextGenerationModel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YCriCYpOjZv0"
      },
      "source": [
        "## Text generation\n",
        "\n",
        "First, you load from Vertex AI LLM, the text-bison model. Once loaded, you make a prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fkAkKn-ijJR5"
      },
      "outputs": [],
      "source": [
        "model = TextGenerationModel.from_pretrained(\"google/text-bison@001\")\n",
        "\n",
        "print(\n",
        "    model.predict(\n",
        "        \"What is the best recipe for banana bread? Recipe:\",\n",
        "        # \"Brainstorm some ideas combining VR and fitness:\",\n",
        "        # Optional:\n",
        "        # max_output_tokens=128,\n",
        "        # temperature=0,\n",
        "        # top_p=1,\n",
        "        # top_k=5,\n",
        "    )\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vltcd5jb5WGe"
      },
      "source": [
        "## Streaming prediction\n",
        "\n",
        "Next, you load from Vertex AI LLM, a version of the text-bison model that support streaming prediction. Once loaded, you make a prediction. Note, how each iteration that model returns the next successive output in the prediction response."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z72f5GTm1weU"
      },
      "outputs": [],
      "source": [
        "import datetime\n",
        "\n",
        "text_generation_model = TextGenerationModel.from_pretrained(\"text-bison\")\n",
        "\n",
        "print(\"Start: \", datetime.datetime.now())\n",
        "for response in text_generation_model.predict_streaming(\n",
        "    prompt=\"Count to 100\", max_output_tokens=1000\n",
        "):\n",
        "    print(datetime.datetime.now())\n",
        "    print(response)\n",
        "print(\"End: \", datetime.datetime.now())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wl2AZceWjXoy"
      },
      "source": [
        "## Chat\n",
        "\n",
        "Next, you load from Vertex AI LLM, a version of the chat-bison model  Once loaded, you do an interactive chat session."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Lco7I7X_Ch2C"
      },
      "outputs": [],
      "source": [
        "from vertexai.language_models import ChatModel, InputOutputTextPair\n",
        "\n",
        "chat_model2 = ChatModel.from_pretrained(\"google/chat-bison@001\")\n",
        "\n",
        "chat2 = chat_model2.start_chat(\n",
        "    # Optional:\n",
        "    context=\"My name is Ned. You are my personal assistant. My favorite movies are Lord of the Rings and Hobbit.\",\n",
        "    examples=[\n",
        "        InputOutputTextPair(\n",
        "            input_text=\"Who do you work for?\",\n",
        "            output_text=\"I work for Ned.\",\n",
        "        ),\n",
        "        InputOutputTextPair(\n",
        "            input_text=\"What do I like?\",\n",
        "            output_text=\"Ned likes watching movies.\",\n",
        "        ),\n",
        "    ],\n",
        ")\n",
        "\n",
        "print(chat2.send_message(\"Are my favorite movies based on a book series?\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G64BrDoxC-K3"
      },
      "outputs": [],
      "source": [
        "print(chat2.send_message(\"When where these books published?\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tmWzgmkUzgzU"
      },
      "source": [
        "## Text embedding\n",
        "\n",
        "Next, you load from Vertex AI LLM, a version of the text-embedding-gecko model, and then make an embedding request."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Eara9dykkJTA"
      },
      "outputs": [],
      "source": [
        "from vertexai.language_models import TextEmbeddingModel\n",
        "\n",
        "model = TextEmbeddingModel.from_pretrained(\"google/textembedding-gecko@001\")\n",
        "embeddings = model.get_embeddings([\"What is life?\"])\n",
        "for embedding in embeddings:\n",
        "    vector = embedding.values\n",
        "    print(len(vector))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7gaFptssn03c"
      },
      "source": [
        "## Batch prediction\n",
        "\n",
        "### text-bison model\n",
        "\n",
        "Now, you do a batch prediction job with a pretrained text-bison model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xHl1XlTpqIFY"
      },
      "outputs": [],
      "source": [
        "dataset = \"gs://cloud-samples-data/vertex-ai/prediction/llm/test_table.jsonl\"\n",
        "destination_uri_prefix = f\"{BUCKET_URI}/text-bison@001_\"\n",
        "! gsutil cp -r gs://cloud-samples-data/vertex-ai/prediction/llm/text-bison@001_/ {destination_uri_prefix}\n",
        "\n",
        "\n",
        "from vertexai.language_models import TextGenerationModel\n",
        "\n",
        "text_generation_model = TextGenerationModel.from_pretrained(\"text-bison\")\n",
        "batch_job_1 = text_generation_model.batch_predict(\n",
        "    dataset=dataset,\n",
        "    destination_uri_prefix=destination_uri_prefix,\n",
        "    model_parameters={},\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3d8c7fc4a4a7"
      },
      "source": [
        "### text-embedding-gecko model\n",
        "\n",
        "Next, you do a batch prediction job with a pretrained text-embedding-gecko model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DfdSZT4SoPxf"
      },
      "outputs": [],
      "source": [
        "dataset = \"gs://cloud-samples-data/vertex-ai/prediction/llm/embedding_input.jsonl\"\n",
        "\n",
        "destination_uri_prefix = f\"{BUCKET_URI}/textembedding-gecko@001_\"\n",
        "\n",
        "from vertexai.preview.language_models import TextEmbeddingModel\n",
        "\n",
        "text_embedding_model = TextEmbeddingModel.from_pretrained(\"textembedding-gecko@001\")\n",
        "batch_job_2 = text_embedding_model.batch_predict(\n",
        "    dataset=dataset,\n",
        "    destination_uri_prefix=destination_uri_prefix,\n",
        "    # Optional:\n",
        "    model_parameters={},\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dt6Cyr9B0ZOT"
      },
      "source": [
        "## Tuning\n",
        "\n",
        "Now, you fine tune with a pretrained text-bison model, and make a prediction with the fine tuned model. Your input data is in JSONL format and stored in a Cloud Storage location."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WxmUoYfC0aJI"
      },
      "outputs": [],
      "source": [
        "model3 = TextGenerationModel.from_pretrained(\"google/text-bison@001\")\n",
        "\n",
        "model3.list_tuned_model_names()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bb4HhsLc0ha4"
      },
      "outputs": [],
      "source": [
        "# Tuning model\n",
        "! gsutil cp gs://cloud-samples-data/vertex-ai/prediction/llm/q_a_train_with_context.jsonl {BUCKET_URI}/q_a_train_with_context.jsonl\n",
        "\n",
        "tuning_job = model3.tune_model(\n",
        "    training_data=f\"{BUCKET_URI}/q_a_train_with_context.jsonl\",\n",
        "    # Optional:\n",
        "    train_steps=1,\n",
        "    tuning_job_location=\"europe-west4\",\n",
        "    tuned_model_location=\"us-central1\",\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xf7tWnMU03ZY"
      },
      "outputs": [],
      "source": [
        "# Make a prediction with the fine tuned model\n",
        "tuned_model = tuning_job.get_tuned_model()\n",
        "\n",
        "print(tuned_model.predict(\"Tell me some ideas combining VR and fitness:\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7UqoxmAd28QZ"
      },
      "outputs": [],
      "source": [
        "# List tuned model names\n",
        "model3.list_tuned_model_names()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uB7RC4Sd2-vK"
      },
      "outputs": [],
      "source": [
        "# Get tuned model\n",
        "tuned_model4 = model3.get_tuned_model(\n",
        "    tuned_model_name=model3.list_tuned_model_names()[0]\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "asNjLVy47_Ox"
      },
      "outputs": [],
      "source": [
        "# Make a prediction with the fine tuned model\n",
        "print(tuned_model4.predict(\"Brainstorm some ideas combining VR and fitness:\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wai6L3fQowJ_"
      },
      "source": [
        "## Tuning from pandas DataFrame\n",
        "\n",
        "Now, you fine tune with a pretrained text-bison model, and make a prediction with the fine tuned model. Your input data is an in-memory pandas DataFrame."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lUsZKDS3nIcJ"
      },
      "outputs": [],
      "source": [
        "import pandas\n",
        "\n",
        "training_data = pandas.DataFrame(\n",
        "    data=[\n",
        "        {\"input_text\": \"Input 1\", \"output_text\": \"Output 1\"},\n",
        "        {\"input_text\": \"Input 2\", \"output_text\": \"Output 2\"},\n",
        "        {\"input_text\": \"Input 3\", \"output_text\": \"Output 3\"},\n",
        "        {\"input_text\": \"Input 4\", \"output_text\": \"Output 4\"},\n",
        "        {\"input_text\": \"Input 5\", \"output_text\": \"Output 5\"},\n",
        "        {\"input_text\": \"Input 6\", \"output_text\": \"Output 6\"},\n",
        "        {\"input_text\": \"Input 7\", \"output_text\": \"Output 7\"},\n",
        "        {\"input_text\": \"Input 8\", \"output_text\": \"Output 8\"},\n",
        "        {\"input_text\": \"Input 9\", \"output_text\": \"Output 9\"},\n",
        "        {\"input_text\": \"Input 10\", \"output_text\": \"Output 10\"},\n",
        "    ]\n",
        ")\n",
        "\n",
        "training_data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RMReD76mnAOC"
      },
      "outputs": [],
      "source": [
        "# Tune pretrained model\n",
        "model4 = TextGenerationModel.from_pretrained(\"google/text-bison@001\")\n",
        "\n",
        "tuning_job = model4.tune_model(\n",
        "    training_data=training_data,\n",
        "    # Optional:\n",
        "    train_steps=10,\n",
        "    tuning_job_location=\"europe-west4\",\n",
        "    tuned_model_location=\"us-central1\",\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KWCaCpTDnL1h"
      },
      "outputs": [],
      "source": [
        "# Make a prediction with the fine tuned model\n",
        "tuned_model = tuning_job.get_tuned_model()\n",
        "\n",
        "print(tuned_model.predict(\"Tell me some ideas combining VR and fitness:\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b1bb76638fb5"
      },
      "source": [
        "## Cleanup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c41c153770e6"
      },
      "outputs": [],
      "source": [
        "delete_bucket = False\n",
        "\n",
        "# Delete job\n",
        "batch_job_1.delete()\n",
        "batch_job_2.delete()\n",
        "\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "llm_streaming_prediction.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}