PyTorch/localGPT_inference/LocalGPT_Inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "646a18d1-97cd-4ffc-8394-3bff4e8bef6e",
   "metadata": {},
   "source": [
    "Copyright (c) 2024 Habana Labs, Ltd. an Intel Company.\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. 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": "4c8dbf1e-14f5-4b1a-a67a-0859d88bf942",
   "metadata": {},
   "source": [
    "# Using LocalGPT with Retrieval Augmented Generation (RAG) on the Intel&reg; Gaudi&reg; 2 AI accelerator with the Llama 2 70B model\n",
    "This tutorial will show how to use the [LocalGPT](https://github.com/PromtEngineer/localGPT) open source initiative on the Intel Gaudi 2 AI accelerator.  LocalGPT allows you to load your own documents and run an interactive chat session with this material using concepts from Retrieval Augmented Generation (RAG).  \n",
    "\n",
    "This allows you to query and summarize your content by loading any .pdf or .txt documents into the `SOURCE DOCUMENTS` folder, using utilities from the ingest.py script to tokenize your content and then the run_localGPT.py script to start the interaction.  \n",
    "\n",
    "The first section shows how RAG works and run thruough the steps of the indexing the local content, retrieval and text generation with a single question and response. \n",
    "\n",
    "The last section uses the full LocalGPT framework with the **meta-llama/Llama-2-70b-chat-hf** model as the reference model that will manage the inference on Gaudi 2.  DeepSpeed inference is used based on the size of the model.\n",
    "\n",
    "To optimize this instantiation of LocalGPT, we have created new content on top of the existing Hugging Face based \"text-generation\" inference task and pipelines, including:\n",
    "\n",
    "1. Using the Hugging Face Optimum Habana Library with the Llama 2 70B model, which is optimized on Gaudi2. \n",
    "2. Using LangChain to import the source document with an embedding model, using the Hugging Face Optimum Habana Library.\n",
    "3. We are using a custom pipeline class, `GaudiTextGenerationPipeline` that optimizes text-generation tasks for padding and indexing for static shapes, to improve performance.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1eab16b-6912-4676-af1f-91b7472c90d4",
   "metadata": {},
   "source": [
    "##### Install DeepSpeed to run inference on the full Llama 2 70B model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1d7fd64-18b0-4a71-8484-813defa94062",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install git+https://github.com/HabanaAI/DeepSpeed.git@1.19.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ad1b812-5032-4a26-867e-155575a14991",
   "metadata": {},
   "source": [
    "##### Go to the LocalGPT folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c604050-d6f5-47a9-adc7-0b9e3c0fe34a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd /root/Gaudi-tutorials/PyTorch/localGPT_inference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "188e2383-deaf-4349-90e5-31bea9a507b0",
   "metadata": {},
   "source": [
    "##### Install the requirements for LocalGPT "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deadcf84-eb2b-4f04-9cbf-4e6d67d4e999",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q --upgrade pip\n",
    "!pip install -q -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "304d111e-2118-4bb0-b62d-214e895479a2",
   "metadata": {},
   "source": [
    "##### Install the Optimum Habana Library from Hugging Face"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "217c2833-0bb7-481e-907d-7b28d88bbf7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q optimum-habana==1.15.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "095322a0",
   "metadata": {},
   "source": [
    "## Retrieval-Augmented Generation (RAG)\n",
    "LocalGPT uses Retrieval-Augmented Generation (RAG) at it's core. RAG is a relatively new AI technique that combines an information retrieval system with text-generation models/LLMs. It provides an effective way to ground LLMs by using retrieved contexts from an external knowledge base, without having to perform retraining or finetuning.\n",
    "The LocalGPT application workflow can be broken down as follows:\n",
    "* Document Ingestion: This step involves creating an external knowledge base via a vector database. The text present in the documents is parsed, split into chunks and converted to embeddings using an embedding model. The vector embeddings are finally stored in the vector database.\n",
    "\n",
    "![](img/ingest.jpg)\n",
    "\n",
    "* Text Generation: This step involves accepting a query from the user, converting the query to embeddings and retrieving appropriate contexts from the knowledge base. The input prompt to the LLM is the concatenation of the query, contexts and chat history.\n",
    "\n",
    "![](img/documentqa.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc2eaba6-2d4b-4007-bb13-ca5113022aad",
   "metadata": {},
   "source": [
    "### Document Ingestion\n",
    "Copy all of your files into the `SOURCE_DOCUMENTS` directory\n",
    "\n",
    "The current default file types are .txt, .pdf, .csv, and .xlsx, if you want to use any other file type, you will need to convert it to one of the default file types.\n",
    "\n",
    "Run the following cells to ingest all the data. This notebook uses LangChain tools to parse the documents and create embeddings locally using the HuggingFace Optimum Habana Library. It then stores the result in a local vector database (DB) using Chroma vector store. \n",
    "\n",
    "If you want to start from an empty database, delete the DB folder and run the next few cells again. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e68ae917-86df-4061-9f9a-d6cb2bb7cd75",
   "metadata": {},
   "source": [
    "##### Load your files as LangChain Documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7db570f4-acce-4e73-8ac1-7742e621c6b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from constants import SOURCE_DIRECTORY\n",
    "from ingest import load_documents\n",
    "\n",
    "documents = load_documents(SOURCE_DIRECTORY)\n",
    "print(f\"Loaded {len(documents)} documents from {SOURCE_DIRECTORY}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b730b93b-eb2c-45a9-9989-203b7de37be4",
   "metadata": {},
   "source": [
    "##### Split the text into chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a90c469-dcae-49aa-8c80-a3043072c4ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)\n",
    "texts = text_splitter.split_documents(documents)\n",
    "print(f\"Created {len(texts)} chunks of text\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c2ab0c9-f379-4870-a36e-6f221a6520a1",
   "metadata": {},
   "source": [
    "##### Create embeddings from chunks of text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63ebe1ee-f066-4083-aa53-eddb73996be5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from constants import EMBEDDING_MODEL_NAME\n",
    "from langchain_huggingface import HuggingFaceEmbeddings\n",
    "\n",
    "from habana_frameworks.torch.utils.library_loader import load_habana_module\n",
    "from optimum.habana.sentence_transformers.modeling_utils import adapt_sentence_transformers_to_gaudi\n",
    "\n",
    "load_habana_module()\n",
    "\n",
    "adapt_sentence_transformers_to_gaudi()\n",
    "embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME, model_kwargs={\"device\": \"hpu\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9079b506-0fcf-49ca-9865-0ebfa8170df3",
   "metadata": {},
   "source": [
    "##### Create a Chroma vector database to store embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68b30209-a897-40b5-aaab-fe0c33585c28",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from constants import PERSIST_DIRECTORY, CHROMA_SETTINGS\n",
    "from langchain_chroma import Chroma\n",
    "\n",
    "start_time = time.perf_counter()\n",
    "db = Chroma.from_documents(texts, embeddings, persist_directory=PERSIST_DIRECTORY, client_settings=CHROMA_SETTINGS)\n",
    "end_time = time.perf_counter()\n",
    "print(f\"Time taken to create vector store: {(end_time-start_time)*1000} ms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc45867b-da73-4fdf-b821-ac932ca2c1d7",
   "metadata": {},
   "source": [
    "### How to access and Use the Llama 2 model\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.\n",
    "\n",
    "To be able to run gated models like this Llama-2-70b-chat-hf, you need the following:\n",
    "\n",
    "* Have a HuggingFace account\n",
    "* Agree to the terms of use of the model in its model card on the HF Hub\n",
    "* Set a read token\n",
    "* Login to your account using the HF CLI: run huggingface-cli login before launching your script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ffeec3a-1671-4cce-940c-3fe962f0b9c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!huggingface-cli login --token <your token here>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01705607",
   "metadata": {},
   "source": [
    "### Text Generation\n",
    "Once the Chroma vector database is ready, we can explore the text-generation component of LocalGPT.\n",
    "\n",
    "The next few cells describe all the steps in the text generation process. We use the smallest Llama 2 model **meta-llama/Llama-2-7b-chat-hf** to perform augmented text-generation after retrieving relevant contexts from the vector database."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2d90d69",
   "metadata": {},
   "source": [
    "##### Load the LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34c8060e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from run_localGPT import load_model\n",
    "\n",
    "model_id = \"meta-llama/Llama-2-7b-chat-hf\"\n",
    "llm, _ = load_model(device_type=\"hpu\", model_id=model_id, temperature=0.2, top_p=0.95, model_basename=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d001a99",
   "metadata": {},
   "source": [
    "##### Define the Retriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b3ff0aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "db = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embeddings, client_settings=CHROMA_SETTINGS)\n",
    "retriever = db.as_retriever()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a3b188c",
   "metadata": {},
   "source": [
    "##### Create the prompt template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a5892a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import PromptTemplate\n",
    "\n",
    "template = \"\"\"Use the following pieces of context to answer the question at the end. If you don't know the answer,\\\n",
    "just say that you don't know, don't try to make up an answer.\n",
    "\n",
    "{context}\n",
    "\n",
    "Question: {question}\n",
    "Answer:\"\"\"\n",
    "\n",
    "prompt = PromptTemplate(input_variables=[\"context\", \"question\"], template=template)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba70acaf",
   "metadata": {},
   "source": [
    "##### Initialize a LangChain object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97272db2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import create_retrieval_chain\n",
    "from langchain.chains.combine_documents import create_stuff_documents_chain\n",
    "\n",
    "qa = create_retrieval_chain(retriever, create_stuff_documents_chain(llm, prompt))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fb87269",
   "metadata": {},
   "source": [
    "##### Ask a question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb908353",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"What is this document about?\"\n",
    "res = qa.invoke({\"question\": query, \"input\": query})\n",
    "print(res[\"answer\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "510a6e92",
   "metadata": {},
   "source": [
    "##### Clean up before running Full LocalGPT below\n",
    "To run the full Local LocalGPT model below, you need to restart the Kernel in the Jupyter Server to ensure that all the Intel Gaudi Accelerators are released.  This can be accomplished by selecting this option in the Kernel menu or the `exit()` command at the bottom of this notebook. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c3ac741b-cb79-4fab-a1aa-efa075503d61",
   "metadata": {},
   "source": [
    "### Running the LocalGPT full example with Llama 2 70B Chat \n",
    "\n",
    "### Set the model Usage\n",
    "\n",
    "To change the model, you can modify the \"LLM_ID = <add model here>\" in the `constants.py` file. For this example, the default is `meta-llama/Llama-2-70b-chat-hf`.  \n",
    "\n",
    "Since this is interactive, it's a better experince to launch this from a terminal window.  This run_localGPT.py script uses a local LLM (Llama 2 in this case) to understand questions and create answers. The context for the answers is extracted from the local vector store using a similarity search to locate the right piece of context from the documentation.  This is the run command to use:\n",
    "\n",
    "`PT_HPU_LAZY_ACC_PAR_MODE=1 PT_HPU_ENABLE_LAZY_COLLECTIVES=true python gaudi_spawn.py --use_deepspeed --world_size 8 run_localGPT.py --device_type hpu --temperature 0.7 --top_p 0.95`\n",
    "\n",
    "Running the full 70B model takes up ~128GB of disk space, so if your system is storage constrained, it may be best to run the Llama 2 7B or 13B chat models.  Change the LLM_ID variable in the `constants.py` file (example: `LLM_ID = \"meta-llama/Llama-2-7b-chat-hf\"`) and use the command below.\n",
    "`python run_localGPT.py --device_type hpu --temperature 0.7 --top_p 0.95`\n",
    "\n",
    "Note: The inference is running sampling mode, so the user can optinally modify the temperature and top_p settings.  The current settings are temperature=0.7, top_p=0.95.  Type \"exit\" at the prompt to stop the execution.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9ee22fd-2bca-49b8-8f8d-501075058b40",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Run this command in a terminal window to start the interactive chat: `PT_HPU_LAZY_ACC_PAR_MODE=1 PT_HPU_ENABLE_LAZY_COLLECTIVES=true python gaudi_spawn.py --use_deepspeed --world_size 8 run_localGPT.py --device_type hpu --temperature 0.7 --top_p 0.95`, the example below is showing the initial output:   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc237333-fbdc-440e-8249-89639af4c6c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!PT_HPU_LAZY_ACC_PAR_MODE=1 PT_HPU_ENABLE_LAZY_COLLECTIVES=true python gaudi_spawn.py --use_deepspeed --world_size 8 run_localGPT.py --device_type hpu --temperature 0.2 --top_p 0.95"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adfe1723-04da-4d59-a299-fe0431a84bc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "exit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ba45e02-5a1f-4dd0-863e-5bdf8e2cbca0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}