import osimport openaiimport ioimport uuidimport base64import timefrom base64 import b64decodeimport numpy as npfrom PIL import Image
from unstructured.partition.pdf import partition_pdf
from langchain.chat_models import ChatOpenAIfrom langchain.schema.messages import HumanMessage, SystemMessagefrom langchain.vectorstores import Chromafrom langchain.storage import InMemoryStorefrom langchain.schema.document import Documentfrom langchain.embeddings import OpenAIEmbeddingsfrom langchain.retrievers.multi_vector import MultiVectorRetrieverfrom langchain.chat_models import ChatOpenAIfrom langchain.prompts import ChatPromptTemplatefrom langchain.schema.output_parser import StrOutputParserfrom langchain.schema.runnable import RunnablePassthrough, RunnableLambda
from operator import itemgetter

# load the pdf file to drive# split the file to text, table and imagesdef doc_partition(path,file_name):raw_pdf_elements = partition_pdf(filename=path + file_name,extract_images_in_pdf=True,infer_table_structure=True,chunking_strategy="by_title",max_characters=4000,new_after_n_chars=3800,combine_text_under_n_chars=2000,image_output_dir_path=path)
return raw_pdf_elementspath = "/content/"file_name = "wildfire_stats.pdf"raw_pdf_elements = doc_partition(path,file_name)

# appending texts and tables from the pdf filedef data_category(raw_pdf_elements): # we may use decorator heretables = []texts = []for element in raw_pdf_elements:if "unstructured.documents.elements.Table" in str(type(element)): tables.append(str(element))elif "unstructured.documents.elements.CompositeElement" in str(type(element)): texts.append(str(element))data_category = [texts,tables]return data_categorytexts = data_category(raw_pdf_elements)[0]tables = data_category(raw_pdf_elements)[1]

os.environ["OPENAI_API_KEY"] = 'sk-xxxxxxxxxxxxxxx'openai.api_key = os.environ["OPENAI_API_KEY"]
# function to take tables as input and then summarize themdef tables_summarize(data_category):prompt_text = """You are an assistant tasked with summarizing tables. \\Give a concise summary of the table. Table chunk: {element} """
prompt = ChatPromptTemplate.from_template(prompt_text)model = ChatOpenAI(temperature=0, model="gpt-4")summarize_chain = {"element": lambda x: x} | prompt | model | StrOutputParser()table_summaries = summarize_chain.batch(tables, {"max_concurrency": 5})
return table_summariestable_summaries = tables_summarize(data_category)text_summaries = texts

def encode_image(image_path):''' Getting the base64 string '''with open(image_path, "rb") as image_file:return base64.b64encode(image_file.read()).decode('utf-8')
def image_captioning(img_base64,prompt):''' Image summary '''chat = ChatOpenAI(model="gpt-4-vision-preview",max_tokens=1024)
msg = chat.invoke([HumanMessage(content=[{"type": "text", "text":prompt},{"type": "image_url","image_url": {"url": f"data:image/jpeg;base64,{img_base64}"},},])])return msg.content

# Store base64 encoded imagesimg_base64_list = []
# Store image summariesimage_summaries = []
# Prompt : Our prompt here is customized to the type of images we have which is chart in our caseprompt = "Describe the image in detail. Be specific about graphs, such as bar plots."
# Read images, encode to base64 stringsfor img_file in sorted(os.listdir(path)):if img_file.endswith('.jpg'):img_path = os.path.join(path, img_file)base64_image = encode_image(img_path)img_base64_list.append(base64_image)img_capt = image_captioning(base64_image,prompt)time.sleep(60)image_summaries.append(image_captioning(img_capt,prompt))
def split_image_text_types(docs):''' Split base64-encoded images and texts '''b64 = []text = []for doc in docs:try:b64decode(doc)b64.append(doc)except Exception as e:text.append(doc)return {"images": b64,"texts": text}

# The vectorstore to use to index the child chunksvectorstore = Chroma(collection_name="multi_modal_rag", embedding_function=OpenAIEmbeddings())
# The storage layer for the parent documentsstore = InMemoryStore()id_key = "doc_id"
# The retriever (empty to start)retriever = MultiVectorRetriever(vectorstore=vectorstore,docstore=store,id_key=id_key,)
# Add textsdoc_ids = [str(uuid.uuid4()) for _ in texts]summary_texts = [Document(page_content=s, metadata={id_key: doc_ids[i]})for i, s in enumerate(text_summaries)]retriever.vectorstore.add_documents(summary_texts)retriever.docstore.mset(list(zip(doc_ids, texts)))
# Add tablestable_ids = [str(uuid.uuid4()) for _ in tables]summary_tables = [Document(page_content=s, metadata={id_key: table_ids[i]})for i, s in enumerate(table_summaries)]retriever.vectorstore.add_documents(summary_tables)retriever.docstore.mset(list(zip(table_ids, tables)))
# Add image summariesimg_ids = [str(uuid.uuid4()) for _ in img_base64_list]summary_img = [Document(page_content=s, metadata={id_key: img_ids[i]})for i, s in enumerate(image_summaries)]retriever.vectorstore.add_documents(summary_img)retriever.docstore.mset(list(zip(img_ids, img_base64_list)))

from operator import itemgetterfrom langchain.schema.runnable import RunnablePassthrough, RunnableLambda

def prompt_func(dict):format_texts = "\\n".join(dict["context"]["texts"])return [HumanMessage(content=[{"type": "text", "text": f"""Answer the question based only on the following context, which can include text, tables, and the below image:Question: {dict["question"]}

Text and tables:{format_texts}"""},{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{dict['context']['images'][0]}"}},])]

model = ChatOpenAI(temperature=0, model="gpt-4-vision-preview", max_tokens=1024)

# RAG pipelinechain = ({"context": retriever | RunnableLambda(split_image_text_types), "question": RunnablePassthrough()}| RunnableLambda(prompt_func)| model| StrOutputParser())

chain.invoke("What is the change in wild fires from 1993 to 2022?")