04_retrieval.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0689733d",
   "metadata": {},
   "source": [
    "# Retrieval\n",
    "\n",
    "Retrieval is the centerpiece of our retrieval augmented generation (RAG) flow. \n",
    "\n",
    "Let's get our vectorDB from before."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed2569c6",
   "metadata": {},
   "source": [
    "## Vectorstore retrieval\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51b15e5c-9b92-4d40-a149-b56335d330d9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import openai\n",
    "import sys\n",
    "sys.path.append('../..')\n",
    "\n",
    "from dotenv import load_dotenv, find_dotenv\n",
    "_ = load_dotenv(find_dotenv()) # read local .env file\n",
    "\n",
    "openai.api_key  = os.environ['OPENAI_API_KEY']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c18f8a7b-62af-403e-9877-18d1c2265b4f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#!pip install lark"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2d552e1",
   "metadata": {},
   "source": [
    "### Similarity Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe368042",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.vectorstores import Chroma\n",
    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
    "persist_directory = 'docs/chroma/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0189dc5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "embedding = OpenAIEmbeddings()\n",
    "vectordb = Chroma(\n",
    "    persist_directory=persist_directory,\n",
    "    embedding_function=embedding\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3659e0f7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(vectordb._collection.count())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a807c758",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "texts = [\n",
    "    \"\"\"The Amanita phalloides has a large and imposing epigeous (aboveground) fruiting body (basidiocarp).\"\"\",\n",
    "    \"\"\"A mushroom with a large fruiting body is the Amanita phalloides. Some varieties are all-white.\"\"\",\n",
    "    \"\"\"A. phalloides, a.k.a Death Cap, is one of the most poisonous of all known mushrooms.\"\"\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "715d54f3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "smalldb = Chroma.from_texts(texts, embedding=embedding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a37b5a5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "question = \"Tell me about all-white mushrooms with large fruiting bodies\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24e3b025",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "smalldb.similarity_search(question, k=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4daa6c0d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "smalldb.max_marginal_relevance_search(question,k=2, fetch_k=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a29e8c9",
   "metadata": {},
   "source": [
    "### Addressing Diversity: Maximum marginal relevance\n",
    "\n",
    "Last class we introduced one problem: how to enforce diversity in the search results.\n",
    " \n",
    "`Maximum marginal relevance` strives to achieve both relevance to the query *and diversity* among the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bb2c0a9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "question = \"what did they say about matlab?\"\n",
    "docs_ss = vectordb.similarity_search(question,k=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f07f8793",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "docs_ss[0].page_content[:100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9f7e165",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "docs_ss[1].page_content[:100]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c4ca1b6",
   "metadata": {},
   "source": [
    "Note the difference in results with `MMR`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "222234c5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "docs_mmr = vectordb.max_marginal_relevance_search(question,k=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93b20226",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "docs_mmr[0].page_content[:100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17d39762",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "docs_mmr[1].page_content[:100]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2b909bc",
   "metadata": {},
   "source": [
    "### Addressing Specificity: working with metadata\n",
    "\n",
    "In last lecture, we showed that a question about the third lecture can include results from other lectures as well.\n",
    "\n",
    "To address this, many vectorstores support operations on `metadata`.\n",
    "\n",
    "`metadata` provides context for each embedded chunk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c1a60b2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "question = \"what did they say about regression in the third lecture?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8612840",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "docs = vectordb.similarity_search(\n",
    "    question,\n",
    "    k=3,\n",
    "    filter={\"source\":\"docs/cs229_lectures/MachineLearning-Lecture03.pdf\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97031876",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for d in docs:\n",
    "    print(d.metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2708f6ae",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "ccc2d784",
   "metadata": {},
   "source": [
    "### Addressing Specificity: working with metadata using self-query retriever\n",
    "\n",
    "But we have an interesting challenge: we often want to infer the metadata from the query itself.\n",
    "\n",
    "To address this, we can use `SelfQueryRetriever`, which uses an LLM to extract:\n",
    " \n",
    "1. The `query` string to use for vector search\n",
    "2. A metadata filter to pass in as well\n",
    "\n",
    "Most vector databases support metadata filters, so this doesn't require any new databases or indexes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1d06084",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.llms import OpenAI\n",
    "from langchain.retrievers.self_query.base import SelfQueryRetriever\n",
    "from langchain.chains.query_constructor.base import AttributeInfo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0aa5e698",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "metadata_field_info = [\n",
    "    AttributeInfo(\n",
    "        name=\"source\",\n",
    "        description=\"The lecture the chunk is from, should be one of `docs/cs229_lectures/MachineLearning-Lecture01.pdf`, `docs/cs229_lectures/MachineLearning-Lecture02.pdf`, or `docs/cs229_lectures/MachineLearning-Lecture03.pdf`\",\n",
    "        type=\"string\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"page\",\n",
    "        description=\"The page from the lecture\",\n",
    "        type=\"integer\",\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7906c15",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "document_content_description = \"Lecture notes\"\n",
    "llm = OpenAI(temperature=0)\n",
    "retriever = SelfQueryRetriever.from_llm(\n",
    "    llm,\n",
    "    vectordb,\n",
    "    document_content_description,\n",
    "    metadata_field_info,\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79d781b9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "question = \"what did they say about regression in the third lecture?\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c51778b0-1fcd-40a4-bd6b-0f13fec8acb1",
   "metadata": {},
   "source": [
    "**You will receive a warning** about predict_and_parse being deprecated the first time you executing the next line. This can be safely ignored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d4f9f7d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "docs = retriever.get_relevant_documents(question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db04374e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for d in docs:\n",
    "    print(d.metadata)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "297b8168",
   "metadata": {},
   "source": [
    "### Additional tricks: compression\n",
    "\n",
    "Another approach for improving the quality of retrieved docs is compression.\n",
    "\n",
    "Information most relevant to a query may be buried in a document with a lot of irrelevant text. \n",
    "\n",
    "Passing that full document through your application can lead to more expensive LLM calls and poorer responses.\n",
    "\n",
    "Contextual compression is meant to fix this. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a060cf74",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.retrievers import ContextualCompressionRetriever\n",
    "from langchain.retrievers.document_compressors import LLMChainExtractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "038649c8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def pretty_print_docs(docs):\n",
    "    print(f\"\\n{'-' * 100}\\n\".join([f\"Document {i+1}:\\n\\n\" + d.page_content for i, d in enumerate(docs)]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc686cf2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Wrap our vectorstore\n",
    "llm = OpenAI(temperature=0)\n",
    "compressor = LLMChainExtractor.from_llm(llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82794397",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=compressor,\n",
    "    base_retriever=vectordb.as_retriever()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cde86848",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "question = \"what did they say about matlab?\"\n",
    "compressed_docs = compression_retriever.get_relevant_documents(question)\n",
    "pretty_print_docs(compressed_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82c4fc4d",
   "metadata": {},
   "source": [
    "## Combining various techniques"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "161ae1ad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=compressor,\n",
    "    base_retriever=vectordb.as_retriever(search_type = \"mmr\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e77ccae1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "question = \"what did they say about matlab?\"\n",
    "compressed_docs = compression_retriever.get_relevant_documents(question)\n",
    "pretty_print_docs(compressed_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c2b63e1",
   "metadata": {},
   "source": [
    "## Other types of retrieval\n",
    "\n",
    "It's worth noting that vectordb as not the only kind of tool to retrieve documents. \n",
    "\n",
    "The `LangChain` retriever abstraction includes other ways to retrieve documents, such as TF-IDF or SVM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83d2e808",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.retrievers import SVMRetriever\n",
    "from langchain.retrievers import TFIDFRetriever\n",
    "from langchain.document_loaders import PyPDFLoader\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcf5b760",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load PDF\n",
    "loader = PyPDFLoader(\"docs/cs229_lectures/MachineLearning-Lecture01.pdf\")\n",
    "pages = loader.load()\n",
    "all_page_text=[p.page_content for p in pages]\n",
    "joined_page_text=\" \".join(all_page_text)\n",
    "\n",
    "# Split\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size = 1500,chunk_overlap = 150)\n",
    "splits = text_splitter.split_text(joined_page_text)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bb0d781",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Retrieve\n",
    "svm_retriever = SVMRetriever.from_texts(splits,embedding)\n",
    "tfidf_retriever = TFIDFRetriever.from_texts(splits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b1cc35f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "question = \"What are major topics for this class?\"\n",
    "docs_svm=svm_retriever.get_relevant_documents(question)\n",
    "docs_svm[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a1659c0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "question = \"what did they say about matlab?\"\n",
    "docs_tfidf=tfidf_retriever.get_relevant_documents(question)\n",
    "docs_tfidf[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7885389e",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}