language/tuning/tuning_text_bison.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "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": "tvgnzT1CKxrO"
      },
      "source": [
        "# Tuning and deploy a foundation model\n",
        "\n",
        "> **NOTE:** This notebook uses the PaLM generative model, which will reach its [discontinuation date in October 2024](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text#model_versions). Please refer to [this updated notebook](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/gemini_supervised_tuning_qa.ipynb) for a version which uses the latest Gemini model.\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/language/tuning/tuning_text_bison.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/language/tuning/tuning_text_bison.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",
        "  <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/generative-ai/main/language/tuning/tuning_text_bison.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "db178f5f0985"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "|Author(s) | [Erwin Huizenga](https://github.com/erwinh85) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "Creating an LLM requires massive amounts of data, significant computing resources, and specialized skills. On Vertex AI, tuning allows you to customize a foundation model for more specific tasks or knowledge domains.\n",
        "\n",
        "While the prompt design is excellent for quick experimentation, if training data is available, you can achieve higher quality by tuning the model. Tuning a model enables you to customize the model response based on examples of the task you want the model to perform.\n",
        "\n",
        "For more details on tuning have a look at the [official documentation](https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-models)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d975e698c9a4"
      },
      "source": [
        "### Objective\n",
        "\n",
        "This tutorial teaches you how to tune a foundational model on new unseen data and you will use the following Google Cloud products:\n",
        "\n",
        "- Vertex AI Generative AI Studio\n",
        "- Vertex AI Pipelines\n",
        "- Vertex AI Model Registry\n",
        "- Vertex AI Endpoints\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Get training data from BQ and generate a JSONL file\n",
        "- Upload training data\n",
        "- Create a pipeline job\n",
        "- Inspect your model on Vertex AI Model Registry\n",
        "- Get predictions from your tuned model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6CZvFRbIaalF"
      },
      "source": [
        "### Quota\n",
        "**important**: Tuning the text-bison model uses the TPU `v3-8` training resources and the accompanying quotas from your Google Cloud project. Each project has a default quota of eight `v3-8` cores, which allows for one to two concurrent tuning jobs. If you want to run more concurrent jobs you need to request additional quota via the [Quotas page](https://console.cloud.google.com/iam-admin/quotas)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6q2bKpVjaalF"
      },
      "source": [
        "### Costs\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI Generative AI Studio\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
        "and 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": {
        "id": "acBlvcGFaalF"
      },
      "source": [
        "### Install Vertex AI SDK"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BEtR1xyRaalG"
      },
      "outputs": [],
      "source": [
        "%pip install google-cloud-aiplatform google-cloud-bigquery sequence-evaluate sentence-transformers rouge --upgrade --user"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qAMVnZC9aalG"
      },
      "source": [
        "**Colab only:** Uncomment the following cell to restart the kernel or use the restart button. For Vertex AI Workbench you can restart the terminal using the button on top."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MdQC6wcuaalG"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2LlxsZrWaalG"
      },
      "source": [
        "### Authenticating your notebook environment\n",
        "* If you are using **Colab** to run this notebook, uncomment the cell below and continue.\n",
        "* If you are using **Vertex AI Workbench**, check out the setup instructions [here](https://github.com/GoogleCloudPlatform/generative-ai/tree/main/setup-env)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "oh-QANoIaalG"
      },
      "outputs": [],
      "source": [
        "from google.colab import auth\n",
        "\n",
        "auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qW8qtGsmaalG"
      },
      "source": [
        "### BigQuery IAM\n",
        "Now you need to add permissions to the service account:\n",
        "- Go to the [IAM page](https://console.cloud.google.com/iam-admin/) in the console\n",
        "- Look for the default compute service account. It should look something like this: `<project-number>-compute@developer.gserviceaccount.com`\n",
        "- Assign the default compute service account with `bigquery.user`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZmhnHOjlaalH"
      },
      "source": [
        "### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`. Otherwise, check the support page: Locate the [project ID](https://support.google.com/googleapi/answer/7014113). Please update `PROJECT_ID` below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j8nXkkYxaalH"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"<your_project_id>\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PrsmSjICaalH"
      },
      "source": [
        "### Create a bucket\n",
        "Now you have to create a bucket that we will use to store our tuning data. To avoid name collisions between users on resources created, you generate a UUID for each instance session and append it to the name of the resources you create in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "LiKRZOgqaalH"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import string\n",
        "\n",
        "\n",
        "# Generate a uuid of a specified length(default=8)\n",
        "def generate_uuid(length: int = 8) -> str:\n",
        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))\n",
        "\n",
        "\n",
        "UUID = generate_uuid()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-D28-KrtaalH"
      },
      "source": [
        "Choose a bucket name and update the `BUCKET_NAME` parameter."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "pxRSNVCYaalH"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = \"<your_bucket_uri>\"  # @param {type:\"string\"}\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "ZpjqMRc-aalH"
      },
      "outputs": [],
      "source": [
        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"<your-bucket-name>\":\n",
        "    BUCKET_NAME = \"vertex-\" + UUID\n",
        "    BUCKET_URI = f\"gs://{BUCKET_NAME}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WtJg8ILPaalH"
      },
      "source": [
        "Only 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": "NSRiXkavaalH"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jNL0oqUJaalH"
      },
      "source": [
        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "leJFL5oIaalH"
      },
      "outputs": [],
      "source": [
        "! gsutil ls -al $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XoEqT2Y4DJmf"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0E7pfl6sjzh_"
      },
      "source": [
        "**Colab only**: Run the following cell to initialize the Vertex AI SDK. For Vertex AI Workbench, you don't need to run this."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "JXlNPFmGjzh_"
      },
      "outputs": [],
      "source": [
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=REGION)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "pRUOFELefqf1"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform, bigquery\n",
        "import pandas as pd\n",
        "from sklearn.model_selection import train_test_split\n",
        "from vertexai.language_models import TextGenerationModel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WdtNETYxaalH"
      },
      "source": [
        "## Tune your Model\n",
        "\n",
        "Now it's time for you to create a tuning job. Tune a foundation model by creating a pipeline job using Generative AI Studio, cURL, or the Python SDK. In this notebook, we will be using the Python SDK. You will be using a Q&A with a context dataset in JSON format.\n",
        "\n",
        "### Training Data\n",
        "💾 Your model tuning dataset must be in a JSONL format where each line contains a single training example. You must make sure that you include instructions.\n",
        "\n",
        "You will use the StackOverflow data on BigQuery Public Datasets, limiting to questions with the `python` tag, and accepted answers for answers since 2020-01-01."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Puc3jl8QaalI"
      },
      "source": [
        "First create a helper function to let you easily query BigQuery and return the results as a Pandas DataFrame."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "Eg60aUgvaalI"
      },
      "outputs": [],
      "source": [
        "def run_bq_query(sql: str) -> str | pd.DataFrame:\n",
        "    \"\"\"\n",
        "    Run a BigQuery query and return the job ID or result as a DataFrame\n",
        "    Args:\n",
        "        sql: SQL query, as a string, to execute in BigQuery\n",
        "    Returns:\n",
        "        df: DataFrame of results from query,  or error, if any\n",
        "    \"\"\"\n",
        "\n",
        "    bq_client = bigquery.Client(project=PROJECT_ID)\n",
        "\n",
        "    # Try dry run before executing query to catch any errors\n",
        "    job_config = bigquery.QueryJobConfig(dry_run=True, use_query_cache=False)\n",
        "    bq_client.query(sql, job_config=job_config)\n",
        "\n",
        "    # If dry run succeeds without errors, proceed to run query\n",
        "    job_config = bigquery.QueryJobConfig()\n",
        "    client_result = bq_client.query(sql, job_config=job_config)\n",
        "\n",
        "    job_id = client_result.job_id\n",
        "\n",
        "    # Wait for query/job to finish running. then get & return DataFrame\n",
        "    df = client_result.result().to_arrow().to_pandas()\n",
        "    print(f\"Finished job_id: {job_id}\")\n",
        "\n",
        "    return df"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1BydoFfTaalI"
      },
      "source": [
        "Next define the query."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9VTaovLtaalI"
      },
      "outputs": [],
      "source": [
        "df = run_bq_query(\n",
        "    \"\"\"SELECT\n",
        "    CONCAT(q.title, q.body) as input_text,\n",
        "    a.body AS output_text\n",
        "FROM\n",
        "    `bigquery-public-data.stackoverflow.posts_questions` q\n",
        "JOIN\n",
        "    `bigquery-public-data.stackoverflow.posts_answers` a\n",
        "ON\n",
        "    q.accepted_answer_id = a.id\n",
        "WHERE\n",
        "    q.accepted_answer_id IS NOT NULL AND\n",
        "    REGEXP_CONTAINS(q.tags, \"python\") AND\n",
        "    a.creation_date >= \"2020-01-01\"\n",
        "LIMIT\n",
        "    10000\n",
        "\"\"\"\n",
        ")\n",
        "\n",
        "df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qYUg8cBbaalJ"
      },
      "source": [
        "There should be 10k questions and answers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6FqbVHoeaalJ"
      },
      "outputs": [],
      "source": [
        "print(len(df))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OftmoPZ6aalJ"
      },
      "source": [
        "Lets split the data into training and evaluation. For Extractive Q&A tasks we advise 100+ training examples. In this case you will use 800."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aXqBwSwaaalJ"
      },
      "outputs": [],
      "source": [
        "# split is set to 80/20\n",
        "train, evaluation = train_test_split(df, test_size=0.2)\n",
        "evaluation = evaluation.sample(n=250, random_state=1)\n",
        "print(len(train))\n",
        "print(len(evaluation))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nf-q8TpnaalJ"
      },
      "source": [
        "For tuning, the training data first needs to be converted into a JSONL format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FqRbOwzEaalJ"
      },
      "outputs": [],
      "source": [
        "tune_jsonl = train.to_json(orient=\"records\", lines=True)\n",
        "\n",
        "print(f\"Length: {len(tune_jsonl)}\")\n",
        "print(tune_jsonl[0:100])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r04PWISCaalJ"
      },
      "source": [
        "Next, you can write it to a local JSONL before transferring it to Google Cloud Storage (GCS)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "vXVV9c0HaalJ"
      },
      "outputs": [],
      "source": [
        "training_data_filename = \"tune_data_stack_overflow_python_qa.jsonl\"\n",
        "\n",
        "with open(training_data_filename, \"w\") as f:\n",
        "    f.write(tune_jsonl)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UHS1lDIrrfQH"
      },
      "outputs": [],
      "source": [
        "tune_jsonl = evaluation.to_json(orient=\"records\", lines=True)\n",
        "\n",
        "print(f\"Length: {len(tune_jsonl)}\")\n",
        "print(tune_jsonl[0:100])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "4eULfrv2rTjO"
      },
      "outputs": [],
      "source": [
        "evaluation_data_filename = \"tune_eval_data_stack_overflow_python_qa.jsonl\"\n",
        "\n",
        "with open(evaluation_data_filename, \"w\") as f:\n",
        "    f.write(tune_jsonl)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FV8Wxz7JaalN"
      },
      "source": [
        "You can then export the local file to GCS, so that it can be used by Vertex AI for the tuning job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vDDLHac5aalN"
      },
      "outputs": [],
      "source": [
        "! gsutil cp $training_data_filename $evaluation_data_filename $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ff68wmzoaalN"
      },
      "source": [
        "You can check to make sure that the file successfully transferred to your Google Cloud Storage bucket:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2-DnKpYlaalN"
      },
      "outputs": [],
      "source": [
        "! gsutil ls -al $BUCKET_URI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "8wE9P7OFaalN"
      },
      "outputs": [],
      "source": [
        "TRAINING_DATA_URI = f\"{BUCKET_URI}/{training_data_filename}\"\n",
        "EVALUATION_DATA_URI = f\"{BUCKET_URI}/{evaluation_data_filename}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-mW7K57BaalN"
      },
      "source": [
        "### Model Tuning\n",
        "Now it's time to start to tune a model. You will use the Vertex AI SDK to submit our tuning job.\n",
        "\n",
        "#### Recommended Tuning Configurations\n",
        "✅ Here are some recommended configurations for tuning a foundation model based on the task, in this example Q&A. You can find more in the [documentation](https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-models).\n",
        "\n",
        "Extractive QA:\n",
        "- Make sure that your train dataset size is 100+\n",
        "- Training steps [100-500]. You can try more than one value to get the best performance on a particular dataset (e.g. 100, 200, 500)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-vP_jeATTnbK"
      },
      "outputs": [],
      "source": [
        "# create tensorboard\n",
        "display_name = \"Adapter tuning - \"\n",
        "\n",
        "tensorboard = aiplatform.Tensorboard.create(\n",
        "    display_name=display_name,\n",
        "    project=PROJECT_ID,\n",
        "    location=REGION,\n",
        ")\n",
        "\n",
        "print(tensorboard.display_name)\n",
        "print(tensorboard.resource_name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0CFVBoFu5Cnx"
      },
      "outputs": [],
      "source": [
        "# Get tensorboard_id thats used in the pipeline\n",
        "tensorboard_id = tensorboard.resource_name.split(\"tensorboards/\")[-1]\n",
        "print(tensorboard_id)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "id": "26HRfld3aalN"
      },
      "outputs": [],
      "source": [
        "MODEL_NAME = f\"genai-workshop-tuned-model-{UUID}\"\n",
        "TRAINING_STEPS = 100"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {
        "id": "DvG1Rp3-iAG_"
      },
      "outputs": [],
      "source": [
        "pipeline_arguments = {\n",
        "    \"model_display_name\": MODEL_NAME,\n",
        "    \"location\": REGION,\n",
        "    \"large_model_reference\": \"text-bison\",\n",
        "    \"project\": PROJECT_ID,\n",
        "    \"train_steps\": TRAINING_STEPS,\n",
        "    \"dataset_uri\": TRAINING_DATA_URI,\n",
        "    \"evaluation_interval\": 20,\n",
        "    \"evaluation_data_uri\": EVALUATION_DATA_URI,\n",
        "    \"tensorboard_resource_id\": tensorboard_id,\n",
        "}\n",
        "\n",
        "pipeline_root = f\"{BUCKET_URI}/{MODEL_NAME}\"\n",
        "template_path = \"https://us-kfp.pkg.dev/ml-pipeline/large-language-model-pipelines/tune-large-model/v2.0.0\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 55,
      "metadata": {
        "id": "on4baTh5aalN"
      },
      "outputs": [],
      "source": [
        "# Function that starts the tuning job\n",
        "\n",
        "\n",
        "def tuned_model(\n",
        "    project_id: str,\n",
        "    location: str,\n",
        "    template_path: str,\n",
        "    model_display_name: str,\n",
        "    pipeline_arguments: str,\n",
        "):\n",
        "    \"\"\"Prompt-tune a new model, based on a prompt-response data.\n",
        "\n",
        "    \"training_data\" can be either the GCS URI of a file formatted in JSONL format\n",
        "    (for example: training_data=f'gs://{bucket}/{filename}.jsonl'), or a pandas\n",
        "    DataFrame. Each training example should be JSONL record with two keys, for\n",
        "    example:\n",
        "      {\n",
        "        \"input_text\": <input prompt>,\n",
        "        \"output_text\": <associated output>\n",
        "      },\n",
        "\n",
        "    Args:\n",
        "      project_id: Google Cloud Project ID, used to initialize aiplatform\n",
        "      location: Google Cloud Region, used to initialize aiplatform\n",
        "      template_path: path to the template\n",
        "      model_display_name: Name for your model.\n",
        "      pipeline_arguments: arguments used during pipeline runtime\n",
        "    \"\"\"\n",
        "\n",
        "    aiplatform.init(project=project_id, location=location)\n",
        "\n",
        "    from google.cloud.aiplatform import PipelineJob\n",
        "\n",
        "    job = PipelineJob(\n",
        "        template_path=template_path,\n",
        "        display_name=model_display_name,\n",
        "        parameter_values=pipeline_arguments,\n",
        "        location=REGION,\n",
        "        pipeline_root=pipeline_root,\n",
        "        enable_caching=True,\n",
        "    )\n",
        "\n",
        "    return job"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o0XNL9ojaalN"
      },
      "source": [
        "Next, it's time to start your tuning job.\n",
        "\n",
        "**Disclaimer:** tuning and deploying a model takes time."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m0m86z20zFgl"
      },
      "outputs": [],
      "source": [
        "job = tuned_model(PROJECT_ID, REGION, template_path, MODEL_NAME, pipeline_arguments)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NPHoXo8UIhlz"
      },
      "outputs": [],
      "source": [
        "job.submit()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PRCkdxXvaalO"
      },
      "source": [
        "Following the link above, you can view your pipeline run. As you can see in the screenshot below, it will execute the following steps:\n",
        "\n",
        "- Validation\n",
        "- Export managed dataset\n",
        "- Convert JSONL to TFRecord\n",
        "- Large language model tuning\n",
        "- Upload LLM Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jq9rL5o32EXs"
      },
      "source": [
        "`job.state` lets you check the state of your pipeline."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u8l8o3z30WO4"
      },
      "outputs": [],
      "source": [
        "job.state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O6JC8XplaalO"
      },
      "source": [
        "## View your tuned foundational model on Vertex AI Model registry\n",
        "When your tuning job is finished, your model will be available on Vertex AI Model Registry. The following Python SDK sample shows you how to list tuned models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GPWX0ITCaalO"
      },
      "outputs": [],
      "source": [
        "def list_tuned_models(project_id, location):\n",
        "    aiplatform.init(project=project_id, location=location)\n",
        "    model = TextGenerationModel.from_pretrained(\"text-bison\")\n",
        "    tuned_model_names = model.list_tuned_model_names()\n",
        "    print(tuned_model_names)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bAIwCGYJaalO"
      },
      "outputs": [],
      "source": [
        "list_tuned_models(PROJECT_ID, REGION)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ZriyF0V-aalO"
      },
      "source": [
        "You can also use the Google Cloud Console UI to view all of your model in [Vertex AI Model Registry](https://console.cloud.google.com/vertex-ai/models). Below you can see an example of a tuned foundational model available on Vertex AI Model Registry."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cFftY6-EaalO"
      },
      "source": [
        "## Use your tuned model to get predictions\n",
        "Now it's time to get predictions. First you need to get the latest tuned model from the Vertex AI Model registry."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vU-K3EIkaalO"
      },
      "outputs": [],
      "source": [
        "def fetch_model(project_id, location):\n",
        "    aiplatform.init(project=project_id, location=location)\n",
        "    model = TextGenerationModel.from_pretrained(\"text-bison\")\n",
        "    list_tuned_models = model.list_tuned_model_names()\n",
        "    tuned_model = list_tuned_models[0]\n",
        "\n",
        "    return tuned_model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j66dr12taalO"
      },
      "outputs": [],
      "source": [
        "deployed_model = fetch_model(PROJECT_ID, REGION)\n",
        "deployed_model = TextGenerationModel.get_tuned_model(deployed_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xDOueoptaalO"
      },
      "source": [
        "Now you can start send a prompt to the API. Feel free to update the following prompt."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2ERbfPJPaalO"
      },
      "outputs": [],
      "source": [
        "PROMPT = \"\"\"\n",
        "How can I store my TensorFlow checkpoint on Google Cloud Storage?\n",
        "\n",
        "Python example:\n",
        "\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "trzon4EyaalO"
      },
      "outputs": [],
      "source": [
        "print(deployed_model.predict(PROMPT))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qtYr_KNPaalO"
      },
      "source": [
        "## Evaluation\n",
        "It's essential to evaluate your model to understand its performance. Evaluation can be done in an automated way using evaluation metrics like F1 or Rouge. You can also leverage human evaluation methods. Human evaluation methods involve asking humans to rate the quality of the LLM's answers. This can be done through crowdsourcing or by having experts evaluate the responses. Some standard human evaluation metrics include fluency, coherence, relevance, and informativeness. Often you want to choose a mix of evaluation metrics to get a good understanding of your model performance. Below you will find an example of how you can do the evaluation.\n",
        "\n",
        "In this example you will be using [sequence-evaluate](https://pypi.org/project/sequence-evaluate/) to evaluation the tuned model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9856CuicaalO"
      },
      "outputs": [],
      "source": [
        "from seq_eval import SeqEval\n",
        "\n",
        "evaluator = SeqEval()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AS10ybdraalO"
      },
      "source": [
        "Earlier in the notebook, you created a train and eval dataset. Now it's time to take some of the eval data. You will use the questions to get a response from our tuned model, and the answers we will use as a reference:\n",
        "\n",
        "- **Candidates**: Answers generated by the tuned model.\n",
        "- **References**: Original answers that we will use to compare."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LKMmIH0XaalO"
      },
      "outputs": [],
      "source": [
        "evaluation = evaluation.head(10)  # you can change the number of rows you want to use\n",
        "evaluation_question = evaluation[\"input_text\"]\n",
        "evaluation_answer = evaluation[\"output_text\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jx-g2molaalP"
      },
      "source": [
        "Now you can go ahead and generate candidates using the tuned model based on the questions you took from the eval dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e5DqVXvEaalP"
      },
      "outputs": [],
      "source": [
        "candidates = []\n",
        "\n",
        "for i in evaluation_question:\n",
        "    response = deployed_model.predict(i)\n",
        "    candidates.append(response.text)\n",
        "\n",
        "len(candidates)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oftLTb0maalP"
      },
      "source": [
        "You will also have to create a list of our references. These will you use to evaluate the model's performance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y7zN70CJaalP"
      },
      "outputs": [],
      "source": [
        "references = evaluation_answer.tolist()\n",
        "\n",
        "len(references)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UwKcyIDdjziD"
      },
      "source": [
        "Next you will generate the evaluation metrics. `evaluator.evaluate` will return a few eval metrics. Some of the important ones are:\n",
        "- [Blue](https://en.wikipedia.org/wiki/BLEU): The BLEU evaluation metric is a measure of the similarity between a machine-generated text and a human-written reference text.\n",
        "- [Rouge](https://en.wikipedia.org/wiki/ROUGE_(metric)): The ROUGE evaluation metric is a measure of the overlap between a machine-generated text and a human-written reference text."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B828sNxUaalP"
      },
      "outputs": [],
      "source": [
        "scores = evaluator.evaluate(candidates, references, verbose=False)\n",
        "print(scores)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "tuning_text_bison.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}