llm_document_answering.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "95853677",
   "metadata": {},
   "source": [
    "# Motivation"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4df2ea35",
   "metadata": {},
   "source": [
    "We want to create an opportunity to chat with a custom knowledge base about amazon sagemaker. For this we use a sagemaker QNA dataset and make it available to a large language model. \n",
    "\n",
    "Since the knowledge is very domain specific about AWS sagemaker, we could ask a simple LLM questions about Amazon or Amazon Web Services (AWS) in general, but we might struggle to find specific details about AWS sagemaker. \n",
    "\n",
    "The general idea is: If we have sensitive data or company data, we can not use the OpenAI API, since we would send our documents to the API. This has two disadvantages:\n",
    "\n",
    "- we share our sensitive data with OpenAI\n",
    "- we need an API key for this, which means it costs us money per requests depending on the number of tokens we send to the API\n",
    "\n",
    "To be able to keep privacy of our data, we show a process how to use a small LLM locally, which will run on our computer, and use a local vector database (ChromaDB) for our domain specific knowledge. In this environment no data is leaving our local machine. This can be tested by setting of the internet connection after the model was downloaded.\n",
    "\n",
    "Let's start by setting up the environment!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5fade883",
   "metadata": {},
   "source": [
    "# Setting up the environment"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "79c0bddf",
   "metadata": {},
   "source": [
    "Requirements:\n",
    "\n",
    "- a notebook or computer with sufficient memory 16GB\n",
    "- python environment with version 3.10.4"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a3d007f8",
   "metadata": {},
   "source": [
    "First we need to setup an environment. For the use of python, we recommend using pyenv virtualenv to setup a specific python version (in this case python 3.10.4) and use poetry to take care of the required python packages.\n",
    "\n",
    "To setup the environment, just use `pyenv virtualenv 3.10.4 llm_chromadb` or a similar method to setup the correct python version.\n",
    "\n",
    "Maybe an installation of python 3.10.4 is required before by using `pyenv install 3.10.4` \n",
    "\n",
    "After the setup, just use `pyenv activate llm_chromadb`. When the python environment is activated, just use `poetry install` inside this repository folder to install all required packages defined in the [pyproject.toml](pyproject.toml) for this project. \n",
    "\n",
    "Now we are ready to start!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a0d0a1c8",
   "metadata": {},
   "source": [
    "# Prepare the documents, setup the LLM and create the vector database \n",
    "\n",
    "First lets import the required classes from the langchain package: \n",
    "- TextLoader to load the texts from our file\n",
    "\n",
    "- RecursiveCharacterTextSplitter to split up the text inside the document\n",
    "\n",
    "- Chroma-DB for the vector database\n",
    "\n",
    "- VectorDBQA to do the retrieval of information from the vector database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "652985d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import TextLoader\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain.vectorstores import Chroma\n",
    "from langchain.chains import VectorDBQA"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "01fe5351",
   "metadata": {},
   "source": [
    "## Load documents\n",
    "\n",
    "Load the documents to do question answering over the sagemaker qna dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "06619a79",
   "metadata": {},
   "outputs": [],
   "source": [
    "loader = TextLoader('Sagemaker_qna.txt')\n",
    "qna = loader.load()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "478be861",
   "metadata": {},
   "source": [
    "## Split documents\n",
    "\n",
    "Now we want to split these documents into smaller chunks. This allows to find the most relevant chunks for a query and pass only those into the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "08eaaf68",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_splitter_qna = RecursiveCharacterTextSplitter(chunk_size=1000,chunk_overlap=0)\n",
    "texts_qna = text_splitter_qna.split_documents(qna)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f9b30aff",
   "metadata": {},
   "source": [
    "## Initialize ChromaDB\n",
    "\n",
    "Create embeddings for each chunk and insert into the Chroma vector database. For this step, we first ne the correct emebdding to the model we want to use later."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "989798d7",
   "metadata": {},
   "source": [
    "We want to use the [google/flan-t5-large](https://huggingface.co/google/flan-t5-large) model from hugginface, since its size is quite small and should be sufficient for this tutorial. For our documents vector database we want to use the same embeddings as for the original model. For this we use the SentenceTransformerEmbeddings.\n",
    "\n",
    "Lets get the model from the huggingface_hub. Warning: this model has a size of about 3GB, so please only download it on your laptop with a stable internet connection. Depending on your internet connection, this might take some minutes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8c37d244",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jonasbechthold/.pyenv/versions/3.10.4/envs/llm_chatbot/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "No sentence-transformers model found with name /Users/jonasbechthold/.cache/torch/sentence_transformers/google_flan-t5-large. Creating a new one with MEAN pooling.\n",
      "Some weights of the model checkpoint at /Users/jonasbechthold/.cache/torch/sentence_transformers/google_flan-t5-large were not used when initializing T5EncoderModel: ['lm_head.weight']\n",
      "- This IS expected if you are initializing T5EncoderModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing T5EncoderModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings\n",
    "model_name = \"google/flan-t5-large\"\n",
    "embedding_function = SentenceTransformerEmbeddings(model_name=model_name)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e5cd23c0",
   "metadata": {},
   "source": [
    "Now this we have the embedding function, we can create the vector database by putting in our QNAs of sagemaker (texts_qna) and the embeddings into the Chroma_from_documents() method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "c0d2a049",
   "metadata": {},
   "outputs": [],
   "source": [
    "vectordb = Chroma.from_documents(texts_qna, embedding_function)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "64967063",
   "metadata": {},
   "source": [
    "## Create the LLM locally"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b73c883",
   "metadata": {},
   "source": [
    "Now, lets get the LLM model of the \"google-flant-t5-large\" and use it locally on our computer. For this we get the model and the tokenizer from huggingface and create a pipeline for the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e5651a71",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms import HuggingFacePipeline\n",
    "from transformers import AutoTokenizer, pipeline, AutoModelForSeq2SeqLM\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(model_name)\n",
    "\n",
    "pipe = pipeline(\n",
    "    \"text2text-generation\",\n",
    "    model=model, \n",
    "    tokenizer=tokenizer, \n",
    "    max_length=100\n",
    ")\n",
    "\n",
    "local_llm = HuggingFacePipeline(pipeline=pipe)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1ddb7866",
   "metadata": {},
   "source": [
    "## Create the LLM chain consisting of prompt and LLM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c2960a6",
   "metadata": {},
   "source": [
    "Now we initialize the chain we will use for question answering. In this case we just use Langchain as a convenience tool, to tell the model which prompt with which input question it should use. We define a simple prompt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ed80ee13",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import PromptTemplate\n",
    "\n",
    "template = \"\"\"Question: {question}\n",
    "\n",
    "Answer: Let's think step by step.\n",
    "\n",
    "Answer: \"\"\"\n",
    "\n",
    "prompt = PromptTemplate(template=template, input_variables=[\"question\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2c2aace4",
   "metadata": {},
   "source": [
    "Now we connect the prompt with the LLM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "fb5164cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import LLMChain\n",
    "\n",
    "llm_chain = LLMChain(prompt=prompt, \n",
    "                     llm=local_llm\n",
    "                     )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0615f7c1",
   "metadata": {},
   "source": [
    "Lets ask the LLM a very general question, it should know the answer, since this is quite general knowledge, which would be available from wikipedia:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fd6c875b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The capital of Germany is Berlin. Berlin is located in Germany. So the final answer is Berlin.\n"
     ]
    }
   ],
   "source": [
    "question1 = \"What is the capital of Germany?\"\n",
    "\n",
    "print(llm_chain.run(question1))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3e1a854f",
   "metadata": {},
   "source": [
    "## Build the information retrieval from our vector db"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fea580f0",
   "metadata": {},
   "source": [
    "Now we build the method to extract relevant knowledge from our vector database to be able to both chat with the LLM and chat with our QNA vector database. For this we put in the LLM (flan-t5-large), a chain type (stuff) and the vector database (vectordb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f1f859b6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jonasbechthold/.pyenv/versions/3.10.4/envs/llm_chatbot/lib/python3.10/site-packages/langchain/chains/retrieval_qa/base.py:243: UserWarning: `VectorDBQA` is deprecated - please use `from langchain.chains import RetrievalQA`\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "qa = VectorDBQA.from_chain_type(llm=local_llm, chain_type=\"stuff\", vectorstore=vectordb)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cff96efe",
   "metadata": {},
   "source": [
    "# Ask questions to the LLM and the vector database\n",
    "\n",
    "To have a comparison between the asked question without the doucment database, we ask a question to the pure llm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f207aec9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The Amazon Sagemaker is a type of adverb. The Amazon Sagemaker is a type of adverb. The answer: adverb.\n"
     ]
    }
   ],
   "source": [
    "question = \"What is amazon sagemaker?\"\n",
    "\n",
    "print(llm_chain.run(question))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b54fc125",
   "metadata": {},
   "source": [
    "The anser we get is quite strange. The llm seems not to have the appropriate knowledge.\n",
    "\n",
    "Now we ask the vector database with the llm on top of it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7f139435",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Amazon SageMaker is a fully managed service that provides every developer and data scientist with the ability to build, train, and deploy machine learning (ML) models quickly'"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = \"What is amazon sagemaker\"\n",
    "qa.run(query)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "18cd929d",
   "metadata": {},
   "source": [
    "Now the answer is much more appropriate."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0b1055d2",
   "metadata": {},
   "source": [
    "Lets ask our documents a few more questions:"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7dd4847a",
   "metadata": {},
   "source": [
    "This seems to be correct. We continue with some more specific questions and get quite promising results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "c0b67fe0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'AWS pricing'"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = \"What do you know about the pricing of aws sagemaker?\"\n",
    "qa.run(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "cb8f16f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Amazon SageMaker pricing is based on the number of instances you need.'"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = \"What do you know about the pricing of aws sagemaker? Give me some example prices!\"\n",
    "qa.run(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3a72bd2e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Yes'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = \"Does Sagemaker support R?\"\n",
    "qa.run(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "08a0f5f6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Yes, R supported with Amazon SageMaker. You can use R within SageMaker Notebook instances, which include a pre-installed R kernel and the reticulate library.'"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = \"Does sagemaker support R and in which applications?\"\n",
    "qa.run(query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "a9fac497",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'With the new notebook experience, you can now quickly launch notebooks without needing to manually provision an instance and waiting for it to be operational.'"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = \"How are SageMaker Studio Notebooks different from the notebooks instances?\"\n",
    "qa.run(query)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c74728d",
   "metadata": {},
   "source": [
    "That's it, we successfully asked our vector database questions and got quite promising question. Considering the small size of the model, a larger model trained on this specific purpose would be a great part for future exploration."
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}