wehub-resource-sync
781 lines
27 KiB
Plaintext
781 lines
27 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "5norOZI0mA6s"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Copyright 2023 Google LLC\n",
|
|
"#\n",
|
|
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
|
|
"# you may not use this file except in compliance with the License.\n",
|
|
"# You may obtain a copy of the License at\n",
|
|
"#\n",
|
|
"# https://www.apache.org/licenses/LICENSE-2.0\n",
|
|
"#\n",
|
|
"# Unless required by applicable law or agreed to in writing, software\n",
|
|
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
|
|
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
|
|
"# See the License for the specific language governing permissions and\n",
|
|
"# limitations under the License."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "XNPE46X8mJj4"
|
|
},
|
|
"source": [
|
|
"# Use Retrieval Augmented Generation (RAG) with Gemini API\n",
|
|
"\n",
|
|
"<table align=\"left\">\n",
|
|
"\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_retrieval_augmented_generation.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",
|
|
"\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%2Fuse-cases%2Fcode%2Fcode_retrieval_augmented_generation.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://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_retrieval_augmented_generation.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/use-cases/code/code_retrieval_augmented_generation.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",
|
|
"</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/use-cases/code/code_retrieval_augmented_generation.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/use-cases/code/code_retrieval_augmented_generation.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/use-cases/code/code_retrieval_augmented_generation.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/use-cases/code/code_retrieval_augmented_generation.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/use-cases/code/code_retrieval_augmented_generation.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": "VrLtlKPFqSxB"
|
|
},
|
|
"source": [
|
|
"| | |\n",
|
|
"|-|-|\n",
|
|
"|Author(s) | [Lavi Nigam](https://github.com/lavinigam-gcp), [Polong Lin](https://github.com/polong-lin) |"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "zNAEdYNFmQcP"
|
|
},
|
|
"source": [
|
|
"### Objective\n",
|
|
"\n",
|
|
"This notebook demonstrates how you augment output from Gemini API by bringing in external knowledge. An example is provided using Code Retrieval Augmented Generation(RAG) pattern using [Google Cloud's Generative AI github repository](https://github.com/GoogleCloudPlatform/generative-ai) as external knowledge. The notebook uses [Gemini API in Vertex AI](https://ai.google.dev/gemini-api), [Embeddings for Text API](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings), FAISS vector store and [LangChain 🦜️🔗](https://python.langchain.com/en/latest/).\n",
|
|
"\n",
|
|
"### Overview\n",
|
|
"\n",
|
|
"Here is overview of what we'll go over.\n",
|
|
"\n",
|
|
"Index Creation:\n",
|
|
"\n",
|
|
"1. Recursively list the files(.ipynb) in github repo\n",
|
|
"2. Extract code and markdown from the files\n",
|
|
"3. Chunk & generate embeddings for each code strings and add initialize the vector store\n",
|
|
"\n",
|
|
"Runtime:\n",
|
|
"\n",
|
|
"4. User enters a prompt or asks a question as a prompt\n",
|
|
"5. Try zero-shot prompt\n",
|
|
"6. Run prompt using RAG Chain & compare results.To generate response we use **gemini-2.0-flash**\n",
|
|
"\n",
|
|
"### Cost\n",
|
|
"\n",
|
|
"This tutorial uses billable components of Google Cloud:\n",
|
|
"\n",
|
|
"- Gemini API in Vertex AI offered by Google Cloud\n",
|
|
"\n",
|
|
"Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage.\n",
|
|
"\n",
|
|
"**Note:** We are using local vector store(FAISS) for this example however recommend managed highly scalable vector store for production usage such as [Vertex AI Vector Search](https://cloud.google.com/vertex-ai/docs/vector-search/overview) or [AlloyDB for PostgreSQL](https://cloud.google.com/alloydb/docs/ai/work-with-embeddings) or [Cloud SQL for PostgreSQL](https://cloud.google.com/sql/docs/postgres/features) using pgvector extension."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "2cab0c8509c9"
|
|
},
|
|
"source": [
|
|
"## Get started"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "b56b5a5d28c1"
|
|
},
|
|
"source": [
|
|
"### Install Vertex AI SDK for Python and other required packages\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "QHaqV20Csqkt"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install --upgrade --user -q google-cloud-aiplatform \\\n",
|
|
" langchain \\\n",
|
|
" langchain_google_vertexai \\\n",
|
|
" langchain-community \\\n",
|
|
" faiss-cpu \\\n",
|
|
" nbformat"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "-VUWOgz6M1rZ"
|
|
},
|
|
"source": [
|
|
"### Restart runtime (Colab only)\n",
|
|
"\n",
|
|
"To use the newly installed packages, you must restart the runtime on Google Colab."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "BIS8EYgkMy8T"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"import sys\n",
|
|
"\n",
|
|
"if \"google.colab\" in sys.modules:\n",
|
|
" import IPython\n",
|
|
"\n",
|
|
" app = IPython.Application.instance()\n",
|
|
" app.kernel.do_shutdown(True)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "0af13c10a26a"
|
|
},
|
|
"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": "uZcP9WBENG0e"
|
|
},
|
|
"source": [
|
|
"### Authenticate your notebook environment (Colab only)\n",
|
|
"\n",
|
|
"Authenticate your environment on Google Colab.\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "1S_HgQXQNcbz"
|
|
},
|
|
"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": "rVmxMr43Nhoo"
|
|
},
|
|
"source": [
|
|
"### Import libraries"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "L-Tljm5asMBc"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"import time\n",
|
|
"\n",
|
|
"from google.cloud import aiplatform\n",
|
|
"from langchain.chains import RetrievalQA\n",
|
|
"from langchain.prompts import PromptTemplate\n",
|
|
"from langchain.schema.document import Document\n",
|
|
"from langchain.text_splitter import Language, RecursiveCharacterTextSplitter\n",
|
|
"from langchain.vectorstores import FAISS\n",
|
|
"\n",
|
|
"# LangChain\n",
|
|
"from langchain_google_vertexai import VertexAI, VertexAIEmbeddings\n",
|
|
"import nbformat\n",
|
|
"import requests\n",
|
|
"\n",
|
|
"# Vertex AI\n",
|
|
"import vertexai\n",
|
|
"\n",
|
|
"# Print the version of Vertex AI SDK for Python\n",
|
|
"print(f\"Vertex AI SDK version: {aiplatform.__version__}\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "4f872cd812d0"
|
|
},
|
|
"source": [
|
|
"### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
|
|
"\n",
|
|
"To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "eNGEcBKG0iK-"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Initialize project\n",
|
|
"# Define project information\n",
|
|
"PROJECT_ID = \"YOUR_PROJECT_ID\" # @param {type:\"string\"}\n",
|
|
"LOCATION = \"us-central1\" # @param {type:\"string\"}\n",
|
|
"\n",
|
|
"vertexai.init(project=PROJECT_ID, location=LOCATION)\n",
|
|
"\n",
|
|
"# Code Generation\n",
|
|
"code_llm = VertexAI(\n",
|
|
" model_name=\"gemini-2.0-flash\",\n",
|
|
" max_output_tokens=2048,\n",
|
|
" temperature=0.1,\n",
|
|
" verbose=False,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "o537exyZk9DI"
|
|
},
|
|
"source": [
|
|
"Next we need to create a GitHub personal token to be able to list all files in a repository.\n",
|
|
"\n",
|
|
"- Follow [this link](https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens) to create GitHub token with repo->public_repo scope and update `GITHUB_TOKEN` variable below."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "Bt9IVDSqk7y4"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# provide GitHub personal access token\n",
|
|
"GITHUB_TOKEN = \"YOUR_GITHUB_TOKEN\" # @param {type:\"string\"}\n",
|
|
"GITHUB_REPO = \"GoogleCloudPlatform/generative-ai\" # @param {type:\"string\"}"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "dqq3GeEbOJbU"
|
|
},
|
|
"source": [
|
|
"# Index Creation\n",
|
|
"\n",
|
|
"We use the Google Cloud Generative AI github repository as the data source. First list all Jupyter Notebook files in the repo and store it in a text file.\n",
|
|
"\n",
|
|
"You can skip this step(#1) if you have executed it once and generated the output text file.\n",
|
|
"\n",
|
|
"### 1. Recursively list the files(.ipynb) in the github repository"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "eTA1Jt0uOX8y"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Crawls a GitHub repository and returns a list of all ipynb files in the repository\n",
|
|
"\n",
|
|
"\n",
|
|
"def crawl_github_repo(url: str, is_sub_dir: bool, access_token: str = GITHUB_TOKEN):\n",
|
|
" ignore_list = [\"__init__.py\"]\n",
|
|
"\n",
|
|
" if not is_sub_dir:\n",
|
|
" api_url = f\"https://api.github.com/repos/{url}/contents\"\n",
|
|
"\n",
|
|
" else:\n",
|
|
" api_url = url\n",
|
|
"\n",
|
|
" headers = {\n",
|
|
" \"Accept\": \"application/vnd.github.v3+json\",\n",
|
|
" \"Authorization\": f\"Bearer {access_token}\",\n",
|
|
" }\n",
|
|
"\n",
|
|
" response = requests.get(api_url, headers=headers)\n",
|
|
" response.raise_for_status() # Check for any request errors\n",
|
|
"\n",
|
|
" files = []\n",
|
|
"\n",
|
|
" contents = response.json()\n",
|
|
"\n",
|
|
" for item in contents:\n",
|
|
" if (\n",
|
|
" item[\"type\"] == \"file\"\n",
|
|
" and item[\"name\"] not in ignore_list\n",
|
|
" and (item[\"name\"].endswith(\".py\") or item[\"name\"].endswith(\".ipynb\"))\n",
|
|
" ):\n",
|
|
" files.append(item[\"html_url\"])\n",
|
|
" elif item[\"type\"] == \"dir\" and not item[\"name\"].startswith(\".\"):\n",
|
|
" sub_files = crawl_github_repo(item[\"url\"], True)\n",
|
|
" time.sleep(0.1)\n",
|
|
" files.extend(sub_files)\n",
|
|
"\n",
|
|
" return files"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "5vaKaxcGO_R6"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"code_files_urls = crawl_github_repo(GITHUB_REPO, False, GITHUB_TOKEN)\n",
|
|
"\n",
|
|
"# Write list to a file so you do not have to download each time\n",
|
|
"with open(\"code_files_urls.txt\", \"w\") as f:\n",
|
|
" for item in code_files_urls:\n",
|
|
" f.write(item + \"\\n\")\n",
|
|
"\n",
|
|
"len(code_files_urls)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "c5hoNYJ5byMJ"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"code_files_urls[0:10]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "mFNVieLnR8Ie"
|
|
},
|
|
"source": [
|
|
"### 2. Extract code from the Jupyter notebooks.\n",
|
|
"\n",
|
|
"You could also include .py file, shell scripts etc."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "ZsM1M4hn4cBu"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Extracts the python code from an ipynb file from github\n",
|
|
"\n",
|
|
"\n",
|
|
"def extract_python_code_from_ipynb(github_url, cell_type=\"code\"):\n",
|
|
" raw_url = github_url.replace(\"github.com\", \"raw.githubusercontent.com\").replace(\n",
|
|
" \"/blob/\", \"/\"\n",
|
|
" )\n",
|
|
"\n",
|
|
" response = requests.get(raw_url)\n",
|
|
" response.raise_for_status() # Check for any request errors\n",
|
|
"\n",
|
|
" notebook_content = response.text\n",
|
|
"\n",
|
|
" notebook = nbformat.reads(notebook_content, as_version=nbformat.NO_CONVERT)\n",
|
|
"\n",
|
|
" python_code = None\n",
|
|
"\n",
|
|
" for cell in notebook.cells:\n",
|
|
" if cell.cell_type == cell_type:\n",
|
|
" if not python_code:\n",
|
|
" python_code = cell.source\n",
|
|
" else:\n",
|
|
" python_code += \"\\n\" + cell.source\n",
|
|
"\n",
|
|
" return python_code\n",
|
|
"\n",
|
|
"\n",
|
|
"def extract_python_code_from_py(github_url):\n",
|
|
" raw_url = github_url.replace(\"github.com\", \"raw.githubusercontent.com\").replace(\n",
|
|
" \"/blob/\", \"/\"\n",
|
|
" )\n",
|
|
"\n",
|
|
" response = requests.get(raw_url)\n",
|
|
" response.raise_for_status() # Check for any request errors\n",
|
|
"\n",
|
|
" python_code = response.text\n",
|
|
"\n",
|
|
" return python_code"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "WCRp5Xtb48is"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"with open(\"code_files_urls.txt\") as f:\n",
|
|
" code_files_urls = f.read().splitlines()\n",
|
|
"len(code_files_urls)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "4Y9SMO7H4xgF"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"code_strings = []\n",
|
|
"\n",
|
|
"for i in range(0, len(code_files_urls)):\n",
|
|
" if code_files_urls[i].endswith(\".ipynb\"):\n",
|
|
" content = extract_python_code_from_ipynb(code_files_urls[i], \"code\")\n",
|
|
" doc = Document(\n",
|
|
" page_content=content, metadata={\"url\": code_files_urls[i], \"file_index\": i}\n",
|
|
" )\n",
|
|
" code_strings.append(doc)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "T1AF3fhBSLOm"
|
|
},
|
|
"source": [
|
|
"### 3. Chunk & generate embeddings for each code strings & initialize the vector store\n",
|
|
"\n",
|
|
"We need to split code into usable chunks that the LLM can use for code generation. Therefore it's crucial to use the right chunking approach and chunk size."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "Rj1cCA2fqx64"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Utility functions for Embeddings API with rate limiting\n",
|
|
"\n",
|
|
"\n",
|
|
"def rate_limit(max_per_minute):\n",
|
|
" period = 60 / max_per_minute\n",
|
|
" print(\"Waiting\")\n",
|
|
" while True:\n",
|
|
" before = time.time()\n",
|
|
" yield\n",
|
|
" after = time.time()\n",
|
|
" elapsed = after - before\n",
|
|
" sleep_time = max(0, period - elapsed)\n",
|
|
" if sleep_time > 0:\n",
|
|
" print(\".\", end=\"\")\n",
|
|
" time.sleep(sleep_time)\n",
|
|
"\n",
|
|
"\n",
|
|
"class CustomVertexAIEmbeddings(VertexAIEmbeddings):\n",
|
|
" requests_per_minute: int\n",
|
|
" num_instances_per_batch: int\n",
|
|
" model_name: str\n",
|
|
"\n",
|
|
" # Overriding embed_documents method\n",
|
|
" def embed_documents(\n",
|
|
" self, texts: list[str], batch_size: int | None = None\n",
|
|
" ) -> list[list[float]]:\n",
|
|
" limiter = rate_limit(self.requests_per_minute)\n",
|
|
" results = []\n",
|
|
" docs = list(texts)\n",
|
|
"\n",
|
|
" while docs:\n",
|
|
" # Working in batches because the API accepts maximum 5\n",
|
|
" # documents per request to get embeddings\n",
|
|
" head, docs = (\n",
|
|
" docs[: self.num_instances_per_batch],\n",
|
|
" docs[self.num_instances_per_batch :],\n",
|
|
" )\n",
|
|
" chunk = self.client.get_embeddings(head)\n",
|
|
" results.extend(chunk)\n",
|
|
" next(limiter)\n",
|
|
"\n",
|
|
" return [r.values for r in results]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "oae37l-pvzZ6"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Chunk code strings\n",
|
|
"text_splitter = RecursiveCharacterTextSplitter.from_language(\n",
|
|
" language=Language.PYTHON, chunk_size=2000, chunk_overlap=200\n",
|
|
")\n",
|
|
"\n",
|
|
"\n",
|
|
"texts = text_splitter.split_documents(code_strings)\n",
|
|
"print(len(texts))\n",
|
|
"\n",
|
|
"# Initialize Embedding API\n",
|
|
"EMBEDDING_QPM = 100\n",
|
|
"EMBEDDING_NUM_BATCH = 5\n",
|
|
"embeddings = CustomVertexAIEmbeddings(\n",
|
|
" requests_per_minute=EMBEDDING_QPM,\n",
|
|
" num_instances_per_batch=EMBEDDING_NUM_BATCH,\n",
|
|
" model_name=\"text-embedding-005\",\n",
|
|
")\n",
|
|
"\n",
|
|
"# Create Index from embedded code chunks\n",
|
|
"db = FAISS.from_documents(texts, embeddings)\n",
|
|
"\n",
|
|
"# Init your retriever.\n",
|
|
"retriever = db.as_retriever(\n",
|
|
" search_type=\"similarity\", # Also test \"similarity\", \"mmr\"\n",
|
|
" search_kwargs={\"k\": 5},\n",
|
|
")\n",
|
|
"\n",
|
|
"retriever"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "Q_gn89IyuHIT"
|
|
},
|
|
"source": [
|
|
"# Runtime\n",
|
|
"### 4. User enters a prompt or asks a question as a prompt"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "1vrvTkO7uFNi"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"user_question = \"Create a Python function that takes a prompt and predicts using langchain.llms interface with Vertex AI Gemini model\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "azbvOUFRvEp5"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Define prompt templates\n",
|
|
"\n",
|
|
"# Zero Shot prompt template\n",
|
|
"prompt_zero_shot = \"\"\"\n",
|
|
" You are a proficient python developer. Respond with the syntactically correct & concise code for to the question below.\n",
|
|
"\n",
|
|
" Question:\n",
|
|
" {question}\n",
|
|
"\n",
|
|
" Output Code :\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
"prompt_prompt_zero_shot = PromptTemplate(\n",
|
|
" input_variables=[\"question\"],\n",
|
|
" template=prompt_zero_shot,\n",
|
|
")\n",
|
|
"\n",
|
|
"\n",
|
|
"# RAG template\n",
|
|
"prompt_RAG = \"\"\"\n",
|
|
" You are a proficient python developer. Respond with the syntactically correct code for to the question below. Make sure you follow these rules:\n",
|
|
" 1. Use context to understand the APIs and how to use it & apply.\n",
|
|
" 2. Do not add license information to the output code.\n",
|
|
" 3. Do not include Colab code in the output.\n",
|
|
" 4. Ensure all the requirements in the question are met.\n",
|
|
"\n",
|
|
" Question:\n",
|
|
" {question}\n",
|
|
"\n",
|
|
" Context:\n",
|
|
" {context}\n",
|
|
"\n",
|
|
" Helpful Response :\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
"prompt_RAG_template = PromptTemplate(\n",
|
|
" template=prompt_RAG, input_variables=[\"context\", \"question\"]\n",
|
|
")\n",
|
|
"\n",
|
|
"qa_chain = RetrievalQA.from_llm(\n",
|
|
" llm=code_llm,\n",
|
|
" prompt=prompt_RAG_template,\n",
|
|
" retriever=retriever,\n",
|
|
" return_source_documents=True,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "3NBaObAQSlIv"
|
|
},
|
|
"source": [
|
|
"### 5. Try zero-shot prompt"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "1svTVwtBS0zP"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"response = code_llm.invoke(input=user_question, max_output_tokens=2048, temperature=0.1)\n",
|
|
"print(response)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "JPm8qdxzwPM0"
|
|
},
|
|
"source": [
|
|
"### 6. Run prompt using RAG Chain & compare results\n",
|
|
"To generate response we use code-bison however can also use code-gecko and codechat-bison"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "ZMz3nPMyVoj_"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"results = qa_chain.invoke(input={\"query\": user_question})\n",
|
|
"print(results[\"result\"])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "HF3lVWK1wjxe"
|
|
},
|
|
"source": [
|
|
"### Let's try another prompt"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "jel0ON68XiU7"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"user_question = \"Create python function that takes text input and returns embeddings using LangChain with Vertex AI text-embedding-005 model\"\n",
|
|
"\n",
|
|
"\n",
|
|
"response = code_llm.invoke(input=user_question, max_output_tokens=2048, temperature=0.1)\n",
|
|
"print(response)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "G9bIkqE8sO6P"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"results = qa_chain.invoke(input={\"query\": user_question})\n",
|
|
"print(results[\"result\"])"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"colab": {
|
|
"name": "code_retrieval_augmented_generation.ipynb",
|
|
"toc_visible": true
|
|
},
|
|
"kernelspec": {
|
|
"display_name": "Python 3",
|
|
"name": "python3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 0
|
|
}
|