diff --git a/templates/rag-multi-modal/README.md b/templates/rag-multi-modal/README.md index e8329223d60..f547f59aff5 100644 --- a/templates/rag-multi-modal/README.md +++ b/templates/rag-multi-modal/README.md @@ -2,17 +2,31 @@ This template performs RAG on documents with images. -It is confiured to ingest a pdf file that contains images: +It is configured to ingest a pdf file that contains images: * It uses uses [Unstructured](https://unstructured-io.github.io/unstructured/) for pdf parsing. * It uses [Chroma](https://www.trychroma.com/) for stroage. +The file is supplied in `docs/` and set in `chain.py`: + +``` +fpath = "../docs/" +fname = "cj.pdf" +``` + +By defaut it runs on a `.pdf` of [this blog post](https://cloudedjudgement.substack.com/p/clouded-judgement-111023) + ## Environment Setup Set the `OPENAI_API_KEY` environment variable to access the OpenAI models. [Unstructured](https://unstructured-io.github.io/unstructured/) requires some system-level package installations: +You will also need these in your system: + +* `poppler` ([installation instructions](https://pdf2image.readthedocs.io/en/latest/installation.html)) +* `tesseract` ([installation instructions](https://tesseract-ocr.github.io/tessdoc/Installation.html)) + On Mac, you can install the necessary packages with the following: ```shell diff --git a/templates/rag-multi-modal/docs/figure-1-1.jpg b/templates/rag-multi-modal/docs/figure-1-1.jpg deleted file mode 100644 index 2ed3804931d..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-1-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-10-1.jpg b/templates/rag-multi-modal/docs/figure-10-1.jpg deleted file mode 100644 index 1d0be959972..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-10-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-11-1.jpg b/templates/rag-multi-modal/docs/figure-11-1.jpg deleted file mode 100644 index d6e34a1fc3a..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-11-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-12-1.jpg b/templates/rag-multi-modal/docs/figure-12-1.jpg deleted file mode 100644 index 48ed8d2feb6..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-12-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-13-1.jpg b/templates/rag-multi-modal/docs/figure-13-1.jpg deleted file mode 100644 index b46524cc48c..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-13-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-14-1.jpg b/templates/rag-multi-modal/docs/figure-14-1.jpg deleted file mode 100644 index 11d920ff1b1..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-14-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-15-1.jpg b/templates/rag-multi-modal/docs/figure-15-1.jpg deleted file mode 100644 index c4d5dd1f71d..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-15-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-16-1.jpg b/templates/rag-multi-modal/docs/figure-16-1.jpg deleted file mode 100644 index 12f15904863..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-16-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-18-1.jpg b/templates/rag-multi-modal/docs/figure-18-1.jpg deleted file mode 100644 index 21c1aac7423..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-18-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-18-2.jpg b/templates/rag-multi-modal/docs/figure-18-2.jpg deleted file mode 100644 index 7ac84ffa2e7..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-18-2.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-19-1.jpg b/templates/rag-multi-modal/docs/figure-19-1.jpg deleted file mode 100644 index db7293315a7..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-19-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-20-1.jpg b/templates/rag-multi-modal/docs/figure-20-1.jpg deleted file mode 100644 index 1c27dbba704..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-20-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-20-2.jpg b/templates/rag-multi-modal/docs/figure-20-2.jpg deleted file mode 100644 index 493c6b6a3e5..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-20-2.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-20-3.jpg b/templates/rag-multi-modal/docs/figure-20-3.jpg deleted file mode 100644 index 3812a77605b..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-20-3.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-20-4.jpg b/templates/rag-multi-modal/docs/figure-20-4.jpg deleted file mode 100644 index ab7cb762c0a..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-20-4.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-20-5.jpg b/templates/rag-multi-modal/docs/figure-20-5.jpg deleted file mode 100644 index 25615744330..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-20-5.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-20-6.jpg b/templates/rag-multi-modal/docs/figure-20-6.jpg deleted file mode 100644 index ea558e96bb1..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-20-6.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-20-7.jpg b/templates/rag-multi-modal/docs/figure-20-7.jpg deleted file mode 100644 index 2a0cbb2a216..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-20-7.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-20-8.jpg b/templates/rag-multi-modal/docs/figure-20-8.jpg deleted file mode 100644 index 2a0cbb2a216..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-20-8.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-5-1.jpg b/templates/rag-multi-modal/docs/figure-5-1.jpg deleted file mode 100644 index b84c3f4ee8d..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-5-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-6-1.jpg b/templates/rag-multi-modal/docs/figure-6-1.jpg deleted file mode 100644 index 3c4c0a95b6a..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-6-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-7-1.jpg b/templates/rag-multi-modal/docs/figure-7-1.jpg deleted file mode 100644 index c637b0b6ebe..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-7-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-8-1.jpg b/templates/rag-multi-modal/docs/figure-8-1.jpg deleted file mode 100644 index 45147696db8..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-8-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/docs/figure-9-1.jpg b/templates/rag-multi-modal/docs/figure-9-1.jpg deleted file mode 100644 index 028ce928f04..00000000000 Binary files a/templates/rag-multi-modal/docs/figure-9-1.jpg and /dev/null differ diff --git a/templates/rag-multi-modal/rag_multi_modal/chain.py b/templates/rag-multi-modal/rag_multi_modal/chain.py index 9d5d145cbdd..7473095ea75 100644 --- a/templates/rag-multi-modal/rag_multi_modal/chain.py +++ b/templates/rag-multi-modal/rag_multi_modal/chain.py @@ -3,13 +3,11 @@ import io import os import re import uuid -from base64 import b64decode -from io import BytesIO -from PIL import Image from langchain.chat_models import ChatOpenAI from langchain.embeddings import OpenAIEmbeddings from langchain.prompts import ChatPromptTemplate +from langchain.pydantic_v1 import BaseModel from langchain.retrievers.multi_vector import MultiVectorRetriever from langchain.schema.document import Document from langchain.schema.messages import HumanMessage @@ -17,10 +15,12 @@ from langchain.schema.output_parser import StrOutputParser from langchain.schema.runnable import RunnableLambda, RunnablePassthrough from langchain.storage import InMemoryStore from langchain.vectorstores import Chroma +from PIL import Image from unstructured.partition.pdf import partition_pdf + # Extract elements from PDF -def extract_pdf_elements(path,fname): +def extract_pdf_elements(path, fname): """ Extract images, tables, and chunk text from a PDF file. path: File path, which is used to dump image files @@ -37,6 +37,7 @@ def extract_pdf_elements(path,fname): image_output_dir_path=path, ) + # Categorize elements by type def categorize_elements(raw_pdf_elements): """ @@ -52,6 +53,7 @@ def categorize_elements(raw_pdf_elements): texts.append(str(element)) return texts, tables + # Generate summaries of text elements def generate_text_summaries(texts, tables, summarize_texts=False): """ @@ -62,11 +64,12 @@ def generate_text_summaries(texts, tables, summarize_texts=False): """ # Prompt - prompt_text = """You are an assistant tasked with summarizing tables and text for retrieval. \ - These summaries will be embedded and used to retrieve the raw text or table elements. \ - Give a concise summary of the table or text that is well optimized for retrieval. Table or text: {element} """ + prompt_text = """You are an assistant tasked with summarizing tables and text for \ + retrieval. These summaries will be embedded and used to retrieve the raw text or \ + table elements. Give a concise summary of the table or text that is well \ + optimized for retrieval. Table or text: {element} """ prompt = ChatPromptTemplate.from_template(prompt_text) - + # Text summary chain model = ChatOpenAI(temperature=0, model="gpt-4") summarize_chain = {"element": lambda x: x} | prompt | model | StrOutputParser() @@ -79,7 +82,9 @@ def generate_text_summaries(texts, tables, summarize_texts=False): if texts and summarize_texts: text_summaries = summarize_chain.batch(texts, {"max_concurrency": 5}) elif texts: - text_summaries = texts # Directly assign texts if summarization is not requested + text_summaries = ( + texts + ) # Directly assign texts if summarization is not requested # Apply to tables if tables are provided if tables: @@ -87,6 +92,7 @@ def generate_text_summaries(texts, tables, summarize_texts=False): return text_summaries, table_summaries + def encode_image(image_path): """Getting the base64 string""" with open(image_path, "rb") as image_file: @@ -112,6 +118,7 @@ def image_summarize(img_base64, prompt): ) return msg.content + def generate_img_summaries(path): """ Generate summaries and base64 encoded strings for images @@ -120,15 +127,15 @@ def generate_img_summaries(path): # Store base64 encoded images img_base64_list = [] - + # Store image summaries image_summaries = [] - + # Prompt prompt = """You are an assistant tasked with summarizing images for retrieval. \ These summaries will be embedded and used to retrieve the raw image. \ Give a concise summary of the image that is well optimized for retrieval.""" - + # Apply to images for img_file in sorted(os.listdir(path)): if img_file.endswith(".jpg"): @@ -139,9 +146,14 @@ def generate_img_summaries(path): return img_base64_list, image_summaries + def create_multi_vector_retriever( vectorstore, text_summaries, texts, table_summaries, tables, image_summaries, images ): + """ + Create retriever that indexes summaries, but returns raw images or texts + """ + # Initialize the storage layer store = InMemoryStore() id_key = "doc_id" @@ -176,13 +188,18 @@ def create_multi_vector_retriever( return retriever + def looks_like_base64(sb): - """Check if the string looks like base64.""" + """ + Check if the string looks like base64. + """ return re.match("^[A-Za-z0-9+/]+[=]{0,2}$", sb) is not None def is_image_data(b64data): - """Check if the base64 data is an image by looking at the start of the data.""" + """ + Check if the base64 data is an image by looking at the start of the data. + """ image_signatures = { b"\xFF\xD8\xFF": "jpg", b"\x89\x50\x4E\x47\x0D\x0A\x1A\x0A": "png", @@ -199,8 +216,31 @@ def is_image_data(b64data): return False +def resize_base64_image(base64_string, size=(128, 128)): + """ + Resize an image encoded as a Base64 string. + base64_string (str): Base64 string of the original image. + size (tuple): Desired size of the image as (width, height). + """ + # Decode the Base64 string + img_data = base64.b64decode(base64_string) + img = Image.open(io.BytesIO(img_data)) + + # Resize the image + resized_img = img.resize(size, Image.LANCZOS) + + # Save the resized image to a bytes buffer + buffered = io.BytesIO() + resized_img.save(buffered, format=img.format) + + # Encode the resized image to Base64 + return base64.b64encode(buffered.getvalue()).decode("utf-8") + + def split_image_text_types(docs): - """Split base64-encoded images and texts.""" + """ + Split base64-encoded images and texts. + """ b64_images = [] texts = [] for doc in docs: @@ -208,6 +248,7 @@ def split_image_text_types(docs): if isinstance(doc, Document): doc = doc.page_content if looks_like_base64(doc) and is_image_data(doc): + doc = resize_base64_image(doc, size=(250, 250)) b64_images.append(doc) else: texts.append(doc) @@ -224,9 +265,7 @@ def img_prompt_func(data_dict): for image in data_dict["context"]["images"]: image_message = { "type": "image_url", - "image_url": { - "url": f"data:image/jpeg;base64,{image}" - }, + "image_url": {"url": f"data:image/jpeg;base64,{image}"}, } messages.append(image_message) @@ -234,8 +273,9 @@ def img_prompt_func(data_dict): text_message = { "type": "text", "text": ( - "Answer the question based only on the provided context, which can include text, tables, and image(s). " - "If an image is provided, analyze it carefully to help answer the question.\n" + "Answer the question based only on the provided context, " + "which can include text, tables, and image(s). If an image is " + "provided, analyze it carefully to help answer the question.\n" f"User-provided question / keywords: {data_dict['question']}\n\n" "Text and / or tables:\n" f"{formatted_texts}" @@ -244,8 +284,11 @@ def img_prompt_func(data_dict): messages.append(text_message) return [HumanMessage(content=messages)] + def multi_modal_rag_chain(retriever): - """Multi-modal RAG chain""" + """ + Multi-modal RAG chain + """ # Multi-modal LLM model = ChatOpenAI(temperature=0, model="gpt-4-vision-preview", max_tokens=1024) @@ -263,13 +306,13 @@ def multi_modal_rag_chain(retriever): return chain + # File path fpath = "../docs/" fname = "cj.pdf" # Get elements -raw_pdf_elements=extract_pdf_elements(fpath, - fname) +raw_pdf_elements = extract_pdf_elements(fpath, fname) # Get text, tables texts, tables = categorize_elements(raw_pdf_elements) @@ -281,8 +324,7 @@ img_base64_list, image_summaries = generate_img_summaries(fpath) # The vectorstore to use to index the summaries vectorstore = Chroma( - collection_name="multi_vector_img", - embedding_function=OpenAIEmbeddings() + collection_name="multi_vector_img", embedding_function=OpenAIEmbeddings() ) # Create retriever @@ -297,11 +339,12 @@ retriever_multi_vector_img = create_multi_vector_retriever( ) # Create RAG chain -chain_multimodal_rag = multi_modal_rag_chain(retriever_multi_vector_img) +chain = multi_modal_rag_chain(retriever_multi_vector_img) + # Add typing for input class Question(BaseModel): __root__: str -chain_multimodal_rag = multi_modal_rag_chain(retriever_multi_vector_img) + chain = chain.with_types(input_type=Question) diff --git a/templates/rag-multi-modal/rag_multi_modal/test.ipynb b/templates/rag-multi-modal/rag_multi_modal/test.ipynb deleted file mode 100644 index 92a067d4a7f..00000000000 --- a/templates/rag-multi-modal/rag_multi_modal/test.ipynb +++ /dev/null @@ -1,440 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": 21, - "id": "f6b055fe-d7e2-4d40-98ee-7648c8fdeb89", - "metadata": {}, - "outputs": [], - "source": [ - "from langchain.chat_models import ChatOpenAI\n", - "from langchain.prompts import ChatPromptTemplate\n", - "from unstructured.partition.pdf import partition_pdf\n", - "from langchain.schema.output_parser import StrOutputParser\n", - "\n", - "import base64\n", - "import io, os\n", - "from PIL import Image\n", - "from io import BytesIO\n", - "from langchain.schema.messages import HumanMessage\n", - "\n", - "import uuid\n", - "from base64 import b64decode\n", - "\n", - "from langchain.embeddings import OpenAIEmbeddings\n", - "from langchain.vectorstores import Chroma\n", - "\n", - "from langchain.retrievers.multi_vector import MultiVectorRetriever\n", - "from langchain.schema.document import Document\n", - "from langchain.storage import InMemoryStore\n", - "\n", - "import uuid\n", - "from base64 import b64decode\n", - "\n", - "from langchain.embeddings import OpenAIEmbeddings\n", - "from langchain.vectorstores import Chroma\n", - "\n", - "from langchain.retrievers.multi_vector import MultiVectorRetriever\n", - "from langchain.schema.document import Document\n", - "from langchain.storage import InMemoryStore\n", - "\n", - "import re\n", - "\n", - "from langchain.schema import Document\n", - "from langchain.schema.runnable import RunnableLambda\n", - "from langchain.schema.runnable import RunnablePassthrough\n" - ] - }, - { - "cell_type": "code", - "execution_count": 22, - "id": "a9807206-b972-4eff-8f40-d263219cb426", - "metadata": {}, - "outputs": [], - "source": [ - "# Extract elements from PDF\n", - "def extract_pdf_elements(path,fname):\n", - " \"\"\"\n", - " Extract images, tables, and chunk text from a PDF file.\n", - " path: File path, which is used to dump image files\n", - " fname: File name\n", - " \"\"\"\n", - " partition_pdf(\n", - " filename=path + fname,\n", - " extract_images_in_pdf=True,\n", - " infer_table_structure=True,\n", - " chunking_strategy=\"by_title\",\n", - " max_characters=4000,\n", - " new_after_n_chars=3800,\n", - " combine_text_under_n_chars=2000,\n", - " image_output_dir_path=path,\n", - " )\n", - "\n", - "# Categorize elements by type\n", - "def categorize_elements(raw_pdf_elements):\n", - " \"\"\"\n", - " Categorize extracted elements from a PDF into tables and texts.\n", - " raw_pdf_elements: List of unstructured.documents.elements\n", - " \"\"\"\n", - " tables = []\n", - " texts = []\n", - " for element in raw_pdf_elements:\n", - " if \"unstructured.documents.elements.Table\" in str(type(element)):\n", - " tables.append(str(element))\n", - " elif \"unstructured.documents.elements.CompositeElement\" in str(type(element)):\n", - " texts.append(str(element))\n", - " return texts, tables\n", - "\n", - "# Generate summaries of text elements\n", - "def generate_text_summaries(texts, tables, summarize_texts=False):\n", - " \"\"\"\n", - " Summarize text elements\n", - " texts: List of str\n", - " tables: List of str\n", - " summarize_texts: Bool to summarize texts\n", - " \"\"\"\n", - "\n", - " # Prompt\n", - " prompt_text = \"\"\"You are an assistant tasked with summarizing tables and text for retrieval. \\\n", - " These summaries will be embedded and used to retrieve the raw text or table elements. \\\n", - " Give a concise summary of the table or text that is well optimized for retrieval. Table or text: {element} \"\"\"\n", - " prompt = ChatPromptTemplate.from_template(prompt_text)\n", - " \n", - " # Text summary chain\n", - " model = ChatOpenAI(temperature=0, model=\"gpt-4\")\n", - " summarize_chain = {\"element\": lambda x: x} | prompt | model | StrOutputParser()\n", - "\n", - " # Initialize empty summaries\n", - " text_summaries = []\n", - " table_summaries = []\n", - "\n", - " # Apply to text if texts are provided and summarization is requested\n", - " if texts and summarize_texts:\n", - " text_summaries = summarize_chain.batch(texts, {\"max_concurrency\": 5})\n", - " elif texts:\n", - " text_summaries = texts # Directly assign texts if summarization is not requested\n", - "\n", - " # Apply to tables if tables are provided\n", - " if tables:\n", - " table_summaries = summarize_chain.batch(tables, {\"max_concurrency\": 5})\n", - "\n", - " return text_summaries, table_summaries\n", - "\n", - "def encode_image(image_path):\n", - " \"\"\"Getting the base64 string\"\"\"\n", - " with open(image_path, \"rb\") as image_file:\n", - " return base64.b64encode(image_file.read()).decode(\"utf-8\")\n", - "\n", - "\n", - "def image_summarize(img_base64, prompt):\n", - " \"\"\"Make image summary\"\"\"\n", - " chat = ChatOpenAI(model=\"gpt-4-vision-preview\", max_tokens=1024)\n", - "\n", - " msg = chat.invoke(\n", - " [\n", - " HumanMessage(\n", - " content=[\n", - " {\"type\": \"text\", \"text\": prompt},\n", - " {\n", - " \"type\": \"image_url\",\n", - " \"image_url\": {\"url\": f\"data:image/jpeg;base64,{img_base64}\"},\n", - " },\n", - " ]\n", - " )\n", - " ]\n", - " )\n", - " return msg.content\n", - "\n", - "def generate_img_summaries(path):\n", - " \"\"\"\n", - " Generate summaries and base64 encoded strings for images\n", - " path: Path to list of .jpg files extracted by Unstructured\n", - " \"\"\"\n", - "\n", - " # Store base64 encoded images\n", - " img_base64_list = []\n", - " \n", - " # Store image summaries\n", - " image_summaries = []\n", - " \n", - " # Prompt\n", - " prompt = \"\"\"You are an assistant tasked with summarizing images for retrieval. \\\n", - " These summaries will be embedded and used to retrieve the raw image. \\\n", - " Give a concise summary of the image that is well optimized for retrieval.\"\"\"\n", - " \n", - " # Apply to images\n", - " for img_file in sorted(os.listdir(path)):\n", - " if img_file.endswith(\".jpg\"):\n", - " img_path = os.path.join(path, img_file)\n", - " base64_image = encode_image(img_path)\n", - " img_base64_list.append(base64_image)\n", - " image_summaries.append(image_summarize(base64_image, prompt))\n", - "\n", - " return img_base64_list, image_summaries\n", - "\n", - "def create_multi_vector_retriever(\n", - " vectorstore, text_summaries, texts, table_summaries, tables, image_summaries, images\n", - "):\n", - " # Initialize the storage layer\n", - " store = InMemoryStore()\n", - " id_key = \"doc_id\"\n", - "\n", - " # Create the multi-vector retriever\n", - " retriever = MultiVectorRetriever(\n", - " vectorstore=vectorstore,\n", - " docstore=store,\n", - " id_key=id_key,\n", - " )\n", - "\n", - " # Helper function to add documents to the vectorstore and docstore\n", - " def add_documents(retriever, doc_summaries, doc_contents):\n", - " doc_ids = [str(uuid.uuid4()) for _ in doc_contents]\n", - " summary_docs = [\n", - " Document(page_content=s, metadata={id_key: doc_ids[i]})\n", - " for i, s in enumerate(doc_summaries)\n", - " ]\n", - " retriever.vectorstore.add_documents(summary_docs)\n", - " retriever.docstore.mset(list(zip(doc_ids, doc_contents)))\n", - "\n", - " # Add texts, tables, and images\n", - " # Check that text_summaries is not empty before adding\n", - " if text_summaries:\n", - " add_documents(retriever, text_summaries, texts)\n", - " # Check that table_summaries is not empty before adding\n", - " if table_summaries:\n", - " add_documents(retriever, table_summaries, tables)\n", - " # Check that image_summaries is not empty before adding\n", - " if image_summaries:\n", - " add_documents(retriever, image_summaries, images)\n", - "\n", - " return retriever\n", - "\n", - "def looks_like_base64(sb):\n", - " \"\"\"Check if the string looks like base64.\"\"\"\n", - " return re.match(\"^[A-Za-z0-9+/]+[=]{0,2}$\", sb) is not None\n", - "\n", - "\n", - "def is_image_data(b64data):\n", - " \"\"\"Check if the base64 data is an image by looking at the start of the data.\"\"\"\n", - " image_signatures = {\n", - " b\"\\xFF\\xD8\\xFF\": \"jpg\",\n", - " b\"\\x89\\x50\\x4E\\x47\\x0D\\x0A\\x1A\\x0A\": \"png\",\n", - " b\"\\x47\\x49\\x46\\x38\": \"gif\",\n", - " b\"\\x52\\x49\\x46\\x46\": \"webp\",\n", - " }\n", - " try:\n", - " header = base64.b64decode(b64data)[:8] # Decode and get the first 8 bytes\n", - " for sig, format in image_signatures.items():\n", - " if header.startswith(sig):\n", - " return True\n", - " return False\n", - " except Exception:\n", - " return False\n", - "\n", - "\n", - "def split_image_text_types(docs):\n", - " \"\"\"Split base64-encoded images and texts.\"\"\"\n", - " b64_images = []\n", - " texts = []\n", - " for doc in docs:\n", - " # Check if the document is of type Document and extract page_content if so\n", - " if isinstance(doc, Document):\n", - " doc = doc.page_content\n", - " if looks_like_base64(doc) and is_image_data(doc):\n", - " b64_images.append(doc)\n", - " else:\n", - " texts.append(doc)\n", - " return {\"images\": b64_images, \"texts\": texts}\n", - "\n", - "\n", - "def img_prompt_func(data_dict):\n", - " # Joining the context texts into a single string\n", - " formatted_texts = \"\\n\".join(data_dict[\"context\"][\"texts\"])\n", - " messages = []\n", - "\n", - " # Adding image(s) to the messages if present\n", - " if data_dict[\"context\"][\"images\"]:\n", - " for image in data_dict[\"context\"][\"images\"]:\n", - " image_message = {\n", - " \"type\": \"image_url\",\n", - " \"image_url\": {\n", - " \"url\": f\"data:image/jpeg;base64,{image}\"\n", - " },\n", - " }\n", - " messages.append(image_message)\n", - "\n", - " # Adding the text message for analysis\n", - " text_message = {\n", - " \"type\": \"text\",\n", - " \"text\": (\n", - " \"Answer the question based only on the provided context, which can include text, tables, and image(s). \"\n", - " \"If an image is provided, analyze it carefully to help answer the question.\\n\"\n", - " f\"User-provided question / keywords: {data_dict['question']}\\n\\n\"\n", - " \"Text and / or tables:\\n\"\n", - " f\"{formatted_texts}\"\n", - " ),\n", - " }\n", - " messages.append(text_message)\n", - " return [HumanMessage(content=messages)]\n", - "\n", - "def multi_modal_rag_chain(retriever):\n", - " \"\"\"Multi-modal RAG chain\"\"\"\n", - "\n", - " # Multi-modal LLM\n", - " model = ChatOpenAI(temperature=0, model=\"gpt-4-vision-preview\", max_tokens=1024)\n", - "\n", - " # RAG pipeline\n", - " chain = (\n", - " {\n", - " \"context\": retriever | RunnableLambda(split_image_text_types),\n", - " \"question\": RunnablePassthrough(),\n", - " }\n", - " | RunnableLambda(img_prompt_func)\n", - " | model\n", - " | StrOutputParser()\n", - " )\n", - "\n", - " return chain" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "d701a0a2-0de6-49a9-9552-55a1d6656ae4", - "metadata": {}, - "outputs": [], - "source": [ - "# File path\n", - "fpath = \"../docs/\"\n", - "fname = \"cj.pdf\"\n", - "\n", - "# Get elements\n", - "raw_pdf_elements=extract_pdf_elements(fpath,\n", - " fname)\n", - "# Get text, tables\n", - "texts, tables = categorize_elements(raw_pdf_elements)\n", - "\n", - "# Get text, table summaries\n", - "text_summaries, table_summaries = generate_text_summaries(texts, tables)\n", - "\n", - "# Image summaries\n", - "img_base64_list, image_summaries = generate_img_summaries(fpath)\n", - "\n", - "# The vectorstore to use to index the summaries\n", - "vectorstore = Chroma(\n", - " collection_name=\"multi_vector_img\", \n", - " embedding_function=OpenAIEmbeddings()\n", - ")\n", - "\n", - "# Create retriever\n", - "retriever_multi_vector_img = create_multi_vector_retriever(\n", - " vectorstore,\n", - " text_summaries,\n", - " texts,\n", - " table_summaries,\n", - " tables,\n", - " image_summaries,\n", - " img_base64_list,\n", - ")\n", - "\n", - "# Create RAG chain\n", - "chain_multimodal_rag = multi_modal_rag_chain(retriever_multi_vector_img)" - ] - }, - { - "cell_type": "code", - "execution_count": 17, - "id": "efd34322-bc82-4549-b795-3e0077072d4c", - "metadata": {}, - "outputs": [], - "source": [ - "# Testing on retrieval\n", - "query = \"What is the change in reported revenue for Squarespace?\"\n", - "suffix_for_images = \" Include any pie charts, graphs, or tables.\"\n", - "docs = retriever_multi_vector_img.get_relevant_documents(query + suffix_for_images)" - ] - }, - { - "cell_type": "code", - "execution_count": 20, - "id": "0a6276fa-b8ce-466b-bba6-f11aa45efdba", - "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "from IPython.display import HTML, display\n", - "\n", - "def plt_img_base64(img_base64):\n", - " # Create an HTML img tag with the base64 string as the source\n", - " image_html = f''\n", - "\n", - " # Display the image by rendering the HTML\n", - " display(HTML(image_html))\n", - "\n", - "plt_img_base64(docs[0])" - ] - }, - { - "cell_type": "code", - "execution_count": 24, - "id": "0bd73d74-4616-4066-9577-2f4d917c8371", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "'The change in reported revenue for Squarespace is 2.1%. This is indicated in the table under the \"Reported Revenue\" section, where the \"Actual\" revenue is $257.1M and the \"Consensus\" was $251.8M, resulting in a delta (Δ) of 2.1%.'" - ] - }, - "execution_count": 24, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "chain_multimodal_rag.invoke(\"What is the change in reported revenue for Squarespace?\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "19269350-1047-4b5e-b128-694b6f54f736", - "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.9.16" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -}