{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7754c133",
   "metadata": {},
   "source": [
    "# Retrieval Augmented Generation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af3dc7d3-e93e-41c1-b373-2f3bbec76616",
   "metadata": {},
   "source": [
    "\n",
    "![rag_2025.png](../rag_2025.png)\n",
    "\n",
    "Our project is the result of team work. We have two sets of\n",
    "documentation pages. Thanks to our [team member for illustrating the\n",
    "process](https://uio-library.github.io/LLM-course/4_RAG.html) .\n",
    "\n",
    "Code view 1:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e4aa497-515e-43a5-93c7-0a6d65ba2407",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['HF_HOME'] = '/fp/projects01/ec443/huggingface/cache/'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "605b7f65-569c-4ead-ab60-5cb612bb9c8d",
   "metadata": {},
   "source": [
    "## The Model\n",
    "\n",
    "First, we check if we have GPU available.\n",
    "\n",
    "code view 3:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35eb4825",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "device = 0 if torch.cuda.is_available() else -1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fb78019",
   "metadata": {},
   "source": [
    "\n",
    "The pipeline from huggingface.\n",
    "code view 4:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4ebed51-7e6a-4aae-8fd8-77e6460bb68a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_huggingface.llms import HuggingFacePipeline\n",
    "\n",
    "llm = HuggingFacePipeline.from_model_id(\n",
    "    model_id='mistralai/Mistral-7B-Instruct-v0.3',\n",
    "    # model_id='meta-llama/Llama-3.2-3B-Instruct',\n",
    "    task='text-generation',\n",
    "    device=device # device = 0\n",
    "    pipeline_kwargs={\n",
    "        'max_new_tokens': 500,\n",
    "        'do_sample': True,\n",
    "        'temperature': 0.3,\n",
    "        'num_beams': 4\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5128a5d7-75e7-4ceb-bb3d-e6b5ca956de6",
   "metadata": {},
   "source": [
    "## Using the model\n",
    "\n",
    "First, we are going to use the language model without connecting it to\n",
    "the documents. The output will be what the model knows from before, on\n",
    "the subject we ask.\n",
    "\n",
    "code view 4:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3298ffa-a43c-46c3-aa83-c09dd6a803c3",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "outputs": [],
   "source": [
    "# Use with models from Meta\n",
    "#llm.pipeline.tokenizer.pad_token_id = llm.pipeline.tokenizer.eos_token_id\n",
    "\n",
    "query = 'What are the major contributions of the Trivandrum Observatory?'\n",
    "output = llm.invoke(query)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77709c8d-6499-414f-a681-86526c715749",
   "metadata": {},
   "source": [
    "## The Vectorizer\n",
    "\n",
    "Text must be vectorized before it can be processed. Our HuggingFace\n",
    "pipeline will do that automatically for the large language model. But we\n",
    "must make a vectorizer for the search index for our documents database.\n",
    "We use a vectorizer called a word embedding model from HuggingFace.\n",
    "Again, the HuggingFace library will automatically download the model.\n",
    "You can read more on the [Embedding model\n",
    "here](https://huggingface.co/BAAI/bge-m3) . The developer writes this in\n",
    "the readme file: \"BGE M3 is an embedding model supporting dense\n",
    "retrieval, lexical matching and multi-vector interaction\".\n",
    "\n",
    "code view 5:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70a867b4-0a48-4051-b956-048f08b3bc66",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_huggingface import HuggingFaceEmbeddings\n",
    "\n",
    "huggingface_embeddings = HuggingFaceEmbeddings(\n",
    "    model_name='BAAI/bge-m3',\n",
    "    model_kwargs = {'device': 'cuda:0'},\n",
    "    #or: model_kwargs={'device':'cpu'},\n",
    "    encode_kwargs={'normalize_embeddings': True}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5244990-d70b-4286-b9f6-0cb2bf346bb9",
   "metadata": {},
   "source": [
    "\n",
    "## The arguments of the embeddings\n",
    "\n",
    "These are the arguments to the embedding model:\n",
    "\n",
    "    ‘model_name’: the name of the model on HuggingFace\n",
    "\n",
    "    ‘device’: the hardware device to use, either a GPU or CPU\n",
    "\n",
    "    ‘normalize_embeddings’: embeddings can have different magnitudes. Normalizing the embeddings makes their magnitudes equal.\n",
    "\n",
    "## Document location\n",
    "\n",
    "Code view 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "459c5c36",
   "metadata": {},
   "outputs": [],
   "source": [
    "document_folder = '/fp/projects01/ec443/documents'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b482de3",
   "metadata": {},
   "source": [
    "Code view 7:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "687022a6-3e10-490c-af24-6f27cc56e04f",
   "metadata": {},
   "source": [
    "## Loading the documents\n",
    "\n",
    "code view 8:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae87b326-0693-4441-bb31-ceba88cd0427",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import DirectoryLoader\n",
    "\n",
    "loader = DirectoryLoader(document_folder)\n",
    "documents = loader.load()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d5c66fe-9197-46b2-a423-cff33efc372d",
   "metadata": {},
   "source": [
    "\n",
    "code view 9:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6a21d0b-99ba-485a-b45a-26c7c8a38185",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'Number of documents:', len(documents))\n",
    "print('Maximum document length: ', max([len(doc.page_content) for doc in documents]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d4b0199-f9b2-465c-b838-dc6d51ecc0f0",
   "metadata": {},
   "source": [
    "The document loader loads each file as a separate document. We can check how long our documents are. For example, we can use the function max() to find the length of the longest document.\n",
    "\n",
    "We can examine one of the documents\n",
    "\n",
    "code view 10:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad016219-dcc8-48d3-86d6-b63438614399",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(documents[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10162757-1717-4835-b0e9-4af6a7ef1fe5",
   "metadata": {},
   "source": [
    "## Splitting the documents\n",
    "\n",
    "code view 11:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36aa90ae-1204-4578-8b69-abe8bdbbea2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size = 700, #  Could be more, for larger models like mistralai/Ministral-8B-Instruct-2410\n",
    "    chunk_overlap  = 200,\n",
    ")\n",
    "documents = text_splitter.split_documents(documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a61a122-5017-45cc-9457-165d861847c5",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "We can check if the maximum document length has changed:\n",
    "\n",
    "code view 10:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0357081-cc4f-4610-ab49-42cee8f93d2d",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "outputs": [],
   "source": [
    "print(f'Number of documents:', len(documents))\n",
    "print('Maximum document length: ', max([len(doc.page_content) for doc in documents]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b2d8ec6-67ea-48c0-a545-f4dd365a83c3",
   "metadata": {},
   "source": [
    "## The Document Index\n",
    "\n",
    "code view 11:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1813dd89-f767-4800-a2b6-0e1d7462f676",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "outputs": [],
   "source": [
    "from langchain_community.vectorstores import FAISS\n",
    "vectorstore = FAISS.from_documents(documents, huggingface_embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cff13375-5f84-4c92-b8da-ee6060f65aa1",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "code view 12:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3ec001a-4281-429a-9dcb-fec123e4e746",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "outputs": [],
   "source": [
    "relevant_documents = vectorstore.similarity_search(query)\n",
    "print(f'Number of documents found: {len(relevant_documents)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7071b74-43b5-4c32-b2ab-196cd9a87fdf",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "code view 13:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8e1f22d-5490-4d2e-b83b-2ce460c89cc9",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "outputs": [],
   "source": [
    "print(relevant_documents[0].page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3faa0eac-9c7c-426b-b7c3-4b4d9da585cf",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "For our RAG application we need to access the search engine through an\n",
    "interface called a retriever:\n",
    "\n",
    "code view 14:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d923004-e2b9-4062-9b76-0bb50ed2aef1",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "outputs": [],
   "source": [
    "retriever = vectorstore.as_retriever(search_kwargs={'k': 3})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4aca2050-9141-4007-8462-7748fb0bcce2",
   "metadata": {},
   "source": [
    "## Retriever arguments\n",
    "\n",
    "These are the arguments to the retriever:\n",
    "\n",
    "    ‘k’: the number of documents to return (kNN search)\n",
    "\n",
    "\n",
    "## Making a prompt\n",
    "\n",
    "code view 15:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fab45e76-3f86-4de2-99bf-511f2892c32e",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "outputs": [],
   "source": [
    "from langchain.prompts import PromptTemplate\n",
    "\n",
    "prompt_template = '''You are an assistant for question-answering tasks.\n",
    "Use the following pieces of retrieved context to answer the question.\n",
    "Context: {context}\n",
    "\n",
    "Question: {input}\n",
    "\n",
    "Answer:\n",
    "'''\n",
    "\n",
    "prompt = PromptTemplate(template=prompt_template,\n",
    "                        input_variables=['context', 'input'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24c861e0-8eaf-4115-ba40-25405ba26d9a",
   "metadata": {},
   "source": [
    "## Making the \"Chatbot\"\n",
    "\n",
    "code view 16:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16b75bb9-6ab7-42bd-a5a1-4d15fd5b4fe0",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "outputs": [],
   "source": [
    "from langchain.chains import create_retrieval_chain\n",
    "from langchain.chains.combine_documents import create_stuff_documents_chain\n",
    "\n",
    "combine_documents_chain = create_stuff_documents_chain(llm, prompt)\n",
    "rag_chain = create_retrieval_chain(retriever, combine_documents_chain)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0a5b7b7-e2c5-4539-a2ad-4974242e63ad",
   "metadata": {},
   "source": [
    "## Asking the \"Chatbot\"\n",
    "\n",
    "code view 17:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7adbad84-a22a-4ee9-b196-508fdc43dd5e",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "outputs": [],
   "source": [
    "result = rag_chain.invoke({'input': query})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca416170-eb13-4982-b83a-0fd2c5a4a1a5",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "code view 18:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc607ad8-7a3e-41c4-8718-3fa617cf749b",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "outputs": [],
   "source": [
    "print(result['answer'])"
   ]
  }
 ],
 "metadata": {
  "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.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}