langchain-vertex-ai/create_text_embeddings_using_langchain.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "84a0b6ea",
   "metadata": {},
   "source": [
    "# Create Text Embeddings for a Vector Store using LangChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d21bd2ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "!pip install --quiet langchain chromadb==0.5.3\n",
    "!pip install langchain-community\n",
    "!pip install langchain-google-vertexai\n",
    "!pip install --upgrade --quiet langchain-google-genai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "84400dd7-e520-438b-bc51-387e94c0219b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'status': 'ok', 'restart': True}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Restart kernel after installs so that your environment can access the new packages\n",
    "import IPython\n",
    "import time\n",
    "\n",
    "app = IPython.Application.instance()\n",
    "app.kernel.do_shutdown(True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fff9cd4-76cd-417d-920a-8dad29cc68c2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mRunning cells with 'Python 3.11.7' requires the ipykernel package.\n",
      "\u001b[1;31mRun the following command to install 'ipykernel' into the Python environment. \n",
      "\u001b[1;31mCommand: 'c:/Python3.11/python.exe -m pip install ipykernel -U --user --force-reinstall'"
     ]
    }
   ],
   "source": [
    "# Define project information\n",
    "import sys\n",
    "\n",
    "PROJECT_ID = \"PROJECT_ID\"  # @param {type:\"string\"} Please set your PROJECT_ID\n",
    "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
    "\n",
    "# if not running on colab, try to get the PROJECT_ID automatically\n",
    "if \"google.colab\" not in sys.modules:\n",
    "    import subprocess\n",
    "\n",
    "    PROJECT_ID = subprocess.check_output(\n",
    "        [\"gcloud\", \"config\", \"get-value\", \"project\"], text=True\n",
    "    ).strip()\n",
    "\n",
    "print(f\"Your project ID is: {PROJECT_ID}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0ea5f25b-0ecc-4d47-b676-1705befc959c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "USER_AGENT environment variable not set, consider setting it to identify your requests.\n"
     ]
    }
   ],
   "source": [
    "from langchain import PromptTemplate\n",
    "from langchain import hub\n",
    "from langchain.docstore.document import Document\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain.document_loaders import WebBaseLoader\n",
    "from langchain.schema import StrOutputParser\n",
    "from langchain.schema.prompt_template import format_document\n",
    "from langchain.schema.runnable import RunnablePassthrough\n",
    "from langchain.vectorstores import Chroma\n",
    "from langchain_google_vertexai import VertexAIEmbeddings\n",
    "from langchain_google_vertexai import VertexAI\n",
    "from langchain.chains import RetrievalQA, ConversationalRetrievalChain\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "969fae3e-703d-4b16-bd82-323bae373d4b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Task 1. Use the WebBaseLoader to load documents related to queries\n",
    "loader = WebBaseLoader(\"https://blog.google/technology/ai/google-gemini-ai/\")\n",
    "docs = loader.load()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6a78e63-2c30-48f6-a4ee-ef78d2645dc2",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Split text into chunks\n",
    "Now that we have the documents we will split them into chunks. Each chunk will become one vector in the vector store. To do this we will define a chunk size (number of characters) and a chunk overlap (amount of overlap i.e. sliding window). The perfect chunk size can be difficult to determine. Too large of a chunk size leads to too much information per chunk (individual chunks not specific enough), however too small of a chunk size leads to not enough information per chunk. In both cases, nearest neighbors lookup with a query/question embedding may struggle to retrieve the actually relevant chunks, or fail altogether if the chunks are too large to use as context with an LLM query.\n",
    "\n",
    "In this notebook we will use a chunk size of 800 chacters and a chunk overlap of 100 characters, but feel free to experiment with other sizes! Note: you can specify a custom `length_function` with `RecursiveCharacterTextSplitter` if you want chunk size/overlap to be determined by something other than Python's `len` function. In addition to `RecursiveCharacterTextSplitter`, there are [other text splitters](https://python.langchain.com/docs/modules/data_connection/document_transformers/text_splitters/split_by_token) you can consider."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3c2f7fa7-06b0-4018-a714-ac89e05adfbf",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(metadata={'source': 'https://blog.google/technology/ai/google-gemini-ai/', 'title': 'Introducing Gemini: Google’s most capable AI model yet', 'description': 'Gemini is our most capable and general model, built to be multimodal and optimized for three different sizes: Ultra, Pro and Nano.', 'language': 'en-us'}, page_content='Introducing Gemini: Google’s most capable AI model yet'),\n",
       " Document(metadata={'source': 'https://blog.google/technology/ai/google-gemini-ai/', 'title': 'Introducing Gemini: Google’s most capable AI model yet', 'description': 'Gemini is our most capable and general model, built to be multimodal and optimized for three different sizes: Ultra, Pro and Nano.', 'language': 'en-us'}, page_content='[{\"model\": \"blogsurvey.survey\", \"pk\": 3, \"fields\": {\"name\": \"General Article Sentiment\", \"survey_id\": \"general-article-sentiment_240906\", \"scroll_depth_trigger\": 50, \"previous_survey\": null, \"display_rate\": 50, \"thank_message\": \"Thank you!\", \"thank_emoji\": \"✅\", \"questions\": \"[{\\\\\"id\\\\\": \\\\\"32a784f8-0a4d-44e9-962b-b9db9422d98b\\\\\", \\\\\"type\\\\\": \\\\\"simple_question\\\\\", \\\\\"value\\\\\": {\\\\\"question\\\\\": \\\\\"Overall, how did you feel about this article?\\\\\", \\\\\"responses\\\\\": [{\\\\\"id\\\\\": \\\\\"d1c06504-53a7-4704-a79b-8deb17a6e072\\\\\", \\\\\"type\\\\\": \\\\\"item\\\\\", \\\\\"value\\\\\": \\\\\"Liked it very much\\\\\"}, {\\\\\"id\\\\\": \\\\\"35c540e2-a10e-491f-b4eb-bebfcf90784b\\\\\", \\\\\"type\\\\\": \\\\\"item\\\\\", \\\\\"value\\\\\": \\\\\"Liked it somewhat\\\\\"}, {\\\\\"id\\\\\": \\\\\"76b8658c-7cc9-43ac-8243-9a6153fb4504\\\\\", \\\\\"type\\\\\": \\\\\"item\\\\\", \\\\\"value\\\\\": \\\\\"Neutral\\\\\"}, {\\\\\"id\\\\\":')]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Task 2. Use the RecursiveCharacterTextSplitter class to split the documents into chunks for embedding\n",
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size = 800,\n",
    "    chunk_overlap  = 100,\n",
    "    length_function = len,\n",
    ")\n",
    "\n",
    "chunks = text_splitter.split_documents(docs)\n",
    "\n",
    "# Look at the first two chunks \n",
    "chunks[0:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f479a543-2382-4359-a15f-96d3e86b732e",
   "metadata": {},
   "source": [
    "### Vectorize/Embed Document Chunks\n",
    "Now we need to embed the document chunks (turn them into vectors) and store them in a vector store. For this, we can use any text embedding model, however we need to be sure to use the same text embedding model when we embed our queries/questions at prediction time. To make things simple we will use the PaLM API for Embeddings. The LangChain library provides a wrapper class around the PaLM Embeddings API, `VertexAIEmbeddings()`.\n",
    "\n",
    "For the purposes of this lab, you will use [Chroma](https://www.trychroma.com/) as the vector store for simplicity. In a real-world scenario with a large private knowledge-base, you may not be able to fit everything in memory. Langchain has a nice wrapper class for Chroma which allows us to pass in a list of documents, and an embedding class to create the vector store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "635f212c-461e-468a-a2c2-4461a1f4cf55",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Task 3. Create vector store using embeddings\n",
    "from langchain_google_genai import GoogleGenerativeAIEmbeddings\n",
    "\n",
    "gemini_embeddings = VertexAIEmbeddings(model=\"text-embedding-004\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6dc41f30-c0d8-4365-b844-a735e5e4c763",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/tmp/ipykernel_83542/1233409443.py:2: LangChainDeprecationWarning: The class `Chroma` was deprecated in LangChain 0.2.9 and will be removed in 1.0. An updated version of the class exists in the langchain-chroma package and should be used instead. To use it run `pip install -U langchain-chroma` and import as `from langchain_chroma import Chroma`.\n",
      "  vectorstore = Chroma(embedding_function=gemini_embeddings, persist_directory=\"./vectorstore\")\n"
     ]
    }
   ],
   "source": [
    "# Save to disk\n",
    "vectorstore = Chroma(embedding_function=gemini_embeddings, persist_directory=\"./vectorstore\")\n",
    "\n",
    "for chunk in chunks:\n",
    "    vectorstore.add_documents([chunk])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2ef299c9-cad7-440f-91f8-9f1f7a065b79",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n"
     ]
    }
   ],
   "source": [
    "# Load from disk\n",
    "vectorstore_disk = Chroma(\n",
    "    embedding_function=gemini_embeddings,   # Embedding model\n",
    "    persist_directory=\"./vectorstore\"       # Directory to save the embeddings\n",
    ")\n",
    "\n",
    "retriever = vectorstore_disk.as_retriever(search_kwargs={\"k\": 1})\n",
    "print(len(retriever.invoke(\"MMLU\")))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bcec8b4-575d-47e4-b0b6-1fa1378fc46c",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Putting it all together\n",
    "Now that everything is in place, we can tie it all together with a langchain chain. A langchain chain simply orchestrates the multiple steps required to use an LLM for a specific use case. In this case the process we will chain together first embeds the query/question, then performs a nearest neighbors lookup to find the relevant chunks, then uses the relevant chunks to formulate a response with an LLM. We will use the Chroma database as our vector store and PaLM as our LLM. Langchain provides a wrapper around PaLM, `VertexAI()`. \n",
    "\n",
    "For this simple Q/A use case we can use langchain's `RetrievalQA` to link together the process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "255dce30-5f7e-49dc-a4f2-635120c85775",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# vector store \n",
    "retriever = vectorstore.as_retriever(\n",
    "    search_type=\"similarity\",\n",
    "    search_kwargs={\"k\":2} # number of nearest neighbors to retrieve  \n",
    ")\n",
    "\n",
    "# You can also set temperature, top_p, top_k \n",
    "llm = VertexAI(\n",
    "    model_name=\"text-bison\",\n",
    "    max_output_tokens=1024\n",
    ")\n",
    "\n",
    "# q/a chain \n",
    "qa = RetrievalQA.from_chain_type(\n",
    "    llm=llm,\n",
    "    chain_type=\"stuff\",\n",
    "    retriever=retriever,\n",
    "    return_source_documents=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b841b15-26fb-4dd1-bc75-16901412f02e",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Query \n",
    "Now that everything is \"chained\" together using LangChain we can send queries and get answers! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c50eb258-abe5-4628-be1d-0b63775d0a8f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def ask_question(question: str):\n",
    "    response = qa.invoke({\"query\": question})\n",
    "    print(f\"Response: {response['result']}\\n\")\n",
    "\n",
    "    citations = {doc.metadata['source'] for doc in response['source_documents']}\n",
    "    print(f\"Citations: {citations}\\n\")\n",
    "\n",
    "    # uncomment below to print source chunks used  \n",
    "    # print(f\"Source Chunks Used: {response['source_documents']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "113eb28a-5af6-4a5c-96f4-2d7d6f5a68c6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response:  MMLU stands for massive multitask language understanding. It is a benchmark that uses a combination of 57 subjects such as math, physics, history, law, medicine, and ethics for testing both world knowledge and problem-solving abilities.\n",
      "\n",
      "Citations: {'https://blog.google/technology/ai/google-gemini-ai/'}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "ask_question(\"What is MMLU?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6b400e66-fc81-4f56-83a9-deac64a49a2d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response:  A TPU is a Tensor Processing Unit, a custom-designed AI accelerator developed by Google.\n",
      "\n",
      "Citations: {'https://blog.google/technology/ai/google-gemini-ai/'}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "ask_question(\"What is a TPU?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbbf9733-f3f4-40fc-93b7-637379733a2e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "tf2-cpu.2-16.m124",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/tf2-cpu.2-16:m124"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}