PyTorch/RAG_Application/RAG-on-Intel-Gaudi.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "70c2adb6-5a7b-48e6-ab13-1f50f8c18572",
   "metadata": {},
   "source": [
    "Copyright (c) 2024 Habana Labs, Ltd. an Intel Company.|"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca62fb01-b39e-4bd5-a70c-426579ad9f19",
   "metadata": {},
   "source": [
    "##### Licensed under the Apache License, Version 2.0 (the \"License\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbdde1fc-2079-4023-9d2f-1f0c7d5adede",
   "metadata": {},
   "source": [
    "you may not use this file except in compliance with the License. You may obtain a copy of the License at https://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64ff5422-73e6-4026-b5d2-57ef2e79b10e",
   "metadata": {},
   "source": [
    "# Retrieval Augmented Generation (RAG) application using Intel&reg; Gaudi&reg; 2 AI Processor\n",
    "A scalable Retrieval Augmented Generation (RAG) application using huggingface tools as an way of deploying optimized applications utilizing the Intel Gaudi 2 acclerator.\n",
    "\n",
    "### Introduction\n",
    "This tutorial will show how to build RAG application using Intel Gaudi 2. The Application will be built from easily accessible huggingface tools such as: text-generation-inference (TGI) and text-embeddings-inference (TEI). To make the code easier to understand, Langchain will be used.  The User interface at the end of the tutorial will use Gradio to submit your queries. This application will be in docker environment, but can be easily deployed to a kubernetes cluster.\n",
    "\n",
    "Retrieval-augmented generation (RAG) is a method that enhances the precision and dependability of generative AI models by incorporating facts from external sources. This technique addresses the limitations of large language models (LLMs), which, despite their ability to generate responses to general prompts rapidly, may not provide in-depth or specific information. By enabling access to external knowledge sources, RAG improves factual consistency, increases the reliability of generated responses, and helps to mitigate the issue of \"hallucination\" in more complex and knowledge-intensive tasks.\n",
    "\n",
    "This Tutorial will show the steps of building the full RAG pipeline on Intel Gaudi 2.  To first ingest a text file and run an embedding model to create a vector store index and database.  Then to start to run a query, it will run the embedding model again on the query, match it with the contents of the database and send the overall prompt and query response to the Llama 2 Large Language model to generate a full formed response. \n",
    "\n",
    "<p style=\"text-align:center;\">\n",
    "    <img src=\"./img/rag-overview.png\" />\n",
    "</p>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d096077b-df4c-4c33-9655-cf2c6b8899d0",
   "metadata": {},
   "source": [
    "## 1. Intial Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcf9afcb",
   "metadata": {},
   "source": [
    "**NOTE: this tutorial requires two Intel Gaudi AI Accelerators, one for text generation and one for text embedding.**  If you only have access to one Intel Gaudi card, please try the [Single Card Version](https://github.com/HabanaAI/Gaudi-tutorials/blob/main/PyTorch/Single_card_tutorials/RAG-on-Intel-Gaudi-single.ipynb) of the RAG tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "635c1f50-008c-48ed-a3f4-8f0853a5f5da",
   "metadata": {},
   "source": [
    "There are the initial steps to ensure that your build environment is set correctly:\n",
    "\n",
    "1. Set the appropriate ports for access when you ssh into the Intel Gaudi 2 node.  you need to ensure that the following ports are open:\n",
    "* 8888 (for running this jupyter notebook)\n",
    "* 7680 (for run the gradio server)\n",
    "Do to this, you need to add the following in your overall ssh commmand when connecting to the Intel Gaudi Node:\n",
    "\n",
    "`ssh -L 8888:localhost:8888 -L 7860:localhost:7860 .... `\n",
    "  \n",
    "   \n",
    "2. Before you load this Notebook, you will run the standard docker image but you need to include the `/var/run/docker.sock` file.  Use these Run and exec commands below to start your docker. \n",
    "\n",
    "`docker run -itd --name RAG --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host -v /var/run/docker.sock:/var/run/docker.sock  vault.habana.ai/gaudi-docker/1.18/ubuntu22.04/habanalabs/pytorch-installer-2.4.0:latest`  \n",
    "\n",
    "`docker exec -it RAG bash`\n",
    "\n",
    "`cd ~ && git clone https://github.com/HabanaAI/Gaudi-tutorials`\n",
    "\n",
    "If you find connectivity issues you can pass proxy information to docker container adding the following to the `docker run` command:\n",
    "`-e https_proxy=<your_proxy_address> -e HTTP_PROXY=<your_proxy_address> -e no_proxy=\"localhost;127.0.0.1\"`\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "256a0ea1-09f9-44ca-aff5-d20b21ee995b",
   "metadata": {},
   "source": [
    "#### Setup the docker environment in this notebook:\n",
    "At this point you have cloned the Gaudi-tutorials notebook inside your docker image and have opened this notebook.  Now start to follow the steps.  Note that you will need to install docker again inside the Intel Gaudi container to manage the execution of the RAG tools. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "085eaa61-6cad-4309-9432-f9ae8edff7c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!apt-get update\n",
    "!apt-get install docker.io curl -y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b77c1151-a505-4797-adb7-7582255bb58d",
   "metadata": {},
   "source": [
    "## 2. Loading the Tools for RAG\n",
    "There are three steps in creating the RAG environment, text generation, text embedding and vectorization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c8617c5-7cc4-4304-9d81-bae2c4e8afbd",
   "metadata": {},
   "source": [
    "### Text Generation Inference (TGI)\n",
    "First building block of application will be text-generation-inference, it's purpose will be serving the LLM model that will answer question based on context. To run it, we will pull the existing TGI-Gaudi docker image:\n",
    "\n",
    "Please note: The Hugging Face Text Generation Interface depends on software that is subject to non-open source licenses.  If you use or redistribute this software, it is your sole responsibility to ensure compliance with such licenses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f98c6bad-efcd-4dc9-b7ee-6749e4d5be04",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker pull ghcr.io/huggingface/tgi-gaudi:2.0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d422efe4-4135-43bb-b36c-0aa5067306b9",
   "metadata": {},
   "source": [
    "### After building image you will run run it:\n",
    "\n",
    "#### How to access and Use the Llama 2 model\n",
    "To use the Llama 2 model, you will need a HuggingFace account, agree to the terms of use of the model in its model card on the HF Hub, and create a read token.  You then copy that token to the HF_TOKEN variable below. \n",
    "\n",
    "Use of the pretrained model is subject to compliance with third party licenses, including the “Llama 2 Community License Agreement” (LLAMAV2). For guidance on the intended use of the LLAMA2 model, what will be considered misuse and out-of-scope uses, who are the intended users and additional terms please review and read the instructions in this link https://ai.meta.com/llama/license/. Users bear sole liability and responsibility to follow and comply with any third party licenses, and Habana Labs disclaims and will bear no liability with respect to users’ use or compliance with third party licenses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbf144b1-7390-4a73-be07-7be1feb683c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker run -d -p 9001:80 \\\n",
    "    --runtime=habana \\\n",
    "    --name gaudi-tgi \\\n",
    "    -e HABANA_VISIBLE_DEVICES=0 \\\n",
    "    -e OMPI_MCA_btl_vader_single_copy_mechanism=none \\\n",
    "    -e HF_TOKEN=\"<enter your hf token>\" \\\n",
    "    --cap-add=sys_nice \\\n",
    "    --ipc=host \\\n",
    "    ghcr.io/huggingface/tgi-gaudi:2.0.5  \\\n",
    "    --model-id meta-llama/Llama-2-7b-chat-hf \\\n",
    "    --max-input-tokens 1024 --max-total-tokens 2048"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a52b470-c992-442f-a3f7-6c2dbbc70d78",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker logs gaudi-tgi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db319c26-0252-4114-b209-24df991deee4",
   "metadata": {},
   "source": [
    "After running the docker server, it will take some time to download the model and load it into the device. To check the status run: `docker logs gaudi-tgi` and you should see:\n",
    "\n",
    "```\n",
    "2024-02-23T16:24:35.125179Z  INFO shard-manager: text_generation_launcher: Waiting for shard to be ready... rank=0\n",
    "2024-02-23T16:24:40.729388Z  INFO shard-manager: text_generation_launcher: Shard ready in 65.710470677s rank=0\n",
    "2024-02-23T16:24:40.796775Z  INFO text_generation_launcher: Starting Webserver\n",
    "2024-02-23T16:24:42.589516Z  WARN text_generation_router: router/src/main.rs:355: `--revision` is not set\n",
    "2024-02-23T16:24:42.589551Z  WARN text_generation_router: router/src/main.rs:356: We strongly advise to set it to a known supported commit.\n",
    "2024-02-23T16:24:42.842098Z  INFO text_generation_router: router/src/main.rs:377: Serving revision e852bc2e78a3fe509ec28c6d76512df3012acba7 of model Intel/neural-chat-7b-v3-1\n",
    "2024-02-23T16:24:42.845898Z  INFO text_generation_router: router/src/main.rs:219: Warming up model\n",
    "2024-02-23T16:24:42.846613Z  WARN text_generation_router: router/src/main.rs:230: Model does not support automatic max batch total tokens\n",
    "2024-02-23T16:24:42.846620Z  INFO text_generation_router: router/src/main.rs:252: Setting max batch total tokens to 16000\n",
    "2024-02-23T16:24:42.846623Z  INFO text_generation_router: router/src/main.rs:253: Connected\n",
    "2024-02-23T16:24:42.846626Z  WARN text_generation_router: router/src/main.rs:258: Invalid hostname, defaulting to 0.0.0.0\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18af7431-24af-4a8c-9234-8ccf3b17aaa1",
   "metadata": {},
   "source": [
    "Once the setup is complete, you can verify that that the text generation is working by sending a request to it (note that first request could be slow due to graph compilation):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "4ad5600d-915f-425b-a979-d603c8ea6270",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"generated_text\":\"\\n\\nThe Earth is round because of the way it was formed and the forces that have acted upon it over time. Here are some reasons why the Earth is round:\\n\\n1. Gravity: The Earth's round shape is primarily due to the force of gravity. Gravity pulls the Earth's material towards its center, causing it to collapse into a sphere.\\n2. Accretion: The Earth formed through a process called accretion, where small particles of rock and dust stuck together to form larger bodies. As these bodies grew, they collided and merged, eventually forming the Earth. The resulting sphere shape is a result of the way these particles stuck together.\\n3. Rotation: The Earth's rotation also plays a role in its shape. As the Earth spins on its axis, it experiences a centrifugal force that pushes material away from the axis, causing it to bulge at the equator. This bulge is more\"}"
     ]
    }
   ],
   "source": [
    "!curl 127.0.0.1:9001/generate \\\n",
    "    -X POST \\\n",
    "    -d '{\"inputs\":\"why is the earth round?\",\"parameters\":{\"max_new_tokens\":200}}' \\\n",
    "    -H 'Content-Type: application/json'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ee60a02-7ea9-490b-b0b2-f9f120b26a2a",
   "metadata": {},
   "source": [
    "### Text Embedding Interface (TEI)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbf035ef-9167-4dad-99d9-10ef906d8de5",
   "metadata": {},
   "source": [
    "Next building block will be text-embeddings-inference, it's purpose will be serving embeddings model that will produce embedings for vector database. To run it, we need to download the docker image:\n",
    "\n",
    "Please note: The Hugging Face Text Embedding Interface depends on software that is subject to non-open source licenses. If you use or redistribute this software, it is your sole responsibility to ensure compliance with such licenses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3dbd4c7-f2bf-405d-aa2f-d19ebcc82ad9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker pull ghcr.io/huggingface/tei-gaudi:1.5.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a899130-f25e-4d89-bb2a-c7878f84e40b",
   "metadata": {},
   "source": [
    "After building the image we can run it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a88b2b90-12de-4441-9c8c-105cb7b8279b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker run \\\n",
    "    -d \\\n",
    "    -p 9002:80 \\\n",
    "    --runtime=habana \\\n",
    "    --name gaudi-tei \\\n",
    "    -e HABANA_VISIBLE_DEVICES=1 \\\n",
    "    -e OMPI_MCA_btl_vader_single_copy_mechanism=none \\\n",
    "    -e MAX_WARMUP_SEQUENCE_LENGTH=512 \\\n",
    "    --cap-add=sys_nice \\\n",
    "    --ipc=host \\\n",
    "    ghcr.io/huggingface/tei-gaudi:1.5.0 \\\n",
    "    --model-id BAAI/bge-large-en-v1.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f6def2f-7883-447f-9924-cba2b10a4c73",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker logs gaudi-tei"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00eb11f7-843f-497e-a10a-8e9481538187",
   "metadata": {},
   "source": [
    "Note that here you need also wait for model to load.  You can run `docker logs gaudi-tei` to confirm that the model is setup and running.  To confirm that the model is working, we can send a request to it.   Running the command below will show embeddings of the input prompt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d780cfd-4688-4d9d-8abc-c0028f345ccb",
   "metadata": {},
   "outputs": [],
   "source": [
    "!curl 127.0.0.1:9002/embed \\\n",
    "    -X POST \\\n",
    "    -d '{\"inputs\":\"What is Deep Learning?\"}' \\\n",
    "    -H 'Content-Type: application/json'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d656f11-bee9-4fa6-bb60-e94254ac67e2",
   "metadata": {},
   "source": [
    "### PGVector\n",
    "Third building block is a vector database, in this tutorial the choice was PGVector. Setting up the docker should be straightforward:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f514fee-0afc-4b9d-9abc-4373f29f0bb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker pull pgvector/pgvector:pg16\n",
    "!docker run \\\n",
    "    --name postgres_vectordb \\\n",
    "    -d \\\n",
    "    -e POSTGRES_PASSWORD=postgres \\\n",
    "    -p 9003:5432 \\\n",
    "    pgvector/pgvector:pg16"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e19bca6-2a8d-4b9c-925f-3b8a325cecc1",
   "metadata": {},
   "source": [
    "### Application Front End \n",
    "Last building block will be a frontend that will serve as a http server. Frontend is implemented in python using the Gradio interface. To setup environment we need to run:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d73d8634-57e6-40c0-a60f-2e05830121f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8ff367a-c01d-4f10-96f8-f9eeb799c9d8",
   "metadata": {},
   "source": [
    "## 3. Data preparation\n",
    "\n",
    "To have a good quality RAG application, we need to prepare data. Data processing for vector database is following, extract text information from documents (for example PDFs, CSVs) than split it into chunks not exceeding max length, with additional metadata (for example filename or file creation date). Than upload preprocessed data to vector database. \n",
    "\n",
    "In the process of data preprocessing, text splitting plays a crucial role. It involves breaking down the text into smaller, semantically meaningful chunks for further processing and analysis. Here are some common methods of text splitting:\n",
    "\n",
    "- **By Character**: This method involves splitting the text into individual characters. It's a straightforward approach, but it may not always be the most effective, as it doesn't take into account the semantic meaning of words or phrases.\n",
    "\n",
    "- **Recursive**: Recursive splitting involves breaking down the text into smaller parts repeatedly until a certain condition is met. This method is particularly useful when dealing with complex structures in the text, as it allows for a more granular level of splitting.\n",
    "\n",
    "- **HTML Specific**: When dealing with HTML content, text splitting can be done based on specific HTML tags or elements. This method is useful for extracting meaningful information from web pages or other HTML documents.\n",
    "\n",
    "- **Code Specific**: In the context of programming code, text can be split based on specific code syntax or structures. This method is particularly useful for code analysis or for building tools that work with code.\n",
    "\n",
    "- **By Tokens**: Tokenization is a common method of text splitting in Natural Language Processing (NLP). It involves breaking down the text into individual words or tokens. This method is effective for understanding the semantic meaning of the text, as it allows for the analysis of individual words and their context.\n",
    "\n",
    "In conclusion, the choice of text splitting method depends largely on the nature of the text and the specific requirements of the task at hand. It's important to choose a method that effectively captures the semantic meaning of the text and facilitates further processing and analysis.\n",
    "\n",
    "In this tutorial we will use **Recursive** method. For better understanding the topic you can check https://langchain-text-splitter.streamlit.app/ app."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bba2f03-4d3a-4e37-8ab5-c7ae89df93fc",
   "metadata": {},
   "source": [
    "### Database Population\n",
    "Database population is a step where we load documents, embed them and than load into database."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5847ab8c-9959-49ee-8f6b-4125ff78be36",
   "metadata": {},
   "source": [
    "#### Data Loading\n",
    "For ease of use, we'll use helper funcitions from langchain. Note that langchain_community is also required. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "45cb0608-0162-4114-91e8-f3e93d1a4e0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from langchain.docstore.document import Document\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain_community.document_loaders import TextLoader\n",
    "from langchain.vectorstores.pgvector import PGVector\n",
    "from langchain_huggingface import HuggingFaceEndpointEmbeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7ab9197-0d1a-42fd-af24-4cdb6264b2c0",
   "metadata": {},
   "source": [
    "#### Loading Documents with embeddings\n",
    "Here we need to create huggingface TEI client and PGVector client. For PGVector, collection name corresponds to table name, within connection string there is connection protocol: `postgresql+psycopg2`, next is user, password, host, port and database name. For ease of use, pre_delete_collection is set to true to prevent duplicates in database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d591325c-ab5d-495c-a012-1a82fbe6b25b",
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = HuggingFaceEndpointEmbeddings(model=\"http://localhost:9002\", huggingfacehub_api_token=\"EMPTY\")\n",
    "store = PGVector(\n",
    "    collection_name=\"documents\",\n",
    "    connection_string=\"postgresql+psycopg2://postgres:postgres@localhost:9003/postgres\",\n",
    "    embedding_function=embeddings,\n",
    "    pre_delete_collection=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "faea35e8-ed40-4bba-a748-732d0a25efab",
   "metadata": {},
   "source": [
    "#### Data Loading and Splitting\n",
    "Data is loaded from text files from `data/`, than documents are splitted into chunks of size: 512 characters and finally loaded into database. Note that documents can have metadata, that can be also stored in vector database. \n",
    "\n",
    "You can load new text file in the `data/` folder to run the RAG pipeline on new content by running the follwoing cell again with new data.  This cell will create a new Database to run your query. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "cc4db026-a84c-4f3c-9b4a-0e6985e3b01e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading data/state_of_the_union.txt...\n",
      "Finished.\n"
     ]
    }
   ],
   "source": [
    "def load_file_to_db(path: str, store: PGVector):\n",
    "    loader = TextLoader(path)\n",
    "    document = loader.load()\n",
    "    text_splitter = CharacterTextSplitter(chunk_size=512, chunk_overlap=0)\n",
    "    for chunk in text_splitter.split_documents(document):\n",
    "        store.add_documents([chunk])\n",
    "\n",
    "for doc in Path(\"data/\").glob(\"*.txt\"):\n",
    "    print(f\"Loading {doc}...\")\n",
    "    load_file_to_db(str(doc), store)\n",
    "\n",
    "print(\"Finished.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b623d46-8c10-40a7-bacf-ee6c7d7511db",
   "metadata": {},
   "source": [
    "## 4. Running the Application\n",
    "To start the application run the following commands below to setup the Gradio Interface.  \n",
    "Load a text file in the `data` folder and the run the cell above and the application will ingest and start the chat application to ask question to the document.  \n",
    "You will see that it's directly accessing the TGI and TEI libraries to ingest, create the embeddings and vector database, the run the query thruough the database and then use the LLM to generate an answer to your query.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "4811b363-65c1-4fce-8ff1-23bfdbfa3cf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext gradio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "198187a7-6c65-4048-a8a9-10758dd6c776",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores.pgvector import PGVector\n",
    "from langchain_huggingface import HuggingFaceEndpointEmbeddings\n",
    "from text_generation import Client\n",
    "\n",
    "rag_prompt_intel_raw = \"\"\"### System: You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise. \n",
    "\n",
    "### User: Question: {question}\n",
    "\n",
    "Context: {context}\n",
    "\n",
    "### Assistant: \"\"\"\n",
    "\n",
    "def get_sources(question):\n",
    "    embeddings = HuggingFaceEndpointEmbeddings(model=\"http://localhost:9002\", huggingfacehub_api_token=\"EMPTY\")\n",
    "    store = PGVector(\n",
    "        collection_name=\"documents\",\n",
    "        connection_string=\"postgresql+psycopg2://postgres:postgres@localhost:9003/postgres\",\n",
    "        embedding_function=embeddings,\n",
    "    )\n",
    "    return store.similarity_search(f\"Represent this sentence for searching relevant passages: {question}\", k=2)\n",
    "\n",
    "def sources_to_str(sources):\n",
    "    return \"\\n\".join(f\"{i+1}. {s.page_content}\" for i, s in enumerate(sources))\n",
    "\n",
    "def get_answer(question, sources):\n",
    "    client = Client(\"http://localhost:9001\") #change this to 9009 for the new model\n",
    "    context = \"\\n\".join(s.page_content for s in sources)\n",
    "    prompt = rag_prompt_intel_raw.format(question=question, context=context)\n",
    "    # return client.generate_stream(prompt, max_new_tokens=1024, stop_sequences=[\"### User:\", \"</s>\"])\n",
    "    return client.generate(prompt, max_new_tokens=1024, stop_sequences=[\"### User:\", \"</s>\"]).generated_text\n",
    "\n",
    "default_question = \"What is this the summary of this document?\"\n",
    "\n",
    "def rag_answer(question, history):\n",
    "    question = question[\"text\"]\n",
    "    sources = get_sources(question)\n",
    "    answer = get_answer(question, sources)\n",
    "    #return f\"Sources:\\n{sources_to_str(sources)}\\nAnswer:\\n{answer}\"\n",
    "    #print(f\" inside rag_answer: Question : {question} ##### Answer : {answer}\")\n",
    "    return f\"{answer}\"\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "444fb610",
   "metadata": {},
   "source": [
    "This step below will launch the Gradio Interface:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc3a5611-3534-4784-b8b9-6139d4813172",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from gui.rag_gui import simple_chat_gui\n",
    "\n",
    "gradio_demo = simple_chat_gui(rag_answer)\n",
    "gradio_demo.launch(server_name=\"0.0.0.0\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a83e9fe",
   "metadata": {},
   "source": [
    "#### Shutdown the RAG Application and Docker images\n",
    "You must shut down the gaudi-tgi and gaudi-tei images so that the Intel Gaudi cards can be released to be used for other tasks or tutorials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "605bced5",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker stop gaudi-tgi && docker rm gaudi-tgi\n",
    "!docker stop gaudi-tei && docker rm gaudi-tei\n",
    "!docker stop postgres_vectordb && docker rm postgres_vectordb\n",
    "exit()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}