Data_NoBlaBla_final.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Po29FbgVhPkM"
      },
      "source": [
        "**Your AI assistant**\n",
        "\n",
        "Vous venez d'être embauché par la NBT (Nouvelle Bibliothèque de Toulouse) qui a entièrement dématérialisé ses ouvrages sous la forme de pdf 😧\n",
        "\n",
        "Elle souhaite mettre en place un service de FAQ sur la base de ces ouvrages et vous a confié cette mission importante.\n",
        "\n",
        "Dans 3 heures, vous devez mettre ce service en ligne et il ne vous reste plus que quelques fonctions à écrire. Vous décidez de tester ces fonctions sur la base de la documentation des documentations des machines à café de votre service ☕ sur la base de ce notebook.\n",
        "\n",
        "Y parviendrez vous?\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "Timing: ~3h (watch out for long processing)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O-wJ451w4ZLr"
      },
      "source": [
        "## Installing dependencies"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G5exRj1P472K"
      },
      "source": [
        "### Langchain and vector storage"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gm1J5fyjTdHJ"
      },
      "outputs": [],
      "source": [
        "%%shell\n",
        "pip install --quiet chromadb==0.4.18 langchain==0.0.349 loguru==0.7.2 pydantic==1.10.13 sentence-transformers==2.2.2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EziWyyOI5AkF"
      },
      "source": [
        "### PDF processing library"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DAq5_ffr3pMM"
      },
      "outputs": [],
      "source": [
        "%%shell\n",
        "pip install --quiet pymupdf==1.23.7"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BNLiJ-Wg5EEa"
      },
      "source": [
        "### LLM usage libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dHJznepbaecl"
      },
      "outputs": [],
      "source": [
        "%%shell\n",
        "pip install --quiet transformers==4.36.0 accelerate==0.25.0 bitsandbytes==0.41.3.post2 einops==0.7.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dnQfDLcTBEXL"
      },
      "source": [
        "### UI packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8h5ZkCOiBDjA"
      },
      "outputs": [],
      "source": [
        "%%shell\n",
        "pip install --quiet gradio==3.44.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5fEKYTAN6B3G"
      },
      "source": [
        "### Update Runtime\n",
        "Some of the libraries installed require a restart of the Kernel/Runtime to be properly loaded. Do it now."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kHeYFxGCF9a3"
      },
      "source": [
        "## Import"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NJ2z1Ib4Ap34"
      },
      "outputs": [],
      "source": [
        "from langchain.globals import set_debug\n",
        "from torch.cuda import empty_cache, ipc_collect, is_available\n",
        "\n",
        "set_debug(True)\n",
        "\n",
        "\n",
        "def flush_gpu_memory():\n",
        "  if is_available():\n",
        "    empty_cache()\n",
        "    ipc_collect()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gno2ny7OGOVD"
      },
      "source": [
        "## Config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MubCzlnH83Hp"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "WORKING_DIR = \"/content\"\n",
        "os.environ['WORKING_DIR'] = WORKING_DIR"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# RAG components implementation"
      ],
      "metadata": {
        "id": "nkodpnO_bUdS"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yKUNoQNUFy0v"
      },
      "source": [
        "## Read documents and chunk them"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VQmGL6y7zt6Y"
      },
      "source": [
        "### Get some documents\n",
        "\n",
        "Here we just download PDF documents from the web. Any PDF would fit and you can get lots of them. Just be cautious that more documents (and expecially pages) will lead to more processing time and extra memory usage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SVK1qEG3z_B8"
      },
      "outputs": [],
      "source": [
        "%%shell\n",
        "curl https://www.nespresso.com/shared_res/manuals/inissia/inissia_C_breville.pdf > $WORKING_DIR/sample_data/inissia_C_breville.pdf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sK95BhXozuts"
      },
      "source": [
        "### Document preprocessing\n",
        "\n",
        "This simple preprocessing function is using Langchain splitter to identify chunk in the document(s). Several splitting approaches are possible with potentially different granularity in the document indexing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lfjy6DW700kO"
      },
      "outputs": [],
      "source": [
        "from langchain.text_splitter import CharacterTextSplitter\n",
        "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
        "\n",
        "\n",
        "def process_docs(docs):\n",
        "  chunk_size = 500\n",
        "  chunk_overlap = 50\n",
        "\n",
        "  #text_splitter = CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)\n",
        "  text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)\n",
        "\n",
        "  chunked_documents = list()\n",
        "  for doc_loader in docs:\n",
        "    chunked_documents.extend(doc_loader.load_and_split(text_splitter))\n",
        "\n",
        "  return chunked_documents"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ansV00211gLS"
      },
      "outputs": [],
      "source": [
        "from langchain.document_loaders import PyMuPDFLoader\n",
        "\n",
        "documents = []\n",
        "for file in os.listdir('sample_data'):\n",
        "    if file.endswith('.pdf'):\n",
        "        pdf_path = f\"{WORKING_DIR}/sample_data/\" + file\n",
        "        doc_loader = PyMuPDFLoader(pdf_path)\n",
        "        documents.append(doc_loader)\n",
        "\n",
        "chunked_documents = process_docs(documents)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VWK7fcRTGgYw"
      },
      "outputs": [],
      "source": [
        "from loguru import logger\n",
        "\n",
        "assert len(chunked_documents) > 0, \"Please load document\"\n",
        "\n",
        "logger.info(f\"Found {len(chunked_documents)} pieces of texts over {len(documents)} documents\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Visualizing a document chunk"
      ],
      "metadata": {
        "id": "0a1xmVJ7Pt4g"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U0bXTNEtGv9Z"
      },
      "outputs": [],
      "source": [
        "for idx, chunk in enumerate(chunked_documents):\n",
        "  if chunk.metadata['page'] == 15:\n",
        "    print(f\"Chunk #{idx}\")\n",
        "    print(chunk.page_content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ioQ7-XuJIbj5"
      },
      "source": [
        "## Compute and store embeddings\n",
        "\n",
        "Selecting sentence a embedding model: https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2\n",
        "\n",
        "This will be used to conver the text chunk into embedding vector to enable vector search (aka semantic). The size and quality of the embedding might lead to various result quality, `all-MiniLM-L12-v2` being quite a solid baseline.\n",
        "\n",
        "Learn more:\n",
        "* https://huggingface.co/blog/mteb\n",
        "* https://huggingface.co/spaces/mteb/leaderboard"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Load embedding model"
      ],
      "metadata": {
        "id": "OSKzu38DOBCy"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KGz9zVFlreeY"
      },
      "outputs": [],
      "source": [
        "embeddings_model_name = \"all-MiniLM-L6-v2\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ov-SFzUt9yMu"
      },
      "outputs": [],
      "source": [
        "from langchain.embeddings import HuggingFaceEmbeddings\n",
        "from langchain.vectorstores import Chroma\n",
        "\n",
        "embedding_model = HuggingFaceEmbeddings(model_name=embeddings_model_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Create retriever\n",
        "\n",
        "We are using Chroma (a simple vector database) to store the embedding and then build a retriever engine on top of it."
      ],
      "metadata": {
        "id": "iEROOUW6ODyX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "vectordb = Chroma.from_documents(chunked_documents, embedding=embedding_model, persist_directory=\"./db\")\n",
        "vectordb.persist()\n",
        "retriever = vectordb.as_retriever(search_type=\"mmr\", search_kwargs={\"k\": 5, \"fetch_k\": 10})"
      ],
      "metadata": {
        "id": "Cjh8fKgYJSxG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0-C98kJaIk6J"
      },
      "source": [
        "### Simple search test"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RXa7VFlQIjZx"
      },
      "outputs": [],
      "source": [
        "retrieved_docs = retriever.invoke(input=\"cold coffee\")\n",
        "\n",
        "for i, doc in enumerate(retrieved_docs):\n",
        "  print(f'\\n<<{i}>> on page {doc.metadata[\"page\"]}: \\n{doc.page_content}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FgaLTE70ALtm"
      },
      "source": [
        "## Local LLM"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Load model"
      ],
      "metadata": {
        "id": "339U0urnUeHP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "flush_gpu_memory()"
      ],
      "metadata": {
        "id": "E54OiTfXUtSY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WFuf0sLOpqDC"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig\n",
        "import torch\n",
        "\n",
        "\n",
        "def load_model(model_name):\n",
        "  flush_gpu_memory()\n",
        "\n",
        "  tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "\n",
        "  model = None\n",
        "  if model_name == \"TinyLlama/TinyLlama-1.1B-Chat-v0.6\":\n",
        "      model = AutoModelForCausalLM.from_pretrained(\n",
        "          model_name, device_map=\"auto\", torch_dtype=\"auto\"\n",
        "      )\n",
        "  elif model_name == \"HuggingFaceH4/zephyr-7b-beta\":\n",
        "      model = AutoModelForCausalLM.from_pretrained(\n",
        "          model_name, device_map=\"auto\", torch_dtype=\"auto\", load_in_4bit=True\n",
        "      )\n",
        "  elif model_name == \"SkunkworksAI/phi-2\":\n",
        "      model = AutoModelForCausalLM.from_pretrained(\n",
        "          model_name, device_map=\"auto\", torch_dtype=torch.float16, trust_remote_code=True\n",
        "      )\n",
        "  else:\n",
        "    assert False, f\"{model_name} unknown\"\n",
        "\n",
        "  generation_config = GenerationConfig.from_pretrained(model_name)\n",
        "  generation_config.max_new_tokens = 1024\n",
        "  generation_config.temperature = 0.0001\n",
        "  generation_config.do_sample = True\n",
        "\n",
        "  return model, tokenizer, generation_config"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# model_name = \"TinyLlama/TinyLlama-1.1B-Chat-v0.6\"\n",
        "#model_name = \"mistralai/Mistral-7B-Instruct-v0.2\"\n",
        "model_name = \"HuggingFaceH4/zephyr-7b-beta\"\n",
        "# model_name = \"microsoft/phi-2\"\n",
        "\n",
        "model, tokenizer, generation_config = load_model(model_name)"
      ],
      "metadata": {
        "id": "Y6aXTNS3Txxc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ctdx4ang8W-D"
      },
      "source": [
        "### Conversational pipeline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7rizvfKC6-p7"
      },
      "outputs": [],
      "source": [
        "from transformers import pipeline\n",
        "\n",
        "pipe_conversation = pipeline(\n",
        "    \"conversational\",\n",
        "    model=model,\n",
        "    tokenizer=tokenizer,\n",
        "    generation_config=generation_config,\n",
        "    num_return_sequences=1,\n",
        "    eos_token_id=tokenizer.eos_token_id,\n",
        "    pad_token_id=tokenizer.eos_token_id,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Basic test"
      ],
      "metadata": {
        "id": "4U4ZkCPnRTIa"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "clx-dqUK6-sh"
      },
      "outputs": [],
      "source": [
        "messages = [\n",
        "    {\n",
        "        \"role\": \"system\",\n",
        "        \"content\": \"\"\"You are a troubleshooting chatbot that talks like a pirate. Your goal is to help user solving the problem they have with\n",
        "their appliance.\n",
        "\n",
        "Make short messages of one or two sentences maximum.\n",
        "\"\"\",\n",
        "    },\n",
        "    {\n",
        "        \"role\": \"user\",\n",
        "        \"content\": \"cold coffee\"\n",
        "    },\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bSKseCRxgbUM"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "conversation = pipe_conversation(messages)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EOeshu9S7XjI"
      },
      "outputs": [],
      "source": [
        "for is_user, text in conversation.iter_texts():\n",
        "  role = \"USER\" if is_user else \"BOT\"\n",
        "  print(role + \"> \" + text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kLBaIc4R8eBg"
      },
      "source": [
        "### RAG pipeline"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SZpOfAasV8OI"
      },
      "outputs": [],
      "source": [
        "from transformers import pipeline\n",
        "\n",
        "pipe_text_generation = pipeline(\n",
        "    \"text-generation\",\n",
        "    model=model,\n",
        "    tokenizer=tokenizer,\n",
        "    generation_config=generation_config,\n",
        "    num_return_sequences=1,\n",
        "    eos_token_id=tokenizer.eos_token_id,\n",
        "    pad_token_id=tokenizer.eos_token_id,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JLUnFYWvHumK"
      },
      "outputs": [],
      "source": [
        "from langchain.llms.huggingface_pipeline import HuggingFacePipeline\n",
        "\n",
        "llm = HuggingFacePipeline(pipeline=pipe_text_generation)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s2z_Ycpx2aqs"
      },
      "outputs": [],
      "source": [
        "flush_gpu_memory()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Basic test"
      ],
      "metadata": {
        "id": "KI5hOo-4RWv7"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WvjsWJlwuEIc"
      },
      "outputs": [],
      "source": [
        "from time import time\n",
        "\n",
        "from langchain.chains import RetrievalQA\n",
        "\n",
        "query = \"cold coffee\"\n",
        "\n",
        "hide_source = True # Switch that to see what has been used to retrieve info\n",
        "qa = RetrievalQA.from_chain_type(llm=llm, chain_type=\"stuff\", retriever=retriever, return_source_documents= not hide_source)\n",
        "\n",
        "# Get the answer from the chain\n",
        "start = time()\n",
        "res = qa(query)\n",
        "answer, docs = res['result'], [] if hide_source else res['source_documents']\n",
        "end = time()\n",
        "\n",
        "# Print the result\n",
        "print(\"\\n\\n> Question:\")\n",
        "print(query)\n",
        "print(\"\\n\\n> Anwser:\")\n",
        "print(answer)\n",
        "\n",
        "# Print the relevant sources used for the answer\n",
        "for document in docs:\n",
        "    print(\"\\n> \" + document.metadata[\"source\"] + \":\")\n",
        "    print(document.page_content)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Uncomment the following cell if you want a continous run without UI"
      ],
      "metadata": {
        "id": "WmfwA5V0Rdlt"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SH8c_NCmTlQ6"
      },
      "outputs": [],
      "source": [
        "# import time\n",
        "\n",
        "# hide_source = True # Switch that to see what has been used to retrieve info\n",
        "# while True:\n",
        "#     query = input(\"\\nEnter a query: \")\n",
        "#     if query == \"exit\":\n",
        "#         break\n",
        "#     if query.strip() == \"\":\n",
        "#         continue\n",
        "\n",
        "#     # Get the answer from the chain\n",
        "#     start = time.time()\n",
        "#     res = qa(query)\n",
        "#     answer, docs = res['result'], [] if hide_source else res['source_documents']\n",
        "#     end = time.time()\n",
        "\n",
        "#     # Print the result\n",
        "#     print(\"\\n\\n> Question:\")\n",
        "#     print(query)\n",
        "#     print(answer)\n",
        "\n",
        "#     # Print the relevant sources used for the answer\n",
        "#     for document in docs:\n",
        "#         print(\"\\n> \" + document.metadata[\"source\"] + \":\")\n",
        "#         print(document.page_content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uCa14bOyeWD-"
      },
      "source": [
        "# Conversational demo"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Recap on RAG architecture\n",
        "\n",
        "![Diagram from wikipedia](https://upload.wikimedia.org/wikipedia/commons/thumb/3/37/RAG_schema.svg/1280px-RAG_schema.svg.png)"
      ],
      "metadata": {
        "id": "WhH5QIEgZumg"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Document encoder & vector database for retriever"
      ],
      "metadata": {
        "id": "Pf2lG7QoKBFh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def process_docs(docs):\n",
        "  chunk_size = 500\n",
        "  chunk_overlap = 50\n",
        "\n",
        "  #text_splitter = CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)\n",
        "  text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)\n",
        "\n",
        "  chunked_documents = list()\n",
        "  for doc in docs:\n",
        "    doc_loader = PyMuPDFLoader(doc)\n",
        "    chunked_documents.extend(doc_loader.load_and_split(text_splitter))\n",
        "\n",
        "  return chunked_documents\n",
        "\n",
        "\n",
        "def doc_retriever(chunked_documents, model=embedding_model):\n",
        "  vectordb = Chroma.from_documents(chunked_documents, embedding=embedding_model, persist_directory=\"./db\")\n",
        "  vectordb.persist()\n",
        "  retriever = vectordb.as_retriever(search_type=\"mmr\", search_kwargs={\"k\": 3, \"fetch_k\": 5})\n",
        "\n",
        "  return retriever"
      ],
      "metadata": {
        "id": "v4tge_D-JN52"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Conversational LLM"
      ],
      "metadata": {
        "id": "0KYTlQNbKFj9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# model_name = \"TinyLlama/TinyLlama-1.1B-Chat-v0.6\"\n",
        "#model_name = \"mistralai/Mistral-7B-Instruct-v0.2\"\n",
        "model_name = \"HuggingFaceH4/zephyr-7b-beta\"\n",
        "# model_name = \"SkunkworksAI/phi-2\"\n",
        "\n",
        "flush_gpu_memory()\n",
        "\n",
        "model, tokenizer, generation_config = load_model(model_name)\n",
        "\n",
        "pipe_conversation = pipeline(\n",
        "    \"conversational\",\n",
        "    model=model,\n",
        "    tokenizer=tokenizer,\n",
        "    generation_config=generation_config,\n",
        "    num_return_sequences=1,\n",
        "    eos_token_id=tokenizer.eos_token_id,\n",
        "    pad_token_id=tokenizer.eos_token_id,\n",
        ")"
      ],
      "metadata": {
        "id": "wjgmqhm5KLrY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def request_chat(messages):\n",
        "  logger.info(\"Generating next message...\")\n",
        "  conversation = pipe_conversation(messages)\n",
        "\n",
        "  turn = 0\n",
        "  messages = list()\n",
        "  for is_user, text in conversation.iter_texts():\n",
        "    role = \"USER\" if is_user else \"BOT\"\n",
        "    if turn == 0:\n",
        "      role='system'\n",
        "    else:\n",
        "      print(turn, '#', role + \"> \" + text)\n",
        "\n",
        "    messages.append(dict(role=role, content=text))\n",
        "    turn += 1\n",
        "\n",
        "  # return last message from conversation\n",
        "  return messages[-1][\"content\"]"
      ],
      "metadata": {
        "id": "_WiWOrXfUkq4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Conversation prompt"
      ],
      "metadata": {
        "id": "4pYQwSKIKQti"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "system_prompt = \"\"\"You are a troubleshooting chatbot that talks like a pirate. Your goal is to help user solving the problem they have with\n",
        "their appliance. One of the first thing is to ask the user what is the product or appliance concerned.\n",
        "\"\"\"\n",
        "\n",
        "contextual_prompt = \"\"\"\\nIn the following, we will provide you contextual information about the product from a\n",
        "document. Use it to answer questions from the user.\n",
        "Your answer MUST ONLY come from this text.\n",
        "DO NOT answer if the information is insufficient.\n",
        "DO NOT be creative in your answer and only cite what from the contextual information.\n",
        "You still talk like a pirate.\n",
        "\"\"\"\n",
        "\n",
        "file_uploaded_prompt = \"Explain briefly that the file has been processed and it's now ready for questions. You still talk like a pirate.\"\n",
        "\n",
        "def raw_llm_mode_prompt(prompt, chunked_documents):\n",
        "  logger.info('Update prompt in RAW_LLM_MODE')\n",
        "  # default RAW mode where the max content is added in prompt as context\n",
        "  total_chars_count = 0\n",
        "  prompt += '\\n\\nHere is the contextual information document which is limited to few pages:'\n",
        "\n",
        "  for chunk in chunked_documents:\n",
        "    content = chunk.page_content\n",
        "    total_chars_count += len(content)\n",
        "    if (total_chars_count + 300) > CONTEXT_MAX_CHARS:\n",
        "      break\n",
        "    prompt += '\\n' + str(content)\n",
        "\n",
        "  return prompt\n",
        "\n",
        "\n",
        "def rag_llm_mode_prompt(query, prompt, chunked_documents, retriever):\n",
        "  logger.info('Update prompt in RAG_MODE')\n",
        "\n",
        "  if isinstance(query, str):\n",
        "    retrieved_docs = retriever.invoke(query)\n",
        "    logger.info(f'Context found:')\n",
        "    for i, doc in enumerate(retrieved_docs):\n",
        "      logger.info(f'\\n<<{i}>> on page {doc.metadata[\"page\"]}: \\n{doc.page_content[0:100]} [...]')\n",
        "\n",
        "    prompt += '\\nWhen providing an answer, you SHALL cite the page number and suggest the user to read it.'\n",
        "    prompt += '\\n\\nHere is the contextual information document which is focused on the key pages:'\n",
        "    total_chars_count = 0\n",
        "    for i, doc in enumerate(retrieved_docs):\n",
        "        added_prompt = '\\n\\n# ON PAGE '+str(doc.metadata[\"page\"])\n",
        "        added_prompt += '\\n' + str(doc.page_content)\n",
        "\n",
        "        total_chars_count += len(added_prompt)\n",
        "        if (total_chars_count + 300) > CONTEXT_MAX_CHARS:\n",
        "            break\n",
        "\n",
        "        prompt += added_prompt\n",
        "  else:\n",
        "    logger.warning(\"No query yet, RAG mode will be started as soon as a query arrives.\")\n",
        "\n",
        "  return prompt\n",
        "\n",
        "def adapt_prompt(history, file_name, messages):\n",
        "    prompt = system_prompt\n",
        "    if file_name is None:\n",
        "        logger.warning('No file.')\n",
        "    else:\n",
        "        chunked_documents, retriever = session_context['files'][file_name]\n",
        "\n",
        "        logger.info(\"Adding content to prompt...\")\n",
        "        prompt += contextual_prompt\n",
        "        prompt += '\\n'\n",
        "\n",
        "        if session_context['llm_mode'] == RAW_LLM_MODE:\n",
        "            prompt += raw_llm_mode_prompt(prompt, chunked_documents)\n",
        "\n",
        "        if session_context['llm_mode'] == RAG_MODE:\n",
        "            # use last user message to query relevant pages as context and add it to the prompt\n",
        "            query = history[-1][0]\n",
        "            prompt += rag_llm_mode_prompt(query, prompt, chunked_documents, retriever)\n",
        "\n",
        "        # update original system prompt\n",
        "        messages[0] = dict(role=\"system\", content=prompt)\n",
        "\n",
        "        logger.info(f\"Updated prompt.\")\n",
        "\n",
        "    return messages"
      ],
      "metadata": {
        "id": "ZJzdxqJgSW5n"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rzPUEtcYc9Un"
      },
      "outputs": [],
      "source": [
        "from random import randint\n",
        "\n",
        "import gradio as gr\n",
        "import time\n",
        "\n",
        "import logging\n",
        "\n",
        "\n",
        "logger = logging.getLogger('rag_chatbot')\n",
        "logger.setLevel(logging.INFO)\n",
        "\n",
        "CONTEXT_MAX_CHARS = 2000\n",
        "RAW_LLM_MODE = \"Fixed context\"\n",
        "RAG_MODE = \"RAG\"\n",
        "\n",
        "session_context = dict(\n",
        "  files=dict(),\n",
        "  llm_mode=RAW_LLM_MODE\n",
        ")\n",
        "\n",
        "def push_response(history, response):\n",
        "  history[-1][1] = \"\"\n",
        "  for character in response:\n",
        "    history[-1][1] += character\n",
        "    time.sleep(randint(25, 75)/1000)\n",
        "    yield history\n",
        "\n",
        "\n",
        "def add_text(history, text):\n",
        "  history = history + [(text, None)]\n",
        "  return history, gr.update(value=\"\", interactive=False)\n",
        "\n",
        "\n",
        "def add_file(history, file):\n",
        "  history = history + [((file.name, None), None)]\n",
        "  return history\n",
        "\n",
        "\n",
        "def change_mode(llm_mode):\n",
        "  logger.info(f'Updating LLm Mode: llm_mode={llm_mode}')\n",
        "  session_context['llm_mode'] = llm_mode\n",
        "\n",
        "\n",
        "def process_message_history(history):\n",
        "  messages = [\n",
        "    dict(role=\"system\", content=system_prompt)\n",
        "  ]\n",
        "  file_name = None\n",
        "  for msg in history:\n",
        "    if isinstance(msg[0], str):\n",
        "      messages.append(dict(role=\"user\", content=msg[0]))\n",
        "    if isinstance(msg[0], tuple):\n",
        "      file_name = msg[0][0]\n",
        "      if file_name not in session_context['files']:\n",
        "        logger.info(\"Processing file...\")\n",
        "\n",
        "        chunked_documents = process_docs([file_name])\n",
        "        retriever = doc_retriever(chunked_documents)\n",
        "\n",
        "        session_context['files'][file_name] = (chunked_documents, retriever)\n",
        "\n",
        "      messages.append(dict(role=\"user\", content=\"The user uploaded a file: \" + file_name))\n",
        "\n",
        "    if msg[1] is not None and isinstance(msg[1], str):\n",
        "      messages.append(dict(role=\"assistant\", content=msg[1]))\n",
        "  return file_name, messages\n",
        "\n",
        "\n",
        "def bot(history):\n",
        "  logger.info(f\"Preparing Bot message turn={len(history)}\")\n",
        "\n",
        "  # building messages sequence from history\n",
        "  file_name, messages = process_message_history(history)\n",
        "\n",
        "  # if a file was processed, then update prompt\n",
        "  adapt_prompt(history, file_name, messages)\n",
        "\n",
        "  if isinstance(history[-1][0], tuple):\n",
        "    # if last message was file processed, answer that it's ready\n",
        "    messages.append(\n",
        "      dict(\n",
        "        role=\"system\",\n",
        "        content=file_uploaded_prompt\n",
        "      )\n",
        "    )\n",
        "\n",
        "  # get LLM to answer\n",
        "  response = request_chat(messages)\n",
        "  history[-1][1] = response\n",
        "  yield history"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Demo runner"
      ],
      "metadata": {
        "id": "XogEHgnBFFcZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with gr.Blocks() as demo:\n",
        "  chatbot = gr.Chatbot(\n",
        "    [[None, \"Hi, I'm your assistant. Ask me anything or upload a product manual to start.\"]],\n",
        "    elem_id=\"chatbot\",\n",
        "    bubble_full_width=False,\n",
        "  )\n",
        "\n",
        "  with gr.Row():\n",
        "    txt = gr.Textbox(\n",
        "      scale=4,\n",
        "      show_label=False,\n",
        "      placeholder=\"Enter text and press enter, or upload an image\",\n",
        "      container=False,\n",
        "    )\n",
        "    with gr.Column():\n",
        "      btn = gr.UploadButton(\"📁\", file_types=[\"pdf\"])\n",
        "      drpdwn = gr.Dropdown(\n",
        "        [RAW_LLM_MODE, RAG_MODE],\n",
        "        value=RAW_LLM_MODE,\n",
        "        label=\"LLM Mode\",\n",
        "      )\n",
        "\n",
        "  txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(\n",
        "    bot, chatbot, chatbot\n",
        "  )\n",
        "  txt_msg.then(lambda: gr.update(interactive=True), None, [txt], queue=False)\n",
        "  file_msg = btn.upload(add_file, [chatbot, btn], [chatbot], queue=False).then(\n",
        "    bot, chatbot, chatbot\n",
        "  )\n",
        "  drpdwn.change(change_mode, drpdwn)\n",
        "\n",
        "demo.queue()\n",
        "if __name__ == \"__main__\":\n",
        "  demo.queue(concurrency_count=3).launch(\n",
        "    server_name=\"0.0.0.0\",\n",
        "    share=True,\n",
        "    debug=True\n",
        "  )"
      ],
      "metadata": {
        "id": "9C8c7zTcFAkZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "vZxQH5GwM-9q"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "gpuType": "T4"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}