notebooks/community/model_garden/model_garden_openai_api_llama3_1.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "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": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Vertex AI Model Garden - Get started with Llama 3.1 models\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_openai_api_llama3_1.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_openai_api_llama3_1.ipynb\"\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>    \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_openai_api_llama3_1.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_openai_api_llama3_1.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to get started with using the OpenAI library and demonstrates how to use Llama 3.1 models as Model-as-service (MaaS) for building translation chain and document question-answer.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Configure OpenAI SDK for the Llama 3.1 Completions API\n",
        "- Chat with Llama 3.1 models with different prompts and model parameters, and apply Llama Guard for safeguarding\n",
        "- Build with Llama 3.1 models\n",
        "  - Translation Chain.\n",
        "  - A RAG application using LLamaIndex on Vertex AI.\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform[langchain] openai\n",
        "! pip3 install --upgrade --quiet langchain-openai"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "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",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"<your-project-id>\"  # @param {type:\"string\"}\n",
        "\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store tutorial artifacts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = \"<your-bucket-name>\"  # @param {type:\"string\"}\n",
        "\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EcIXiGsCePi"
      },
      "source": [
        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Wn8ZkcV86KR"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "B8DawN9D9NLU"
      },
      "outputs": [],
      "source": [
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jVYoyDl165EE"
      },
      "source": [
        "### Import libraries\n",
        "\n",
        "Import libraries to use in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "c1tEW-U968h8"
      },
      "outputs": [],
      "source": [
        "# Chat completions API\n",
        "import openai\n",
        "from google.auth import default, transport\n",
        "from langchain import PromptTemplate\n",
        "# Build\n",
        "from langchain_openai import ChatOpenAI\n",
        "from vertexai.preview import rag"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uqYCG2Fw7D3L"
      },
      "source": [
        "### Configure OpenAI SDK for the Llama 3.1 Chat Completions API\n",
        "\n",
        "To configure the OpenAI SDK for the Llama 3.1 Chat Completions API, you need to request the access token and initialize the client pointing to the Llama 3.1 endpoint.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W0K6VSJRHhH2"
      },
      "source": [
        "#### Authentication\n",
        "\n",
        "You can request an access token from the default credentials for the current environment. Note that the access token lives for [1 hour by default](https://cloud.google.com/docs/authentication/token-types#at-lifetime); after expiration, it must be refreshed.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "i0qceuiQEPHv"
      },
      "outputs": [],
      "source": [
        "credentials, _ = default()\n",
        "auth_request = transport.requests.Request()\n",
        "credentials.refresh(auth_request)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q04wJmA0HT6X"
      },
      "source": [
        "Then configure the OpenAI SDK to point to the Llama 3.1 Chat Completions API endpoint.\n",
        "\n",
        "Notice, only `us-central1` is supported region for Llama 3.1 models using Model-as-a-Service (MaaS)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "c-MRhsnlj6iw"
      },
      "outputs": [],
      "source": [
        "MODEL_LOCATION = \"us-central1\"\n",
        "\n",
        "client = openai.OpenAI(\n",
        "    base_url=f\"https://{MODEL_LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{MODEL_LOCATION}/endpoints/openapi/chat/completions?\",\n",
        "    api_key=credentials.token,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UGokrtdiIHrX"
      },
      "source": [
        "#### Llama 3.1 models\n",
        "\n",
        "You can experiment with various supported Llama 3.1 models.\n",
        "\n",
        "This tutorial use Llama 3 8B Instruct, 70B Instruct, and 405B Instruct using Model-as-a-Service (MaaS). Using Model-as-a-Service (MaaS), you can access Llama 3.1 models in just a few clicks without any setup or infrastructure hassles. Model-as-a-Service (MaaS) integrates [Llama Guard](https://huggingface.co/meta-llama/Llama-Guard-3-8B) as a safety filter. It is switched on by default and can be switched off. Llama Guard enables us to safeguard model inputs and outputs. If a response is filtered, it will be populated with a `finish_reason` field (with value `content_filtered`) and a `refusal` field (stating the filtering reason). \n",
        "\n",
        "You can also access Llama models for self-service in Vertex AI Model Garden, allowing you to choose your preferred infrastructure. [Check out Llama 3.1 model card](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama3_1?_ga=2.31261500.2048242469.1721714335-1107467625.1721655511) to learn how to deploy a Llama 3.1 models on Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "r7OhyH46H2H5"
      },
      "outputs": [],
      "source": [
        "MODEL_ID = \"meta/llama-3.1-405b-instruct-maas\"  # @param {type:\"string\"} [\"meta/llama-3.1-8b-instruct-maas\", \"meta/llama-3.1-70b-instruct-maas\", \"meta/llama-3.1-405b-instruct-maas\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1xD62NTpqHXd"
      },
      "source": [
        "### Chat with Llama 3.1\n",
        "\n",
        "Use the Chat Completions API to send a request to the Llama 3.1 model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tkyp9kZSuJGx"
      },
      "source": [
        "#### Hello, Llama 3!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "CKVOZ1HEqRbY"
      },
      "outputs": [],
      "source": [
        "apply_llama_guard = True  # @param {type:\"boolean\"}\n",
        "\n",
        "response = client.chat.completions.create(\n",
        "    model=MODEL_ID,\n",
        "    messages=[{\"role\": \"user\", \"content\": \"Hello, Llama 3.1!\"}],\n",
        "    extra_body={\n",
        "        \"extra_body\": {\n",
        "            \"google\": {\n",
        "                \"model_safety_settings\": {\n",
        "                    \"enabled\": apply_llama_guard,\n",
        "                    \"llama_guard_settings\": {},\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "LxpdxYCxH51u"
      },
      "outputs": [],
      "source": [
        "print(response.choices[0].message.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B1rKbHUQt605"
      },
      "source": [
        "#### Ask Llama 3.1 using different model configuration\n",
        "\n",
        "Use the following parameters to generate different answers:\n",
        "\n",
        "*   `temperature` to control the randomness of the response\n",
        "*   `max_tokens` to limit the response length\n",
        "*   `top_p` to control the quality of the response\n",
        "*   `stream` to stream the response back or not\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "owv-5Sz5rIEU"
      },
      "outputs": [],
      "source": [
        "temperature = 1.0  # @param {type:\"number\"}\n",
        "max_tokens = 50  # @param {type:\"integer\"}\n",
        "top_p = 1.0  # @param {type:\"number\"}\n",
        "stream = True  # @param {type:\"boolean\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a-qBuhcK-G1V"
      },
      "source": [
        "Get the answer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "O1YU8bSivH0B"
      },
      "outputs": [],
      "source": [
        "apply_llama_guard = True  # @param {type:\"boolean\"}\n",
        "\n",
        "response = client.chat.completions.create(\n",
        "    model=MODEL_ID,\n",
        "    messages=[\n",
        "        {\"role\": \"user\", \"content\": \"What is Vertex AI?\"},\n",
        "        {\"role\": \"assistant\", \"content\": \"Sure, Vertex AI is:\"},\n",
        "    ],\n",
        "    temperature=temperature,\n",
        "    max_tokens=max_tokens,\n",
        "    top_p=top_p,\n",
        "    stream=stream,\n",
        "    extra_body={\n",
        "        \"extra_body\": {\n",
        "            \"google\": {\n",
        "                \"model_safety_settings\": {\n",
        "                    \"enabled\": apply_llama_guard,\n",
        "                    \"llama_guard_settings\": {},\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-o9-gF0U-Kba"
      },
      "source": [
        "Depending if `stream` parameter is enabled or not, you can print the response entirely or chunk by chunk."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "CoDHLGhyyt8d"
      },
      "outputs": [],
      "source": [
        "if stream:\n",
        "    for chunk in response:\n",
        "        print(chunk.choices[0].delta.content, end=\"\")\n",
        "else:\n",
        "    print(response.choices[0].message.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BkoaelaKxm1r"
      },
      "source": [
        "#### Use Llama 3.1 with different tasks\n",
        "\n",
        "In this section, you will use Llama 3.1 to perform different tasks including text generation, text summarization, and code generation.\n",
        "\n",
        "For each task, you'll define a different prompt and submit a request to the model as you did before."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-en7AYQDyONt"
      },
      "source": [
        "##### Text Generation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "QANInNvizWbi"
      },
      "outputs": [],
      "source": [
        "prompt = \"Write a poem about a cat who loves to code\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2x8ML1Y_yfom"
      },
      "outputs": [],
      "source": [
        "apply_llama_guard = True  # @param {type:\"boolean\"}\n",
        "\n",
        "response = client.chat.completions.create(\n",
        "    model=MODEL_ID,\n",
        "    messages=[\n",
        "        {\"role\": \"user\", \"content\": prompt},\n",
        "    ],\n",
        "    extra_body={\n",
        "        \"extra_body\": {\n",
        "            \"google\": {\n",
        "                \"model_safety_settings\": {\n",
        "                    \"enabled\": apply_llama_guard,\n",
        "                    \"llama_guard_settings\": {},\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "qQ6UUgpHztXZ"
      },
      "outputs": [],
      "source": [
        "print(response.choices[0].message.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kBLESIw4zhto"
      },
      "source": [
        "##### Text summarization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "UrAybklfzhtz"
      },
      "outputs": [],
      "source": [
        "article = \"\"\"\n",
        "Vertex AI: Google's Unified Platform for Machine Learning\n",
        "\n",
        "Google Cloud's Vertex AI is a comprehensive platform that simplifies the process of building, deploying, and managing machine learning (ML) models and AI applications. It provides a single environment for all your AI needs, from data preparation to model deployment and monitoring.\n",
        "\n",
        "Vertex AI offers a range of features to cater to various user levels, including:\n",
        "\n",
        "AutoML: This feature allows you to train models on tabular, image, text, or video data without writing code. It's ideal for users without extensive ML expertise.\n",
        "Custom Training: For advanced users, Vertex AI provides custom training options, allowing you to use your preferred ML framework and write your own code.\n",
        "Model Garden: This feature lets you discover, test, and deploy pre-trained models from Vertex AI and open-source sources.\n",
        "Generative AI: Access Google's powerful large language models (LLMs) to generate text, code, images, and speech, which can be customized and deployed for your applications.\n",
        "Vertex AI seamlessly integrates with other Google Cloud services like BigQuery for data warehousing, Cloud Storage for data management, and Cloud AI Platform for custom model training. It provides managed infrastructure that can be tailored to your performance and budget needs.\n",
        "\n",
        "Whether you're a seasoned data scientist or just starting out with AI, Vertex AI simplifies the entire ML lifecycle and empowers you to build and deploy AI solutions effectively.\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "prompt = (\"Summarize the following article in one sentence: \" + article).replace(\n",
        "    \"\\n\", \"\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "VJYIZbGyzhtz"
      },
      "outputs": [],
      "source": [
        "apply_llama_guard = True  # @param {type:\"boolean\"}\n",
        "\n",
        "response = client.chat.completions.create(\n",
        "    model=MODEL_ID,\n",
        "    messages=[\n",
        "        {\"role\": \"user\", \"content\": prompt},\n",
        "    ],\n",
        "    extra_body={\n",
        "        \"extra_body\": {\n",
        "            \"google\": {\n",
        "                \"model_safety_settings\": {\n",
        "                    \"enabled\": apply_llama_guard,\n",
        "                    \"llama_guard_settings\": {},\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "9vfWA2i9zwOZ"
      },
      "outputs": [],
      "source": [
        "print(response.choices[0].message.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-_mNB0Fezh6G"
      },
      "source": [
        "##### Code generation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ltbhrGiwzh6H"
      },
      "outputs": [],
      "source": [
        "prompt = \"Write a Python function that takes a list of numbers and returns the average. Include error handling for empty lists.\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "q66yE4Pszh6H"
      },
      "outputs": [],
      "source": [
        "apply_llama_guard = True  # @param {type:\"boolean\"}\n",
        "\n",
        "response = client.chat.completions.create(\n",
        "    model=MODEL_ID,\n",
        "    messages=[\n",
        "        {\"role\": \"user\", \"content\": prompt},\n",
        "    ],\n",
        "    extra_body={\n",
        "        \"extra_body\": {\n",
        "            \"google\": {\n",
        "                \"model_safety_settings\": {\n",
        "                    \"enabled\": apply_llama_guard,\n",
        "                    \"llama_guard_settings\": {},\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "mcAf5tXrtPIu"
      },
      "outputs": [],
      "source": [
        "print(response.choices[0].message.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gnrXpv5Y3yFK"
      },
      "source": [
        "### Build with Llama 3.1\n",
        "\n",
        "In this section, you use Llama 3.1 to build two simple applications. In order:\n",
        "\n",
        "1.   **Translation Chain** to translate text across multiple languages using Llama 3.1 and LangChain Expression Language (LCEL).\n",
        "\n",
        "2.   **Document Q&A with RAG using LLamaIndex on Vertex AI** to answer questions about documents with retrieval augmented generation, powered by Llama 3.1."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IowGcZq95HqZ"
      },
      "source": [
        "#### Translation chain\n",
        "\n",
        "In this scenario, you use LangChain Expression Language (LCEL) to build a simple chain which translates some `text_to_translate` to the specified `target_language`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yLAwaYPzFqDQ"
      },
      "source": [
        "##### Initialize the chat interface and the translation prompt template using LangChain"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "CE2KxIrG5xKC"
      },
      "outputs": [],
      "source": [
        "llm = ChatOpenAI(\n",
        "    model=MODEL_ID,\n",
        "    base_url=f\"https://{MODEL_LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT_ID}/locations/{MODEL_LOCATION}/endpoints/openapi/chat/completions?\",\n",
        "    api_key=credentials.token,\n",
        ")\n",
        "\n",
        "template = \"\"\"Translate the following {text} to {target_language}:\"\"\"\n",
        "\n",
        "prompt = PromptTemplate(input_variables=[\"text\", \"target_language\"], template=template)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "odTLqzLiF8h_"
      },
      "source": [
        "##### Initialize the chain"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "74Mywe4W9MmE"
      },
      "outputs": [],
      "source": [
        "chain = prompt | llm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QE3sizbzGFER"
      },
      "source": [
        "##### Translate a text"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "SRJ-xuSI9ZQl"
      },
      "outputs": [],
      "source": [
        "text_to_translate = \"Hello Llama 3.1!\"  # @param {type:\"string\"}\n",
        "target_language = \"Italian\"  # @param {type:\"string\"}\n",
        "\n",
        "response = chain.invoke({\"text\": text_to_translate, \"target_language\": target_language})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "yYzc1kCjEHGP"
      },
      "outputs": [],
      "source": [
        "print(response.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y1VIksOZx4pP"
      },
      "source": [
        "#### Document Q&A using LlamaIndex on Vertex AI for RAG\n",
        "\n",
        "In this scenario, you'll use LlamaIndex on Vertex AI for RAG to build a document Q&A.\n",
        "\n",
        "LlamaIndex on Vertex AI helps you with the end-to-end process of building and deploying  context-augmented large language model (LLM) applications including retrieval-augmented generation (RAG), from ingesting data from various sources, to transforming it for indexing, and creating numerical representations (embeddings) for semantic understanding. Then, when a user provides a query, LlamaIndex on Vertex AI retrieves relevant information and uses it as context to generate accurate and relevant responses.\n",
        "\n",
        "[Refer to the documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/llamaindex-on-vertexai) to find detailed instructions on how to get started with the RAG API with LlamaIndex on Vertex AI."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cmxhW-LArK2L"
      },
      "source": [
        "##### Create a RAG corpus\n",
        "\n",
        "You'll start by creating an index (i.e., corpus) using one of the supported embedding models, for example `text-embedding-004`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "5E1tVMx3rAXF"
      },
      "outputs": [],
      "source": [
        "# Currently supports Google first-party embedding models\n",
        "embedding_model_config = rag.EmbeddingModelConfig(\n",
        "    publisher_model=\"publishers/google/models/text-embedding-004\"\n",
        ")\n",
        "\n",
        "# Name your corpus\n",
        "DISPLAY_NAME = \"<your-corpus-display-name>\"  # @param {type:\"string\"}\n",
        "\n",
        "rag_corpus = rag.create_corpus(\n",
        "    display_name=DISPLAY_NAME, embedding_model_config=embedding_model_config\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "16i1ZInQrFnL"
      },
      "outputs": [],
      "source": [
        "# Check the corpus just created\n",
        "rag.list_corpora()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sSLWYGF8rfMf"
      },
      "source": [
        "##### Import documents to an existing corpus from Google Cloud Storage.\n",
        "\n",
        "For the purpose of this notebook, you'll create a document of interesting facts about llamas and upload it directly to the newly created corpus."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "4G5uyvbdraMY"
      },
      "outputs": [],
      "source": [
        "%%writefile llama_facts.txt\n",
        "Beyond the Spit: Unpacking the Wondrous World of Llamas\n",
        "They're the fluffy, doe-eyed denizens of the Andes, known for their luxurious wool and, let's face it, their rather pungent spitting habits. But beyond these surface-level quirks lies a fascinating world of adaptation, intelligence, and surprising utility. Buckle up, folks, because we're diving deep into the wondrous world of llamas!\n",
        "\n",
        "From Mountain Tops to Your Backyard: Llama Evolution is No Myth\n",
        "Forget unicorns and dragons, the llama's origin story is where the real magic lies. Evolving in the unforgiving terrains of the North American plains millions of years ago, these camelids (yes, they're related to camels!) migrated southwards, eventually conquering the challenging Andean highlands. Their impressive adaptation to high altitude, with their unique blood composition that efficiently transports oxygen, is a testament to their evolutionary prowess.\n",
        "\n",
        "Pack Animal Extraordinaire: More Than Just a Pretty Fleece\n",
        "While their luxurious wool, prized for its warmth and softness, is a major draw, llamas are much more than walking sweaters. For centuries, indigenous communities have relied on these gentle giants as pack animals, capable of carrying impressive loads (up to 100 pounds!) across treacherous mountain paths. Forget horsepower, in the Andes, it's all about llama-power!\n",
        "\n",
        "Social Butterflies with a Side of Spit: Decoding Llama Communication\n",
        "Llamas are highly social creatures, living in herds led by a dominant male. Their communication is a fascinating mix of soft hums, alarm calls, and yes, the infamous spit. But here's the kicker: llamas primarily reserve their spitting for each other, often to settle disputes or establish dominance within the herd. So, unless you're challenging a llama to a staring contest (not recommended!), you're unlikely to be on the receiving end of that projectile saliva.\n",
        "\n",
        "Guardian of the Flock: The Llama's Unexpected Talent\n",
        "Move over, sheepdogs, there's a new sheriff in town! Farmers have discovered the llama's remarkable ability to protect livestock from predators. Their size, assertive nature, and surprisingly powerful kicks are enough to deter coyotes, foxes, and even stray dogs. This unexpected talent has earned them the title of \"Guardian Llamas,\" a role they take very seriously, patrolling their territory and fiercely defending their woolly companions.\n",
        "\n",
        "The Future is Fuzzy: Llamas Beyond the Andes\n",
        "With their gentle nature, intelligence, and surprising versatility, llamas are steadily gaining popularity beyond their traditional Andean home. From therapy animals providing comfort and emotional support to eco-friendly lawnmowers, the possibilities seem endless. So, the next time you encounter a llama, remember, there's more to them than meets the eye (or the spit!). They are a testament to nature's ingenuity, captivating us with their unique charm and leaving us in awe of their remarkable journey through time."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PcpEoi0yCpAq"
      },
      "source": [
        "Upload the text file to the RAG corpus."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "c6nkMtmp75Le"
      },
      "outputs": [],
      "source": [
        "rag_file = rag.upload_file(\n",
        "    corpus_name=rag_corpus.name,\n",
        "    path=\"llama_facts.txt\",\n",
        "    display_name=\"llama_facts\",\n",
        "    description=\"Llama Facts for RAG test\",\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "YiTAFiEasHLX"
      },
      "outputs": [],
      "source": [
        "# Check the files just uploaded.\n",
        "list(rag.list_files(corpus_name=rag_corpus.name))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cje1WHtOtt2q"
      },
      "source": [
        "##### Generate Content\n",
        "\n",
        "You can provide a query and use RAG API to retrieve relevant content from the corpus and concantenate them as the context of the input prompt to generate accurate and relevant responses with Llama 3.1."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "BEGIIcY3_Db7"
      },
      "outputs": [],
      "source": [
        "question = \"What about llama spitting?\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qfZ2ph8fyw5N"
      },
      "source": [
        "Retrieve relevant content from the corpus and concantenate them as the context."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "dK7YmoIGtyki"
      },
      "outputs": [],
      "source": [
        "context = \" \".join(\n",
        "    [\n",
        "        context.text\n",
        "        for context in rag.retrieval_query(\n",
        "            rag_resources=[\n",
        "                rag.RagResource(\n",
        "                    rag_corpus=rag_corpus.name,\n",
        "                )\n",
        "            ],\n",
        "            text=question,\n",
        "            similarity_top_k=1,\n",
        "            vector_distance_threshold=0.5,\n",
        "        ).contexts.contexts\n",
        "    ]\n",
        ").replace(\"\\n\", \"\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "OSUL8QtDtYO2"
      },
      "outputs": [],
      "source": [
        "context"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sEWxni06zSMv"
      },
      "source": [
        "Generate content with the retrieved context."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "R5eLmNQS4TFh"
      },
      "outputs": [],
      "source": [
        "apply_llama_guard = True  # @param {type:\"boolean\"}\n",
        "\n",
        "response = client.chat.completions.create(\n",
        "    model=MODEL_ID,\n",
        "    messages=[\n",
        "        {\n",
        "            \"role\": \"system\",\n",
        "            \"content\": \"\"\"You are an AI assistant. Your goal is to answer questions using the pieces of context. If you don't know the answer, say that you don't know.\"\"\",\n",
        "        },\n",
        "        {\"role\": \"user\", \"content\": question},\n",
        "        {\"role\": \"assistant\", \"content\": context},\n",
        "    ],\n",
        "    extra_body={\n",
        "        \"extra_body\": {\n",
        "            \"google\": {\n",
        "                \"model_safety_settings\": {\n",
        "                    \"enabled\": apply_llama_guard,\n",
        "                    \"llama_guard_settings\": {},\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "AYUGyginvU6h"
      },
      "outputs": [],
      "source": [
        "print(response.choices[0].message.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NF6hj3XlzZEm"
      },
      "source": [
        "For comparison, generate content without the context."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "JCoPy3ApupcG"
      },
      "outputs": [],
      "source": [
        "apply_llama_guard = True  # @param {type:\"boolean\"}\n",
        "\n",
        "response = client.chat.completions.create(\n",
        "    model=MODEL_ID,\n",
        "    messages=[\n",
        "        {\n",
        "            \"role\": \"system\",\n",
        "            \"content\": \"\"\"You are an AI assistant. Your goal is to answer questions using the pieces of context. If you don't know the answer, say that you don't know.\"\"\",\n",
        "        },\n",
        "        {\"role\": \"user\", \"content\": question},\n",
        "    ],\n",
        "    extra_body={\n",
        "        \"extra_body\": {\n",
        "            \"google\": {\n",
        "                \"model_safety_settings\": {\n",
        "                    \"enabled\": apply_llama_guard,\n",
        "                    \"llama_guard_settings\": {},\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Wc6FtOJaD7BG"
      },
      "outputs": [],
      "source": [
        "print(response.choices[0].message.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a4e033321ad"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "Clean up resources created in this notebook.\n",
        "\n",
        "To delete to the search engine in Vertex AI, check out the following [documentation](https://cloud.google.com/generative-ai-app-builder/docs/delete-engine)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "OC7Ypb05ccUE"
      },
      "outputs": [],
      "source": [
        "delete_rag_corpus = False  # @param {type:\"boolean\"}\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "\n",
        "if delete_rag_corpus:\n",
        "    rag_corpus_list = rag.list_corpora()\n",
        "    for rag_corpus in rag_corpus_list:\n",
        "        rag.delete_corpus(name=rag_corpus.name)\n",
        "\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_NAME"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_openai_api_llama3_1.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}