Fmt

templates/rag-chroma-multi-modal-multi-vector/ingest.py

@@ -1,28 +1,18 @@
-import os
-from pathlib import Path
-
 import base64
 import io
+import uuid
 from io import BytesIO
-
-from PIL import Image
-
+from pathlib import Path
 
 import pypdfium2 as pdfium
-from langchain.vectorstores import Chroma
-from langchain_experimental.open_clip import OpenCLIPEmbeddings
-
-import uuid
-
+from langchain.chat_models import ChatOpenAI
 from langchain.embeddings import OpenAIEmbeddings
 from langchain.retrievers.multi_vector import MultiVectorRetriever
 from langchain.schema.document import Document
-from langchain.schema.output_parser import StrOutputParser
-from langchain.storage import UpstashRedisByteStore
-
-
-from langchain.chat_models import ChatOpenAI
 from langchain.schema.messages import HumanMessage
+from langchain.storage import UpstashRedisByteStore
+from langchain.vectorstores import Chroma
+from PIL import Image
 
 
 def image_summarize(img_base64, prompt):
@@ -80,12 +70,11 @@ def generate_img_summaries(img_base64_list):
     return image_summaries, processed_images
 
 
-def get_images_from_pdf(pdf_path, img_dump_path):
+def get_images_from_pdf(pdf_path):
     """
     Extract images from each page of a PDF document and save as JPEG files.
 
     :param pdf_path: A string representing the path to the PDF file.
-    :param img_dump_path: A string representing the path to dummp images.
     """
     pdf = pdfium.PdfDocument(pdf_path)
     n_pages = len(pdf)
@@ -94,7 +83,6 @@ def get_images_from_pdf(pdf_path, img_dump_path):
         page = pdf.get_page(page_number)
         bitmap = page.render(scale=1, rotation=0, crop=(0, 0, 0, 0))
         pil_image = bitmap.to_pil()
-        pil_image.save(f"{img_dump_path}/img_{page_number + 1}.jpg", format="JPEG")
         pil_images.append(pil_image)
     return pil_images
 
@@ -136,6 +124,7 @@ def convert_to_base64(pil_image):
     img_str = resize_base64_image(img_str, size=(960, 540))
     return img_str
 
+
 def create_multi_vector_retriever(vectorstore, image_summaries, images):
     """
     Create retriever that indexes summaries, but returns raw images or texts
@@ -147,10 +136,9 @@ def create_multi_vector_retriever(vectorstore, image_summaries, images):
     """
 
     # Initialize the storage layer for images
-    UPSTASH_URL = "https://usw1-bright-beagle-34178.upstash.io"
-    UPSTASH_TOKEN = "AYWCACQgNzk3OTJjZTItMGIxNy00MTEzLWIyZTAtZWI0ZmI1ZGY0NjFhNGRhMGZjNDE4YjgxNGE4MTkzOWYxMzllM2MzZThlOGY="
-    store = UpstashRedisByteStore(url=UPSTASH_URL,
-                                token=UPSTASH_TOKEN)
+    UPSTASH_URL = "xxx"
+    UPSTASH_TOKEN = "xxx"
+    store = UpstashRedisByteStore(url=UPSTASH_URL, token=UPSTASH_TOKEN)
     id_key = "doc_id"
 
     # Create the multi-vector retriever
@@ -174,6 +162,7 @@ def add_documents(retriever, doc_summaries, doc_contents):
 
     return retriever
 
+
 # Load PDF
 doc_path = Path(__file__).parent / "docs/DDOG_Q3_earnings_deck.pdf"
 img_dump_path = Path(__file__).parent / "docs/"
@@ -193,11 +182,13 @@ def add_documents(retriever, doc_summaries, doc_contents):
 vectorstore_mvr = Chroma(
     collection_name="image_summaries",
     persist_directory=str(Path(__file__).parent / "chroma_db_multi_modal"),
-    embedding_function=OpenAIEmbeddings()
+    embedding_function=OpenAIEmbeddings(),
 )
 
 # Create documents
-images_base_64_processed_documents = [Document(page_content=i) for i in images_base_64_processed]
+images_base_64_processed_documents = [
+    Document(page_content=i) for i in images_base_64_processed
+]
 
 # Create retriever
 retriever_multi_vector_img = create_multi_vector_retriever(

templates/rag-chroma-multi-modal-multi-vector/rag_chroma_multi_modal_multi_vector/chain.py

@@ -1,18 +1,17 @@
 import base64
 import io
-import json 
+import json
 from pathlib import Path
 
-from langchain.retrievers.multi_vector import MultiVectorRetriever
 from langchain.chat_models import ChatOpenAI
+from langchain.embeddings import OpenAIEmbeddings
 from langchain.pydantic_v1 import BaseModel
-from langchain.schema.document import Document
+from langchain.retrievers.multi_vector import MultiVectorRetriever
 from langchain.schema.messages import HumanMessage
 from langchain.schema.output_parser import StrOutputParser
 from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
-from langchain.vectorstores import Chroma
-from langchain.embeddings import OpenAIEmbeddings
 from langchain.storage import UpstashRedisByteStore
+from langchain.vectorstores import Chroma
 from PIL import Image
 
 
@@ -41,7 +40,7 @@ def get_resized_images(docs):
     """
     b64_images = []
     for doc in docs:
-        doc = json.loads(doc.decode('utf-8'))['kwargs']['page_content']
+        doc = json.loads(doc.decode("utf-8"))["kwargs"]["page_content"]
         resized_image = resize_base64_image(doc, size=(1280, 720))
         b64_images.append(resized_image)
     return {"images": b64_images}
@@ -104,14 +103,13 @@ def multi_modal_rag_chain(retriever):
 vectorstore_mvr = Chroma(
     collection_name="image_summaries",
     persist_directory=str(Path(__file__).parent.parent / "chroma_db_multi_modal"),
-    embedding_function=OpenAIEmbeddings()
+    embedding_function=OpenAIEmbeddings(),
 )
 
 # Load redis
-UPSTASH_URL = "https://usw1-bright-beagle-34178.upstash.io"
-UPSTASH_TOKEN = "AYWCACQgNzk3OTJjZTItMGIxNy00MTEzLWIyZTAtZWI0ZmI1ZGY0NjFhNGRhMGZjNDE4YjgxNGE4MTkzOWYxMzllM2MzZThlOGY="
-store = UpstashRedisByteStore(url=UPSTASH_URL,
-                            token=UPSTASH_TOKEN)
+UPSTASH_URL = "xxx"
+UPSTASH_TOKEN = "xxx"
+store = UpstashRedisByteStore(url=UPSTASH_URL, token=UPSTASH_TOKEN)
 id_key = "doc_id"
 
 # Create the multi-vector retriever
@@ -124,8 +122,10 @@ def multi_modal_rag_chain(retriever):
 # Create RAG chain
 chain = multi_modal_rag_chain(retriever)
 
+
 # Add typing for input
 class Question(BaseModel):
     __root__: str
 
+
 chain = chain.with_types(input_type=Question)