10_rag_pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Building a Retrieval-Augmented Generation (RAG) System using Pytorch\n",
    "\n",
    "## Learning Objectives\n",
    "### Remember\n",
    "- Define the key components of a RAG system\n",
    "- Identify the essential libraries required for building a RAG pipeline\n",
    "- List the main steps in the RAG process\n",
    "\n",
    "### Understand\n",
    "- Explain how document embedding works\n",
    "- Describe the role of vector databases in RAG\n",
    "- Interpret the relationship between chunks, embeddings, and retrieval\n",
    "\n",
    "### Apply\n",
    "- Implement a document loading and chunking pipeline\n",
    "- Configure a language model for generation\n",
    "- Set up a vector store for document retrieval\n",
    "\n",
    "### Analyze\n",
    "- Compare different chunking strategies\n",
    "- Examine the impact of various parameter settings\n",
    "- Debug common issues in RAG systems\n",
    "\n",
    "### Evaluate\n",
    "- Assess the quality of generated responses\n",
    "- Test system performance with different configurations\n",
    "- Judge the effectiveness of retrieval strategies\n",
    "\n",
    "### Create\n",
    "- Design a complete RAG pipeline\n",
    "- Develop custom prompt templates\n",
    "- Build an interactive question-answering system\n",
    "\n",
    "## Prerequisites\n",
    "- Basic understanding of Python\n",
    "- Familiarity with machine learning concepts\n",
    "- Understanding of basic NLP terminology\n",
    "\n",
    "Let's begin our journey into building a RAG system!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# What is RAG (Retrieval-augmented generation)\n",
    "\n",
    "**Retrieval-augmented generation (RAG)** is a technique for augmenting LLM knowledge with additional, often private or real-time, data. LLMs can reason about wide-ranging topics, but their knowledge is limited to the public data up to a specific point in time that they were trained on. If you want to build AI applications that can reason about private data or data introduced after a model’s cutoff date, you need to augment the knowledge of the model with the specific information it needs. The process of bringing the appropriate information and inserting it into the model prompt is known as Retrieval Augmented Generation (RAG)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -r rag/requirements.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run QA over Document\n",
    "\n",
    "Now, when model created, we can setup Chatbot\n",
    "\n",
    "A typical RAG application has two main components:\n",
    "\n",
    "- **Indexing**: a pipeline for ingesting data from a source and indexing it. This usually happen offline.\n",
    "\n",
    "- **Retrieval and generation**: the actual RAG chain, which takes the user query at run time and retrieves the relevant data from the index, then passes that to the model.\n",
    "\n",
    "The most common full sequence from raw data to answer looks like:\n",
    "\n",
    "**Indexing**\n",
    "\n",
    "1. `Load`: First we need to load our data. We’ll use DocumentLoaders for this.\n",
    "2. `Split`: Text splitters break large Documents into smaller chunks. This is useful both for indexing data and for passing it in to a model, since large chunks are harder to search over and won’t in a model’s finite context window.\n",
    "3. `Store`: We need somewhere to store and index our splits, so that they can later be searched over. This is often done using a VectorStore and Embeddings model.\n",
    "\n",
    "![Indexing pipeline](https://github.com/openvinotoolkit/openvino_notebooks/assets/91237924/dfed2ba3-0c3a-4e0e-a2a7-01638730486a)\n",
    "\n",
    "**Retrieval and generation**\n",
    "\n",
    "1. `Retrieve`: Given a user input, relevant splits are retrieved from storage using a Retriever.\n",
    "2. `Generate`: A LLM produces an answer using a prompt that includes the question and the retrieved data.\n",
    "\n",
    "![Retrieval and generation pipeline](https://github.com/openvinotoolkit/openvino_notebooks/assets/91237924/f0545ddc-c0cd-4569-8c86-9879fdab105a)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setting Up the Environment\n",
    "\n",
    "First, let's import all necessary libraries. We'll go through each import and understand its role in our RAG system.\n",
    "\n",
    "Key Components:\n",
    "- LangChain: Framework for developing applications powered by language models\n",
    "- HuggingFace: Platform for accessing pre-trained models and embeddings\n",
    "- Chroma: Vector store for efficient similarity search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import chains, text_splitter, PromptTemplate\n",
    "from langchain_community.embeddings.fastembed import FastEmbedEmbeddings\n",
    "from langchain_community import document_loaders, embeddings, vectorstores\n",
    "from langchain.llms.base import LLM\n",
    "from typing import Any, List, Optional, Dict\n",
    "from pydantic import Field\n",
    "import torch\n",
    "import shutil\n",
    "import intel_extension_for_pytorch as ipex\n",
    "import time\n",
    "import os\n",
    "import warnings\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import display, clear_output\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "VECTOR_DB_DIR = \"vector_dbs\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LoadLLM(LLM):\n",
    "    model_path: str = Field(default=\"meta-llama/Llama-2-7b-chat-hf\")\n",
    "    model: Optional[Any] = Field(default=None, exclude=True)\n",
    "    tokenizer: Optional[Any] = Field(default=None, exclude=True)\n",
    "\n",
    "    class Config:\n",
    "        arbitrary_types_allowed = True\n",
    "\n",
    "    def __init__(self, model_path: str = \"meta-llama/Llama-2-7b-chat-hf\", **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.model_path = model_path\n",
    "        self.model = AutoModelForCausalLM.from_pretrained(\n",
    "            self.model_path,\n",
    "            trust_remote_code=True,\n",
    "            torch_dtype=torch.float16,\n",
    "            low_cpu_mem_usage=True\n",
    "        )\n",
    "        self.model = self.model.to('xpu')\n",
    "        self.model = ipex.optimize(self.model)\n",
    "        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)\n",
    "\n",
    "    @property\n",
    "    def _llm_type(self) -> str:\n",
    "        return \"intel_llm\"\n",
    "\n",
    "    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:\n",
    "        with torch.inference_mode():\n",
    "            input_ids = self.tokenizer(prompt, return_tensors=\"pt\").to('xpu')\n",
    "            output = self.model.generate(\n",
    "                **input_ids,\n",
    "                do_sample=True,\n",
    "                max_new_tokens=500\n",
    "            )\n",
    "            torch.xpu.synchronize()\n",
    "            return self.tokenizer.decode(output[0], skip_special_tokens=True)\n",
    "\n",
    "    @property\n",
    "    def _identifying_params(self) -> Dict[str, Any]:\n",
    "        return {\"model_path\": self.model_path}\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Document Loading and Processing\n",
    "\n",
    "A crucial part of any RAG system is how it handles document loading and processing. We'll create functions to:\n",
    "1. Load documents from URLs\n",
    "2. Split documents into manageable chunks\n",
    "3. Initialize our embedding model\n",
    "\n",
    "Understanding these processes is crucial as they directly impact the quality of our retrieval system.\n",
    "\n",
    "* Document loaders in RAG are used to load and preprocess the documents that will be used for retrieval during the question answering process.\n",
    "* Document loaders are responsible for preprocessing the documents. This includes tokenizing the text, converting it to the format expected by the retriever, and creating batches of documents.\n",
    "* Document loaders work in conjunction with the retriever in RAG. The retriever uses the documents loaded by the document loader to find the most relevant documents for a given query.\n",
    "* The WebBaseLoader in Retrieval Augmented Generation (RAG) is a type of document loader that is designed to load documents from the web.\n",
    "* The WebBaseLoader is used when the documents for retrieval are not stored locally or in a Hugging Face dataset, but are instead located on the web. This can be useful when you want to use the most up-to-date information available on the internet for your question answering system"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Text splitter\n",
    "\n",
    "* RecursiveCharacterTextSplitter is used to split text into smaller pieces recursively at the character level. \n",
    "* split_documents fuctions splits larger documents into smaller chunks, for easier processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_document(url):\n",
    "    print(\"Loading document from URL...\")\n",
    "    loader = document_loaders.WebBaseLoader(url)\n",
    "    return loader.load()\n",
    "\n",
    "def split_document(text, chunk_size=3000, overlap=200):\n",
    "    print(\"Splitting document into chunks...\")\n",
    "    text_splitter_instance = text_splitter.RecursiveCharacterTextSplitter(\n",
    "        chunk_size=chunk_size, chunk_overlap=overlap)\n",
    "    return text_splitter_instance.split_documents(text)\n",
    "\n",
    "def initialize_embedding_fn(embedding_type=\"huggingface\"):\n",
    "    print(f\"Initializing {embedding_type} embeddings...\")\n",
    "    if embedding_type == \"huggingface\":\n",
    "        return embeddings.HuggingFaceEmbeddings(\n",
    "            model_name=\"sentence-transformers/paraphrase-MiniLM-L3-v2\")\n",
    "    elif embedding_type == \"fastembed\":\n",
    "        return FastEmbedEmbeddings(threads=16)\n",
    "    else:\n",
    "        raise ValueError(f\"Unsupported embedding type: {embedding_type}\")\n",
    "\n",
    "def clear_vector_store(persist_dir=VECTOR_DB_DIR):\n",
    "    import shutil\n",
    "    if os.path.exists(persist_dir):\n",
    "        try:\n",
    "            # Create a temporary Chroma instance to properly close any open connections\n",
    "            temp_db = vectorstores.Chroma(\n",
    "                persist_directory=persist_dir,\n",
    "                embedding_function=FastEmbedEmbeddings()\n",
    "            )\n",
    "            temp_db.persist()\n",
    "            del temp_db  # Explicitly delete the instance\n",
    "            \n",
    "            # Add a small delay to ensure connections are closed\n",
    "            time.sleep(1)\n",
    "            \n",
    "            shutil.rmtree(persist_dir)\n",
    "        except Exception as e:\n",
    "            print(f\"Warning: Error while clearing vector store: {e}\")\n",
    "            # If regular deletion fails, try forcing it\n",
    "            os.system(f\"rm -rf {persist_dir}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Vector Store Management\n",
    "\n",
    "The vector store is a crucial component in RAG systems. It enables efficient similarity search over document embeddings, allowing us to retrieve relevant context for any given query.\n",
    "\n",
    "Key Concepts:\n",
    "- Vector databases\n",
    "- Persistence and caching\n",
    "- Similarity search\n",
    "- Document retrieval\n",
    "- In Retrieval Augmented Generation (RAG) embeddings play a crucial role in the retrieval of relevant documents for a given query.\n",
    "\n",
    "* In RAG, each document in the knowledge base is represented as a dense vector, also known as an embedding. These embeddings are typically generated by a transformer model.\n",
    "* When a query is received, it is also converted into an embedding using the same transformer model. This ensures that the query and the documents are in the same vector space, making it possible to compare them.\n",
    "* Retrieval: The retrieval step in RAG involves finding the documents whose embeddings are most similar to the query embedding. This is typically done using a nearest neighbor search.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_or_create_embeddings(document_url, embedding_fn, persist_dir=VECTOR_DB_DIR):\n",
    "    # Create a unique directory name based on timestamp\n",
    "    import datetime\n",
    "    timestamp = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
    "    unique_persist_dir = os.path.join(persist_dir, timestamp)\n",
    "    \n",
    "    os.makedirs(unique_persist_dir, exist_ok=True)\n",
    "    \n",
    "    start_time = time.time()\n",
    "    print(\"Creating new vector store...\")\n",
    "    document = load_document(document_url)\n",
    "    documents = split_document(document)\n",
    "    \n",
    "    vector_store = vectorstores.Chroma.from_documents(\n",
    "        documents=documents,\n",
    "        embedding=embedding_fn,\n",
    "        persist_directory=unique_persist_dir,\n",
    "        collection_metadata={\"hnsw:space\": \"cosine\"}\n",
    "    )\n",
    "    vector_store.persist()\n",
    "    print(f\"Embedding time: {time.time() - start_time:.2f} seconds\")\n",
    "    return vector_store, unique_persist_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Question-Answering Interface\n",
    "\n",
    "The question-answering interface provides the user interaction layer of our RAG system. It handles:\n",
    "1. User input processing\n",
    "2. Response generation\n",
    "3. Error handling\n",
    "4. Performance monitoring"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def handle_user_interaction(vector_store, llm, question):\n",
    "    prompt_template = \"\"\"\n",
    "    Use the following pieces of context to answer the question at the end.\n",
    "    If you do not know the answer, answer 'I don't know', limit your response to the answer and nothing more.\n",
    "\n",
    "    {context}\n",
    "\n",
    "    Question: {question}\n",
    "    \"\"\"\n",
    "    prompt = PromptTemplate(\n",
    "        template=prompt_template,\n",
    "        input_variables=[\"context\", \"question\"])\n",
    "    \n",
    "    chain_type_kwargs = {\"prompt\": prompt}\n",
    "    retriever = vector_store.as_retriever(search_kwargs={\"k\": 4})\n",
    "    \n",
    "    qachain = chains.RetrievalQA.from_chain_type(\n",
    "        llm=llm,\n",
    "        retriever=retriever,\n",
    "        chain_type=\"stuff\",\n",
    "        chain_type_kwargs=chain_type_kwargs,\n",
    "        return_source_documents=False\n",
    "    )\n",
    "\n",
    "    start_time = time.time()\n",
    "    response = qachain.invoke({\"query\": question})\n",
    "    # The response is a dict with 'result' key containing the actual answer\n",
    "    return response['result'] if isinstance(response, dict) and 'result' in response else response\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Main Application Logic\n",
    "\n",
    "The main function orchestrates all components of our RAG system. It:\n",
    "1. Initializes components\n",
    "2. Sets up the RAG pipeline\n",
    "3. Manages the interaction loop\n",
    "4. Handles errors and cleanup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retrievers\n",
    "\n",
    "* Retrievers are responsible for fetching relevant documents from a document store or knowledge base given a query. The retrieved documents are then used by the generator to produce a response.\n",
    "* RetrievalQA is a type of question answering system that uses a retriever to fetch relevant documents given a question, and then uses a reader to extract the answer from the retrieved documents.\n",
    "* RetrievalQA can be seen as a two-step process:\n",
    "    * Retrieval: The retriever fetches relevant documents from the document store given a query.    \n",
    "    * Generation: The generator uses the retrieved documents to generate a response.\n",
    "* This two-step process allows RAG to leverage the strengths of both retrieval-based and generation-based approaches to question answering. The retriever allows RAG to efficiently search a large document store, while the generator allows RAG to generate detailed and coherent responses.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#from langchain_community.vectorstores.chroma import Chroma\n",
    "def main(document_url, question, embedding_type=\"huggingface\", model_path=\"meta-llama/Llama-2-7b-chat-hf\"):\n",
    "    try:\n",
    "        llm = LoadLLM(model_path=model_path)\n",
    "        embedding_fn = initialize_embedding_fn(embedding_type)\n",
    "        vector_store, persist_dir = get_or_create_embeddings(document_url, embedding_fn)\n",
    "        answer = handle_user_interaction(vector_store, llm, question)\n",
    "        vector_store.persist()\n",
    "        del vector_store\n",
    "        # Cleanup old directories if needed\n",
    "        cleanup_old_directories(VECTOR_DB_DIR, keep_last=5)\n",
    "        return answer\n",
    "    except Exception as e:\n",
    "        print(f\"An error occurred: {e}\")\n",
    "        return None\n",
    "\n",
    "def cleanup_old_directories(base_dir, keep_last=5):\n",
    "    \"\"\"Clean up old vector store directories, keeping only the most recent ones.\"\"\"\n",
    "    try:\n",
    "        dirs = [os.path.join(base_dir, d) for d in os.listdir(base_dir)]\n",
    "        dirs = [d for d in dirs if os.path.isdir(d)]\n",
    "        dirs.sort(key=lambda x: os.path.getctime(x), reverse=True)\n",
    "        \n",
    "        # Remove all but the last keep_last directories\n",
    "        for dir_path in dirs[keep_last:]:\n",
    "            try:\n",
    "                shutil.rmtree(dir_path)\n",
    "            except Exception as e:\n",
    "                print(f\"Warning: Could not remove directory {dir_path}: {e}\")\n",
    "    except Exception as e:\n",
    "        print(f\"Warning: Error during cleanup: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run QA with Ipywidgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create input widgets\n",
    "url_input = widgets.Text(\n",
    "    value='',\n",
    "    placeholder='Enter URL here',\n",
    "    description='URL:',\n",
    "    style={'description_width': 'initial'},\n",
    "    layout=widgets.Layout(width='80%')\n",
    ")\n",
    "\n",
    "question_input = widgets.Text(\n",
    "    value='',\n",
    "    placeholder='Enter your question here',\n",
    "    description='Question:',\n",
    "    style={'description_width': 'initial'},\n",
    "    layout=widgets.Layout(width='80%')\n",
    ")\n",
    "\n",
    "submit_button = widgets.Button(\n",
    "    description='Get Answer',\n",
    "    button_style='primary',\n",
    "    layout=widgets.Layout(width='20%')\n",
    ")\n",
    "\n",
    "output_area = widgets.Output()\n",
    "\n",
    "def on_submit_button_clicked(b):\n",
    "    with output_area:\n",
    "        clear_output()\n",
    "        if not url_input.value or not question_input.value:\n",
    "            print(\"Please enter both URL and question.\")\n",
    "            return\n",
    "        \n",
    "        response = main(url_input.value, question_input.value)\n",
    "        print(response)  # Now this will be just the answer text\n",
    "\n",
    "submit_button.on_click(on_submit_button_clicked)\n",
    "\n",
    "# Display the widgets\n",
    "display(url_input)\n",
    "display(question_input)\n",
    "display(submit_button)\n",
    "display(output_area)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Workshop Summary\n",
    "\n",
    "In this workshop, we've built a complete RAG system from scratch, covering:\n",
    "\n",
    "### Technical Components\n",
    "1. Document Processing\n",
    "   - Web document loading\n",
    "   - Text chunking strategies\n",
    "   - Embedding generation\n",
    "\n",
    "2. Vector Store Management\n",
    "   - Persistent storage\n",
    "   - Efficient retrieval\n",
    "   - Similarity search\n",
    "\n",
    "3. Language Model Integration\n",
    "   - Model configuration\n",
    "   - Parameter optimization\n",
    "   - Response generation\n",
    "\n",
    "### Key Learnings\n",
    "1. System Architecture\n",
    "   - Understanding RAG pipeline components\n",
    "   - Component interaction\n",
    "   - Error handling\n",
    "\n",
    "2. Performance Optimization\n",
    "   - Memory management\n",
    "   - GPU utilization\n",
    "   - Response time optimization\n",
    "\n",
    "3. Best Practices\n",
    "   - Code organization\n",
    "   - Documentation\n",
    "   - Error handling\n",
    "\n",
    "### Next Steps\n",
    "To further improve the system, consider:\n",
    "1. Implementing different embedding models\n",
    "2. Experimenting with chunk sizes and overlap\n",
    "3. Adding evaluation metrics\n",
    "4. Implementing caching mechanisms\n",
    "5. Adding support for different document types\n",
    "\n",
    "This workshop provides a foundation for building and understanding RAG systems, which you can extend and customize for your specific use cases."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "PyTorch GPU",
   "language": "python",
   "name": "pytorch-gpu"
  },
  "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}