{
 "cells": [
  {
   "attachments": {
    "7b5c5a30-393c-4b27-8fa1-688306ef2aef.png": {
     "image/png": "[encoded data, 131504 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "b6d466cc-aa8b-4baf-a80a-fef01921ca8d",
   "metadata": {},
   "source": [
    "## Semi-structured RAG\n",
    "\n",
    "Many documents contain a mixture of content types, including text and tables. \n",
    "\n",
    "Semi-structured data can be challenging for conventional RAG for at least two reasons: \n",
    "\n",
    "* Text splitting may break up tables, corrupting the data in retrieval\n",
    "* Embedding tables may pose challenges for semantic similarity search \n",
    "\n",
    "This cookbook shows how to perform RAG on documents with semi-structured data: \n",
    "\n",
    "* We will use [Unstructured](https://unstructured.io/) to parse both text and tables from documents (PDFs).\n",
    "* We will use the [multi-vector retriever](https://python.langchain.com/docs/modules/data_connection/retrievers/multi_vector) to store raw tables, text along with table summaries better suited for retrieval.\n",
    "* We will use [LCEL](https://python.langchain.com/docs/expression_language/) to implement the chains used.\n",
    "\n",
    "The overall flow is here:\n",
    "\n",
    "![MVR.png](attachment:7b5c5a30-393c-4b27-8fa1-688306ef2aef.png)\n",
    "\n",
    "## Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5740fc70-c513-4ff4-9d72-cfc098f85fef",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install langchain unstructured[all-docs] pydantic lxml langchainhub"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44349a83-e1dc-4eed-ba75-587f309d8c88",
   "metadata": {},
   "source": [
    "The PDF partitioning used by Unstructured will use: \n",
    "\n",
    "* `tesseract` for Optical Character Recognition (OCR)\n",
    "*  `poppler` for PDF rendering and processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7880871-4949-4ea2-aed8-540a09188a41",
   "metadata": {},
   "outputs": [],
   "source": [
    "! brew install tesseract\n",
    "! brew install poppler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c24efa9-b6f6-4dc2-bfe3-70819ba3ef75",
   "metadata": {},
   "source": [
    "## Data Loading\n",
    "\n",
    "### Partition PDF tables and text\n",
    "\n",
    "Apply to the [`LLaMA2`](https://arxiv.org/pdf/2307.09288.pdf) paper. \n",
    "\n",
    "We use the Unstructured [`partition_pdf`](https://unstructured-io.github.io/unstructured/bricks/partition.html#partition-pdf), which segments a PDF document by using a layout model. \n",
    "\n",
    "This layout model makes it possible to extract elements, such as tables, from pdfs. \n",
    "\n",
    "We also can use `Unstructured` chunking, which:\n",
    "\n",
    "* Tries to identify document sections (e.g., Introduction, etc)\n",
    "* Then, builds text blocks that maintain sections while also honoring user-defined chunk sizes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "62cf502b-407d-4645-a72c-24498fd55130",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"/Users/rlm/Desktop/Papers/LLaMA2/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3867a654-61ba-4759-9a64-de953a429ced",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any\n",
    "\n",
    "from pydantic import BaseModel\n",
    "from unstructured.partition.pdf import partition_pdf\n",
    "\n",
    "# Get elements\n",
    "raw_pdf_elements = partition_pdf(\n",
    "    filename=path + \"LLaMA2.pdf\",\n",
    "    # Unstructured first finds embedded image blocks\n",
    "    extract_images_in_pdf=False,\n",
    "    # Use layout model (YOLOX) to get bounding boxes (for tables) and find titles\n",
    "    # Titles are any sub-section of the document\n",
    "    infer_table_structure=True,\n",
    "    # Post processing to aggregate text once we have the title\n",
    "    chunking_strategy=\"by_title\",\n",
    "    # Chunking params to aggregate text blocks\n",
    "    # Attempt to create a new chunk 3800 chars\n",
    "    # Attempt to keep chunks > 2000 chars\n",
    "    max_characters=4000,\n",
    "    new_after_n_chars=3800,\n",
    "    combine_text_under_n_chars=2000,\n",
    "    image_output_dir_path=path,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b09cd727-aeab-49af-8a51-0dc377321e7c",
   "metadata": {},
   "source": [
    "We can examine the elements extracted by `partition_pdf`.\n",
    "\n",
    "`CompositeElement` are aggregated chunks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "628abfc6-4057-434b-b880-d88e3ba44657",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{\"<class 'unstructured.documents.elements.CompositeElement'>\": 184,\n",
       " \"<class 'unstructured.documents.elements.Table'>\": 47,\n",
       " \"<class 'unstructured.documents.elements.TableChunk'>\": 2}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create a dictionary to store counts of each type\n",
    "category_counts = {}\n",
    "\n",
    "for element in raw_pdf_elements:\n",
    "    category = str(type(element))\n",
    "    if category in category_counts:\n",
    "        category_counts[category] += 1\n",
    "    else:\n",
    "        category_counts[category] = 1\n",
    "\n",
    "# Unique_categories will have unique elements\n",
    "unique_categories = set(category_counts.keys())\n",
    "category_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "5462f29e-fd59-4e0e-9493-ea3b560e523e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "49\n",
      "184\n"
     ]
    }
   ],
   "source": [
    "class Element(BaseModel):\n",
    "    type: str\n",
    "    text: Any\n",
    "\n",
    "\n",
    "# Categorize by type\n",
    "categorized_elements = []\n",
    "for element in raw_pdf_elements:\n",
    "    if \"unstructured.documents.elements.Table\" in str(type(element)):\n",
    "        categorized_elements.append(Element(type=\"table\", text=str(element)))\n",
    "    elif \"unstructured.documents.elements.CompositeElement\" in str(type(element)):\n",
    "        categorized_elements.append(Element(type=\"text\", text=str(element)))\n",
    "\n",
    "# Tables\n",
    "table_elements = [e for e in categorized_elements if e.type == \"table\"]\n",
    "print(len(table_elements))\n",
    "\n",
    "# Text\n",
    "text_elements = [e for e in categorized_elements if e.type == \"text\"]\n",
    "print(len(text_elements))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "731b3dfc-7ddf-4a11-9a30-9a79b7c66e16",
   "metadata": {},
   "source": [
    "## Multi-vector retriever\n",
    "\n",
    "Use [multi-vector-retriever](https://python.langchain.com/docs/modules/data_connection/retrievers/multi_vector#summary) to produce summaries of tables and, optionally, text. \n",
    "\n",
    "With the summary, we will also store the raw table elements.\n",
    "\n",
    "The summaries are used to improve the quality of retrieval, [as explained in the multi vector retriever docs](https://python.langchain.com/docs/modules/data_connection/retrievers/multi_vector).\n",
    "\n",
    "The raw tables are passed to the LLM, providing the full table context for the LLM to generate the answer.  \n",
    "\n",
    "### Summaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "8e275736-3408-4d7a-990e-4362c88e81f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.prompts import ChatPromptTemplate\n",
    "from langchain.schema.output_parser import StrOutputParser"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37b65677-aeb4-44fd-b06d-4539341ede97",
   "metadata": {},
   "source": [
    "We create a simple summarize chain for each element.\n",
    "\n",
    "You can also see, re-use, or modify the prompt in the Hub [here](https://smith.langchain.com/hub/rlm/multi-vector-retriever-summarization).\n",
    "\n",
    "```\n",
    "from langchain import hub\n",
    "obj = hub.pull(\"rlm/multi-vector-retriever-summarization\")\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "1b12536a-1303-41ad-9948-4eb5a5f32614",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prompt\n",
    "prompt_text = \"\"\"You are an assistant tasked with summarizing tables and text. \\ \n",
    "Give a concise summary of the table or text. Table or text chunk: {element} \"\"\"\n",
    "prompt = ChatPromptTemplate.from_template(prompt_text)\n",
    "\n",
    "# Summary chain\n",
    "model = ChatOpenAI(temperature=0, model=\"gpt-4\")\n",
    "summarize_chain = {\"element\": lambda x: x} | prompt | model | StrOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d8b567c-b442-4bf0-b639-04bd89effc62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply to tables\n",
    "tables = [i.text for i in table_elements]\n",
    "table_summaries = summarize_chain.batch(tables, {\"max_concurrency\": 5})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "3e9c176c-3d46-4034-b169-0d7305d42d27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply to texts\n",
    "texts = [i.text for i in text_elements]\n",
    "text_summaries = summarize_chain.batch(texts, {\"max_concurrency\": 5})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60524010-754f-4924-ad75-78cb54ca7257",
   "metadata": {},
   "source": [
    "### Add to vectorstore\n",
    "\n",
    "Use [Multi Vector Retriever](https://python.langchain.com/docs/modules/data_connection/retrievers/multi_vector#summary) with summaries: \n",
    "\n",
    "* `InMemoryStore` stores the raw text, tables\n",
    "* `vectorstore` stores the embedded summaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "346c3a02-8fea-4f75-a69e-fc9542b99dbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "from langchain.retrievers.multi_vector import MultiVectorRetriever\n",
    "from langchain.schema.document import Document\n",
    "from langchain.storage import InMemoryStore\n",
    "from langchain.vectorstores import Chroma\n",
    "\n",
    "# The vectorstore to use to index the child chunks\n",
    "vectorstore = Chroma(collection_name=\"summaries\", embedding_function=OpenAIEmbeddings())\n",
    "\n",
    "# The storage layer for the parent documents\n",
    "store = InMemoryStore()\n",
    "id_key = \"doc_id\"\n",
    "\n",
    "# The retriever (empty to start)\n",
    "retriever = MultiVectorRetriever(\n",
    "    vectorstore=vectorstore,\n",
    "    docstore=store,\n",
    "    id_key=id_key,\n",
    ")\n",
    "\n",
    "# Add texts\n",
    "doc_ids = [str(uuid.uuid4()) for _ in texts]\n",
    "summary_texts = [\n",
    "    Document(page_content=s, metadata={id_key: doc_ids[i]})\n",
    "    for i, s in enumerate(text_summaries)\n",
    "]\n",
    "retriever.vectorstore.add_documents(summary_texts)\n",
    "retriever.docstore.mset(list(zip(doc_ids, texts)))\n",
    "\n",
    "# Add tables\n",
    "table_ids = [str(uuid.uuid4()) for _ in tables]\n",
    "summary_tables = [\n",
    "    Document(page_content=s, metadata={id_key: table_ids[i]})\n",
    "    for i, s in enumerate(table_summaries)\n",
    "]\n",
    "retriever.vectorstore.add_documents(summary_tables)\n",
    "retriever.docstore.mset(list(zip(table_ids, tables)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d8bbbd9-009b-4b34-a206-5874a60adbda",
   "metadata": {},
   "source": [
    "## RAG\n",
    "\n",
    "Run [RAG pipeline](https://python.langchain.com/docs/expression_language/cookbook/retrieval)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "f2489de4-51e3-48b4-bbcd-ed9171deadf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.schema.runnable import RunnablePassthrough\n",
    "\n",
    "# Prompt template\n",
    "template = \"\"\"Answer the question based only on the following context, which can include text and tables:\n",
    "{context}\n",
    "Question: {question}\n",
    "\"\"\"\n",
    "prompt = ChatPromptTemplate.from_template(template)\n",
    "\n",
    "# LLM\n",
    "model = ChatOpenAI(temperature=0, model=\"gpt-4\")\n",
    "\n",
    "# RAG pipeline\n",
    "chain = (\n",
    "    {\"context\": retriever, \"question\": RunnablePassthrough()}\n",
    "    | prompt\n",
    "    | model\n",
    "    | StrOutputParser()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "90e3d100-10e8-4ee6-ae46-2480b1524ec8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The number of training tokens for LLaMA2 is 2.0T.'"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke(\"What is the number of training tokens for LLaMA2?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37f46054-e239-4ba8-af81-22d0d6a9bc32",
   "metadata": {},
   "source": [
    "We can check the [trace](https://smith.langchain.com/public/4739ae7c-1a13-406d-bc4e-3462670ebc01/r) to see what chunks were retrieved:\n",
    "\n",
    "This includes Table 1 of the paper, showing the Tokens used for training.\n",
    "\n",
    "```\n",
    "Training Data Params Context GQA Tokens LR Length 7B 2k 1.0T 3.0x 10-4 See Touvron et al. 13B 2k 1.0T 3.0 x 10-4 LiaMa 1 (2023) 33B 2k 14T 1.5 x 10-4 65B 2k 1.4T 1.5 x 10-4 7B 4k 2.0T 3.0x 10-4 Liama 2 A new mix of publicly 13B 4k 2.0T 3.0 x 10-4 available online data 34B 4k v 2.0T 1.5 x 10-4 70B 4k v 2.0T 1.5 x 10-4\n",
    "```"
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
