gemini/tuning/sft_gemini_qa.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 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": "ojoyvz6mH1Hv"
      },
      "source": [
        "# Supervised fine-tuning with Gemini 2.0 Flash for Q&A using the Google Gen AI SDK\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/gemini/tuning/sft_gemini_qa.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" 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%2Fgenerative-ai%2Fmain%2Fgemini%2Ftuning%2Fsft_gemini_qa.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" 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/generative-ai/main/gemini/tuning/sft_gemini_qa.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/generative-ai/blob/main/gemini/tuning/sft_gemini_qa.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://goo.gle/4jjzEpM\">\n",
        "      <img width=\"32px\" src=\"https://cdn.qwiklabs.com/assets/gcp_cloud-e3a77215f0b8bfa9b3f611c0d2208c7e8708ed31.svg\" alt=\"Google Cloud logo\"><br> Open in  Cloud Skills Boost\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/sft_gemini_qa.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/sft_gemini_qa.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/sft_gemini_qa.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/sft_gemini_qa.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/sft_gemini_qa.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>            "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84f0f73a0f76"
      },
      "source": [
        "| Author(s) |\n",
        "| --- |\n",
        "| [Erwin Huizenga](https://github.com/erwinh85) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "## Overview\n",
        "\n",
        "**Gemini** is a family of generative AI models developed by Google DeepMind designed for multimodal use cases. The Gemini API gives you access to the various Gemini models, such as Gemini 2.0 and Gemini 2.0.\n",
        "This notebook demonstrates fine-tuning the Gemini 2.0 using the Vertex AI Supervised Tuning feature. Supervised Tuning allows you to use your own labeled training data to further refine the base model's capabilities toward your specific tasks.\n",
        "Supervised Tuning uses labeled examples to tune a model. Each example demonstrates the output you want from your text model during inference.\n",
        "First, ensure your training data is of high quality, well-labeled, and directly relevant to the target task. This is crucial as low-quality data can adversely affect the performance and introduce bias in the fine-tuned model.\n",
        "Training: Experiment with different configurations to optimize the model's performance on the target task.\n",
        "Evaluation:\n",
        "Metric: Choose appropriate evaluation metrics that accurately reflect the success of the fine-tuned model for your specific task\n",
        "Evaluation Set: Use a separate set of data to evaluate the model's performance\n",
        "\n",
        "\n",
        "Refer to public [documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/models/gemini-supervised-tuning) for more details.\n",
        "\n",
        "<hr/>\n",
        "\n",
        "Before running this notebook, ensure you have:\n",
        "\n",
        "- A Google Cloud project: Provide your project ID in the `PROJECT_ID` variable.\n",
        "\n",
        "- Authenticated your Colab environment: Run the authentication code block at the beginning.\n",
        "\n",
        "- Prepared training data (Test with your own data or use the one in the notebook): Data should be formatted in JSONL with prompts and corresponding completions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f7SS5pzuIA-1"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI\n",
        "pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage.\n",
        "\n",
        "To estimate the cost of token please have a look at this [notebook](https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/tuning/vertexai_supervised_tuning_token_count_and_cost_estimation.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install the Google Gen AI SDK and other required packages\n",
        "\n",
        "The new Google Gen AI SDK provides a unified interface to Gemini through both the Gemini Developer API and the Gemini API on Vertex AI. With a few exceptions, code that runs on one platform will run on both. This means that you can prototype an application using the Developer API and then migrate the application to Vertex AI without rewriting your code.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet google-cloud-aiplatform google-genai"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n",
        "\n",
        "The restart might take a minute or longer. After it's restarted, continue to the next step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "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": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set the Google Cloud project information and initialize the Google Gen AI SDK\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).\n",
        "\n",
        "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": {
        "cellView": "code",
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "# Use the environment variable if the user doesn't provide Project ID.\n",
        "import os\n",
        "\n",
        "from google import genai\n",
        "from google.genai import types\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\", isTemplate: true}\n",
        "if PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5303c05f7aa6"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6fc324893334"
      },
      "outputs": [],
      "source": [
        "from collections import Counter\n",
        "import json\n",
        "import random\n",
        "\n",
        "# Vertex AI SDK\n",
        "from google.cloud import aiplatform\n",
        "from google.cloud.aiplatform.metadata import context\n",
        "from google.cloud.aiplatform.metadata import utils as metadata_utils\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import plotly.graph_objects as go\n",
        "from plotly.subplots import make_subplots\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6bBZa2I-c-x8"
      },
      "source": [
        "### Data\n",
        "\n",
        "#### SQuAD dataset\n",
        "Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span, from the corresponding reading passage, or the question might be unanswerable.\n",
        "\n",
        "You can find more information on the SQuAD [github page](https://rajpurkar.github.io/SQuAD-explorer/)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MllVfoJfOkDQ"
      },
      "source": [
        "```bibtex\n",
        "@inproceedings{rajpurkar-etal-2016-squad,\n",
        "    title = \"{SQ}u{AD}: 100,000+ Questions for Machine Comprehension of Text\",\n",
        "    author = \"Rajpurkar, Pranav  and\n",
        "      Zhang, Jian  and\n",
        "      Lopyrev, Konstantin  and\n",
        "      Liang, Percy\",\n",
        "    editor = \"Su, Jian  and\n",
        "      Duh, Kevin  and\n",
        "      Carreras, Xavier\",\n",
        "    booktitle = \"Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing\",\n",
        "    month = nov,\n",
        "    year = \"2016\",\n",
        "    address = \"Austin, Texas\",\n",
        "    publisher = \"Association for Computational Linguistics\",\n",
        "    url = \"https://aclanthology.org/D16-1264\",\n",
        "    doi = \"10.18653/v1/D16-1264\",\n",
        "    pages = \"2383--2392\",\n",
        "    eprint={1606.05250},\n",
        "    archivePrefix={arXiv},\n",
        "    primaryClass={cs.CL},\n",
        "}\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KhebDJjRKePL"
      },
      "source": [
        "First update the `BUCKET_NAME` parameter below. You can either use an existing bucket or create a new one."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lit30Cktbfvo"
      },
      "outputs": [],
      "source": [
        "# Provide a bucket name\n",
        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type:\"string\"}\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
        "print(BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ed-G-9cyKmPY"
      },
      "source": [
        "Only run the code below if you want to create a new Google Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0UJ8S9YFA1pZ"
      },
      "outputs": [],
      "source": [
        "# ! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "izjwF63tLLEq"
      },
      "source": [
        "Next you will copy the data into your bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wjvcxx_sA3xP"
      },
      "outputs": [],
      "source": [
        "!gsutil cp gs://github-repo/generative-ai/gemini/tuning/qa/squad_test.csv .\n",
        "!gsutil cp gs://github-repo/generative-ai/gemini/tuning/qa/squad_train.csv .\n",
        "!gsutil cp gs://github-repo/generative-ai/gemini/tuning/qa/squad_validation.csv ."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3F10LuZeL3kt"
      },
      "source": [
        "### Baseline\n",
        "\n",
        "Next you will prepare some data that you will use to establish a baseline.  This means evaluating the out of the box default model on a representative sample of your dataset before any fine-tuning. A baseline allows you to quantify the improvements achieved through fine-tuning."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LkOmXpegA8CW"
      },
      "outputs": [],
      "source": [
        "test_df = pd.read_csv(\"squad_test.csv\")\n",
        "test_df.head(2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Px6HUVKn8ErZ"
      },
      "source": [
        "First you need to prepare some data to evaluate the out of the box model and set a baseline. In this case, we will lower the text and remove extra whitespace, but preserve newlines."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6HXBgGFZB3hx"
      },
      "outputs": [],
      "source": [
        "row_dataset = random.randint(0, 100)  # lets take a random example from the dataset\n",
        "\n",
        "\n",
        "def normalize_answer(s):\n",
        "    \"\"\"Lower text and remove extra whitespace, but preserve newlines.\"\"\"\n",
        "\n",
        "    def white_space_fix(text):\n",
        "        return \" \".join(text.split())  # Splits by any whitespace, including \\n\n",
        "\n",
        "    def lower(text):\n",
        "        return text.lower()\n",
        "\n",
        "    return white_space_fix(lower(s))\n",
        "\n",
        "\n",
        "test_df[\"answers\"] = test_df[\"answers\"].apply(normalize_answer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mrLxcVVcMsNO"
      },
      "source": [
        "You want to make sure that you test data looks the same as your training data to prevent training / serving skew. We will add a system instruction to the dataset:\n",
        "\n",
        "- `SystemInstruct`: System instructions are a set of instructions that the model processes before it processes prompts. We recommend that you use system instructions to tell the model how you want it to behave and respond to prompts.\n",
        "- You will also combine the `context` and `question`. Both will be send to the model to generate a response."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jdKUBaA4ZOFg"
      },
      "outputs": [],
      "source": [
        "few_shot_examples = test_df.sample(3)\n",
        "# Get the indices of the sampled rows\n",
        "dropped_indices = few_shot_examples.index\n",
        "# Remove the sampled rows from the original DataFrame\n",
        "test_df = test_df.drop(dropped_indices)\n",
        "\n",
        "few_shot_prompt = \"\"\n",
        "for _, row in few_shot_examples.iterrows():\n",
        "    few_shot_prompt += (\n",
        "        f\"Context: {row.context}\\nQuestion: {row.question}\\nAnswer: {row.answers}\\n\\n\"\n",
        "    )\n",
        "\n",
        "print(few_shot_prompt)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-uJSrdU3aDi9"
      },
      "outputs": [],
      "source": [
        "# Incorporate few-shot examples into the system instruction\n",
        "systemInstruct = f\"\"\"Answer the question with a concise extract from the given context. Do not add any additional information, capital letters (only for names) or a punctuation mark in the end.\\n\\n\n",
        "Here are some examples: \\n\\n\n",
        "{few_shot_prompt}\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N_u3VzUMsyqj"
      },
      "outputs": [],
      "source": [
        "# combine the systeminstruct + context + question into one column. This will be your input prompt.\n",
        "test_df[\"systemInstruct\"] = systemInstruct\n",
        "\n",
        "test_df[\"input_question\"] = (\n",
        "    \"\\n\\n **Below the question with context that you need to answer**\"\n",
        "    + \"\\n Context: \"\n",
        "    + test_df[\"context\"]\n",
        "    + \"\\n Question: \"\n",
        "    + test_df[\"question\"]\n",
        ")\n",
        "\n",
        "test_systemInstruct = test_df[\"systemInstruct\"].iloc[row_dataset]\n",
        "print(test_systemInstruct)\n",
        "test_question = test_df[\"input_question\"].iloc[row_dataset]\n",
        "print(test_question)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FSxYYqMGWrmj"
      },
      "source": [
        "Next, set the model that you will use. In this example you will use `\"gemini-2.0-flash-001\"`, a multimodal model that is designed for high-volume, cost-effective applications, and which delivers speed and efficiency to build fast, lower-cost applications that don't compromise on quality.\n",
        "\n",
        "For the latest Gemini models and versions, please have a look at our [documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t-5X4goiqqBQ"
      },
      "outputs": [],
      "source": [
        "base_model = \"gemini-2.0-flash-001\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cfZM3RQ3OJ_t"
      },
      "outputs": [],
      "source": [
        "y_true = test_df[\"answers\"].values\n",
        "y_pred_question = test_df[\"input_question\"].values\n",
        "\n",
        "# Check two pairs of question and answer.\n",
        "for i in range(2):  # Loop through the first two indices\n",
        "    print(f\"Pair {i+1}:\")\n",
        "    print(f\"  True Answer: {y_true[i]}\")\n",
        "    print(f\"  Predicted Question: {y_pred_question[i]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wyscyIenW4WZ"
      },
      "source": [
        "Next lets take a question and get a prediction from Gemini that we can compare to the actual answer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "udTxzY8mpGYf"
      },
      "outputs": [],
      "source": [
        "def get_predictions(question: str, model_version: str) -> str:\n",
        "\n",
        "    prompt = question\n",
        "    base_model = model_version\n",
        "\n",
        "    response = client.models.generate_content(\n",
        "        model=base_model,\n",
        "        contents=prompt,\n",
        "        config={\n",
        "            \"system_instruction\": systemInstruct,\n",
        "            \"temperature\": 0.3,\n",
        "        },\n",
        "    )\n",
        "\n",
        "    return response.text"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PFvwmGll3MIv"
      },
      "outputs": [],
      "source": [
        "test_answer = test_df[\"answers\"].iloc[row_dataset]\n",
        "response = get_predictions(test_question, base_model)\n",
        "\n",
        "print(f\"Gemini response: {response}\")\n",
        "print(f\"Actual answer: {test_answer}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OGRJTHKrdujw"
      },
      "source": [
        "Sometimes you might get an answer from Gemini is more lengthy. However, answers in the SQuAD dataset are typically concise and clear.\n",
        "\n",
        "Fine-tuning is a great way to control the type of output your use case requires. In this instance, you would want the model to provide short, clear answers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oQEEaRZZFgCD"
      },
      "outputs": [],
      "source": [
        "# Apply the get_prediction() function to the 'question_column'\n",
        "test_df[\"predicted_answer\"] = test_df[\"input_question\"].apply(get_predictions)\n",
        "test_df.head(2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T9ZnoFRm-uUZ"
      },
      "source": [
        "You also need to make sure that the predicted answer is in the same format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B-kcXWYWHM3n"
      },
      "outputs": [],
      "source": [
        "test_df[\"predicted_answer\"] = test_df[\"predicted_answer\"].apply(normalize_answer)\n",
        "test_df.head(4)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hvi7m8pKE8WB"
      },
      "source": [
        "Next, let's establish a baseline using evaluation metrics.\n",
        "\n",
        "Evaluating the performance of a Question Answering (QA) system requires specific metrics. Two commonly used metrics are Exact Match (EM) and F1 score.\n",
        "\n",
        "EM is a strict measure that only considers an answer correct if it perfectly matches the ground truth, even down to the punctuation. It's a binary metric - either 1 for a perfect match or 0 otherwise. This makes it sensitive to minor variations in phrasing.\n",
        "\n",
        "F1 score is more flexible. It considers the overlap between the predicted answer and the true answer in terms of individual words or tokens. It calculates the harmonic mean of precision (proportion of correctly predicted words out of all predicted words) and recall (proportion of correctly predicted words out of all true answer words). This allows for partial credit and is less sensitive to minor wording differences.\n",
        "\n",
        "In practice, EM is useful when exact wording is crucial, while F1 is more suitable when evaluating the overall understanding and semantic accuracy of the QA system. Often, both metrics are used together to provide a comprehensive evaluation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6AI57wAgIB6A"
      },
      "outputs": [],
      "source": [
        "def f1_score_squad(prediction, ground_truth):\n",
        "    prediction_tokens = normalize_answer(prediction).split()\n",
        "    ground_truth_tokens = normalize_answer(ground_truth).split()\n",
        "    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)\n",
        "    num_same = sum(common.values())\n",
        "    if num_same == 0:\n",
        "        return 0\n",
        "    precision = 1.0 * num_same / len(prediction_tokens)\n",
        "    recall = 1.0 * num_same / len(ground_truth_tokens)\n",
        "    f1 = (2 * precision * recall) / (precision + recall)\n",
        "    return f1\n",
        "\n",
        "\n",
        "def exact_match_score(prediction, ground_truth):\n",
        "    return normalize_answer(prediction) == normalize_answer(ground_truth)\n",
        "\n",
        "\n",
        "def calculate_em_and_f1(y_true, y_pred):\n",
        "    \"\"\"Calculates EM and F1 scores for DataFrame columns.\"\"\"\n",
        "\n",
        "    # Ensure inputs are Series\n",
        "    if not isinstance(y_true, pd.Series):\n",
        "        y_true = pd.Series(y_true)\n",
        "    if not isinstance(y_pred, pd.Series):\n",
        "        y_pred = pd.Series(y_pred)\n",
        "\n",
        "    em = np.mean(y_true.combine(y_pred, exact_match_score))\n",
        "    f1 = np.mean(y_true.combine(y_pred, f1_score_squad))\n",
        "\n",
        "    # # Print non-matching pairs (using index for clarity)\n",
        "    # for i, (t, p) in enumerate(zip(y_true, y_pred)):\n",
        "    #     if not exact_match_score(p, t):\n",
        "    #         print(f\"No EM Match at index {i}:\\nTrue: {t}\\nPred: {p}\\n\")\n",
        "\n",
        "    return em, f1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iuHUevIX_fMn"
      },
      "outputs": [],
      "source": [
        "em, f1 = calculate_em_and_f1(test_df[\"answers\"], test_df[\"predicted_answer\"])\n",
        "print(f\"EM score: {em}\")\n",
        "print(f\"F1 score: {f1}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "22DfexbNfUHm"
      },
      "source": [
        "### Prepare the data for fine-tuning\n",
        "\n",
        "To optimize the supervised fine-tuning process for a foundation model, ensure your dataset includes examples that reflect the desired task. Each record in the dataset pairs an input text (or prompt) with its corresponding expected output. This supervised tuning approach uses the dataset to effectively teach the model the specific behavior or task you need it to perform, by providing numerous illustrative examples.\n",
        "\n",
        "The size of your dataset will vary depending on the complexity of the task, but as a general rule, the more examples you include, the better the model's performance. For fine-tuning Gemini on Vertex AI, the minimum number of examples are 100.\n",
        "\n",
        "#### Dataset Format\n",
        "Your training data should be structured in a JSONL file and stored at a Google Cloud Storage (GCS) URI.  Each line in the JSONL file must adhere to the following schema:\n",
        "\n",
        "A `contents` array containing objects that define:\n",
        "- A `role` (\"user\" for user input or \"model\" for model output)\n",
        "- `parts` containing the input data.\n",
        "\n",
        "```json\n",
        "{\n",
        "   \"contents\":[\n",
        "      {\n",
        "         \"role\":\"user\",  # This indicate input content\n",
        "         \"parts\":[\n",
        "            {\n",
        "               \"text\":\"How are you?\"\n",
        "            }\n",
        "         ]\n",
        "      },\n",
        "      {\n",
        "         \"role\":\"model\", # This indicate target content\n",
        "         \"parts\":[ # text only\n",
        "            {\n",
        "               \"text\":\"I am good, thank you!\"\n",
        "            }\n",
        "         ]\n",
        "      }\n",
        "      #  ... repeat \"user\", \"model\" for multi turns.\n",
        "   ]\n",
        "}\n",
        "```\n",
        "\n",
        "Refer to the public [documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/models/gemini-supervised-tuning-prepare#about-datasets) for more details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I0YcID-jCKLk"
      },
      "outputs": [],
      "source": [
        "# combine the systeminstruct + context + question into one column.\n",
        "train_df = pd.read_csv(\"squad_train.csv\")\n",
        "validation_df = pd.read_csv(\"squad_validation.csv\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4DqrQp4cLqRy"
      },
      "outputs": [],
      "source": [
        "# combine the systeminstruct + context + question into one column.\n",
        "train_df[\"input_question\"] = (\n",
        "    \"\\n\\n **Below the question with context that you need to answer**\"\n",
        "    + \"\\n Context: \"\n",
        "    + train_df[\"context\"]\n",
        "    + \"\\n Question: \"\n",
        "    + train_df[\"question\"]\n",
        ")\n",
        "validation_df[\"input_question\"] = (\n",
        "    \"\\n\\n **Below the question with context that you need to answer**\"\n",
        "    + \"\\n Context: \"\n",
        "    + validation_df[\"context\"]\n",
        "    + \"\\n Question: \"\n",
        "    + validation_df[\"question\"]\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pmzyz1migvHN"
      },
      "outputs": [],
      "source": [
        "def df_to_jsonl(df, output_file):\n",
        "    \"\"\"Converts a Pandas DataFrame to JSONL format and saves it to a file.\n",
        "\n",
        "    Args:\n",
        "      df: The DataFrame to convert.\n",
        "      output_file: The name of the output file.\n",
        "    \"\"\"\n",
        "\n",
        "    with open(output_file, \"w\") as f:\n",
        "        for row in df.itertuples(index=False):\n",
        "            jsonl_obj = {\n",
        "                \"systemInstruction\": {\"parts\": [{\"text\": f\"{systemInstruct}\"}]},\n",
        "                \"contents\": [\n",
        "                    {\n",
        "                        \"role\": \"user\",\n",
        "                        \"parts\": [{\"text\": f\"{row.input_question}\"}],\n",
        "                    },\n",
        "                    {\"role\": \"model\", \"parts\": [{\"text\": row.answers}]},\n",
        "                ],\n",
        "            }\n",
        "            f.write(json.dumps(jsonl_obj) + \"\\n\")\n",
        "\n",
        "\n",
        "# Process the DataFrames\n",
        "df_to_jsonl(train_df, \"squad_train.jsonl\")\n",
        "df_to_jsonl(validation_df, \"squad_validation.jsonl\")\n",
        "\n",
        "print(f\"JSONL data written to squad_train.jsonl\")\n",
        "print(f\"JSONL data written to squad_validation.jsonl\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5OQv-ZMpJDhi"
      },
      "source": [
        "Next you will copy the files into your Google Cloud bucket"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O5k1jYJ10IeW"
      },
      "outputs": [],
      "source": [
        "!gsutil cp ./squad_train.jsonl {BUCKET_URI}\n",
        "!gsutil cp ./squad_validation.jsonl {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UAHMYgFJJHjm"
      },
      "source": [
        "### Start fine-tuning job\n",
        "Next you can start the fine-tuning job.\n",
        "\n",
        "- `source_model`: Specifies the base Gemini model version you want to fine-tune.\n",
        " - `train_dataset`: Path to your training data in JSONL format.\n",
        "\n",
        "  *Optional parameters*\n",
        " - `validation_dataset`: If provided, this data is used to evaluate the model during tuning.\n",
        " - `tuned_model_display_name`: Display name for the tuned model.\n",
        " - `epochs`: The number of training epochs to run.\n",
        " - `learning_rate_multiplier`: A value to scale the learning rate during training.\n",
        " - `adapter_size` : Gemini 2.0 supports Adapter length [1, 4], default value is 4.\n",
        "\n",
        " **Important**: The default hyperparameter settings are optimized for optimal performance based on rigorous testing and are recommended for initial use. Users may customize these parameters to address specific performance requirements.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gdcy4umfpGZE"
      },
      "outputs": [],
      "source": [
        "train_dataset = f\"\"\"{BUCKET_URI}/squad_train.jsonl\"\"\"\n",
        "validation_dataset = f\"\"\"{BUCKET_URI}/squad_train.jsonl\"\"\"\n",
        "\n",
        "training_dataset = {\n",
        "    \"gcs_uri\": train_dataset,\n",
        "}\n",
        "\n",
        "validation_dataset = types.TuningValidationDataset(gcs_uri=validation_dataset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NkboVUkoqWSp"
      },
      "outputs": [],
      "source": [
        "sft_tuning_job = client.tunings.tune(\n",
        "    base_model=base_model,\n",
        "    training_dataset=training_dataset,\n",
        "    config=types.CreateTuningJobConfig(\n",
        "        adapter_size=\"ADAPTER_SIZE_EIGHT\",\n",
        "        epoch_count=1,  # set to one to keep time and cost low\n",
        "        tuned_model_display_name=\"gemini-flash-1.5-qa\",\n",
        "    ),\n",
        ")\n",
        "sft_tuning_job"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UKo8cwF2KVM5"
      },
      "source": [
        "**Important:** Tuning time depends on several factors, such as training data size, number of epochs, learning rate multiplier, etc."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8NiZnPkIKcwm"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ It will take ~30 mins for the model tuning job to complete on the provided dataset and set configurations/hyperparameters. ⚠️</b>\n",
        "</div>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WECSLyPRth6M"
      },
      "outputs": [],
      "source": [
        "sft_tuning_job.state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_iwz4lhUDC_f"
      },
      "outputs": [],
      "source": [
        "tuning_job = client.tunings.get(name=sft_tuning_job.name)\n",
        "tuning_job"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gV1ukBznKmlN"
      },
      "source": [
        "#### Model tuning metrics\n",
        "\n",
        "- `/train_total_loss`: Loss for the tuning dataset at a training step.\n",
        "- `/train_fraction_of_correct_next_step_preds`: The token accuracy at a training step. A single prediction consists of a sequence of tokens. This metric measures the accuracy of the predicted tokens when compared to the ground truth in the tuning dataset.\n",
        "- `/train_num_predictions`: Number of predicted tokens at a training step\n",
        "\n",
        "#### Model evaluation metrics:\n",
        "\n",
        "- `/eval_total_loss`: Loss for the evaluation dataset at an evaluation step.\n",
        "- `/eval_fraction_of_correct_next_step_preds`: The token accuracy at an evaluation step. A single prediction consists of a sequence of tokens. This metric measures the accuracy of the predicted tokens when compared to the ground truth in the evaluation dataset.\n",
        "- `/eval_num_predictions`: Number of predicted tokens at an evaluation step.\n",
        "\n",
        "The metrics visualizations are available after the model tuning job completes. If you don't specify a validation dataset when you create the tuning job, only the visualizations for the tuning metrics are available."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_IoiiRH5Lhpf"
      },
      "outputs": [],
      "source": [
        "experiment_name = tuning_job.experiment\n",
        "experiment_name"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DH0guHM---Jo"
      },
      "outputs": [],
      "source": [
        "# Locate Vertex AI Experiment and Vertex AI Experiment Run\n",
        "experiment = aiplatform.Experiment(experiment_name=experiment_name)\n",
        "filter_str = metadata_utils._make_filter_string(\n",
        "    schema_title=\"system.ExperimentRun\",\n",
        "    parent_contexts=[experiment.resource_name],\n",
        ")\n",
        "experiment_run = context.Context.list(filter_str)[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hggHQFIl_FXC"
      },
      "outputs": [],
      "source": [
        "# Read data from Tensorboard\n",
        "tensorboard_run_name = f\"{experiment.get_backing_tensorboard_resource().resource_name}/experiments/{experiment.name}/runs/{experiment_run.name.replace(experiment.name, '')[1:]}\"\n",
        "tensorboard_run = aiplatform.TensorboardRun(tensorboard_run_name)\n",
        "metrics = tensorboard_run.read_time_series_data()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BdHKZdqG_bHf"
      },
      "outputs": [],
      "source": [
        "def get_metrics(metric: str = \"/train_total_loss\"):\n",
        "    \"\"\"\n",
        "    Get metrics from Tensorboard.\n",
        "\n",
        "    Args:\n",
        "      metric: metric name, eg. /train_total_loss or /eval_total_loss.\n",
        "    Returns:\n",
        "      steps: list of steps.\n",
        "      steps_loss: list of loss values.\n",
        "    \"\"\"\n",
        "    loss_values = metrics[metric].values\n",
        "    steps_loss = []\n",
        "    steps = []\n",
        "    for loss in loss_values:\n",
        "        steps_loss.append(loss.scalar.value)\n",
        "        steps.append(loss.step)\n",
        "    return steps, steps_loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_pDrlpA7_e9o"
      },
      "outputs": [],
      "source": [
        "# Get Train and Eval Loss\n",
        "train_loss = get_metrics(metric=\"/train_total_loss\")\n",
        "eval_loss = get_metrics(metric=\"/eval_total_loss\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DL07j7u__iZx"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<html>\n",
              "<head><meta charset=\"utf-8\" /></head>\n",
              "<body>\n",
              "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
              "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.35.2.min.js\"></script>                <div id=\"26f084c4-23e9-42bd-9090-155ccc4c6174\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"26f084c4-23e9-42bd-9090-155ccc4c6174\")) {                    Plotly.newPlot(                        \"26f084c4-23e9-42bd-9090-155ccc4c6174\",                        [{\"mode\":\"lines\",\"name\":\"Train Loss\",\"x\":[1,2,3,4,5,6,7,8,9,10,11,12],\"y\":[5.9871954917907715,6.5085577964782715,6.738699913024902,5.503544330596924,6.144056797027588,5.722049713134766,5.719411849975586,5.6194376945495605,4.656821250915527,4.039559364318848,3.0081570148468018,1.6816469430923462],\"type\":\"scatter\",\"xaxis\":\"x\",\"yaxis\":\"y\"},{\"mode\":\"lines\",\"name\":\"Eval Loss\",\"x\":[],\"y\":[],\"type\":\"scatter\",\"xaxis\":\"x2\",\"yaxis\":\"y2\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,0.45],\"title\":{\"text\":\"Steps\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"Loss\"}},\"xaxis2\":{\"anchor\":\"y2\",\"domain\":[0.55,1.0],\"title\":{\"text\":\"Steps\"}},\"yaxis2\":{\"anchor\":\"x2\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"Loss\"}},\"annotations\":[{\"font\":{\"size\":16},\"showarrow\":false,\"text\":\"Train Loss\",\"x\":0.225,\"xanchor\":\"center\",\"xref\":\"paper\",\"y\":1.0,\"yanchor\":\"bottom\",\"yref\":\"paper\"},{\"font\":{\"size\":16},\"showarrow\":false,\"text\":\"Eval Loss\",\"x\":0.775,\"xanchor\":\"center\",\"xref\":\"paper\",\"y\":1.0,\"yanchor\":\"bottom\",\"yref\":\"paper\"}],\"title\":{\"text\":\"Train and Eval Loss\"}},                        {\"responsive\": true}                    ).then(function(){\n",
              "                            \n",
              "var gd = document.getElementById('26f084c4-23e9-42bd-9090-155ccc4c6174');\n",
              "var x = new MutationObserver(function (mutations, observer) {{\n",
              "        var display = window.getComputedStyle(gd).display;\n",
              "        if (!display || display === 'none') {{\n",
              "            console.log([gd, 'removed!']);\n",
              "            Plotly.purge(gd);\n",
              "            observer.disconnect();\n",
              "        }}\n",
              "}});\n",
              "\n",
              "// Listen for the removal of the full notebook cells\n",
              "var notebookContainer = gd.closest('#notebook-container');\n",
              "if (notebookContainer) {{\n",
              "    x.observe(notebookContainer, {childList: true});\n",
              "}}\n",
              "\n",
              "// Listen for the clearing of the current output cell\n",
              "var outputEl = gd.closest('.output');\n",
              "if (outputEl) {{\n",
              "    x.observe(outputEl, {childList: true});\n",
              "}}\n",
              "\n",
              "                        })                };                            </script>        </div>\n",
              "</body>\n",
              "</html>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# Plot the train and eval loss metrics using Plotly python library\n",
        "fig = make_subplots(\n",
        "    rows=1, cols=2, shared_xaxes=True, subplot_titles=(\"Train Loss\", \"Eval Loss\")\n",
        ")\n",
        "\n",
        "# Add traces\n",
        "fig.add_trace(\n",
        "    go.Scatter(x=train_loss[0], y=train_loss[1], name=\"Train Loss\", mode=\"lines\"),\n",
        "    row=1,\n",
        "    col=1,\n",
        ")\n",
        "fig.add_trace(\n",
        "    go.Scatter(x=eval_loss[0], y=eval_loss[1], name=\"Eval Loss\", mode=\"lines\"),\n",
        "    row=1,\n",
        "    col=2,\n",
        ")\n",
        "\n",
        "# Add figure title\n",
        "fig.update_layout(title=\"Train and Eval Loss\", xaxis_title=\"Steps\", yaxis_title=\"Loss\")\n",
        "\n",
        "# Set x-axis title\n",
        "fig.update_xaxes(title_text=\"Steps\")\n",
        "\n",
        "# Set y-axes titles\n",
        "fig.update_yaxes(title_text=\"Loss\")\n",
        "\n",
        "# Show plot\n",
        "fig.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pivmh4Lwbgy1"
      },
      "source": [
        "### Use the fine-tuned model and evaluation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qO6ln4teagw1"
      },
      "outputs": [],
      "source": [
        "prompt = \"\"\"\n",
        "Answer the question based on the context\n",
        "\n",
        "Context: In the 1840s and 50s, there were attempts to overcome this problem by means of various patent valve gears with a separate, variable cutoff expansion valve riding on the back of the main slide valve; the latter usually had fixed or limited cutoff.\n",
        "The combined setup gave a fair approximation of the ideal events, at the expense of increased friction and wear, and the mechanism tended to be complicated.\n",
        "The usual compromise solution has been to provide lap by lengthening rubbing surfaces of the valve in such a way as to overlap the port on the admission side, with the effect that the exhaust side remains open for a longer period after cut-off on the admission side has occurred.\n",
        "This expedient has since been generally considered satisfactory for most purposes and makes possible the use of the simpler Stephenson, Joy and Walschaerts motions.\n",
        "Corliss, and later, poppet valve gears had separate admission and exhaust valves driven by trip mechanisms or cams profiled so as to give ideal events; most of these gears never succeeded outside of the stationary marketplace due to various other issues including leakage and more delicate mechanisms.\n",
        "\n",
        "Question: How is lap provided by overlapping the admission side port?\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QJgK_ZfL5dsj"
      },
      "outputs": [],
      "source": [
        "tuned_model = tuning_job.tuned_model.endpoint\n",
        "tuned_model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ifhRboiCOBje"
      },
      "outputs": [],
      "source": [
        "get_predictions(prompt, tuned_model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SyrEDGAoA4f5"
      },
      "outputs": [],
      "source": [
        "# Apply the get_prediction() function to the 'question_column'\n",
        "test_df[\"predicted_answer\"] = test_df[\"input_question\"].apply(get_predictions)\n",
        "test_df.head(2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IjQO_tqVFKf3"
      },
      "outputs": [],
      "source": [
        "test_df[\"predicted_answer\"] = test_df[\"predicted_answer\"].apply(normalize_answer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kBawjkvKQ_Q-"
      },
      "source": [
        "After running the evaluation you can see that the model generally performs better on our use case after fine-tuning. Of course, depending on things like use case or data quality performance will differ."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zjoRgjg6fr0P"
      },
      "outputs": [],
      "source": [
        "em, f1 = calculate_em_and_f1(test_df[\"answers\"], test_df[\"predicted_answer\"])\n",
        "print(f\"EM score: {em}\")\n",
        "print(f\"F1 score: {f1}\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "sft_gemini_qa.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}