Upload folder using huggingface_hub

app.py

@@ -0,0 +1,206 @@
+import streamlit as st
+import os
+from PyPDF2 import PdfReader
+import pymupdf
+import numpy as np
+import cv2
+import shutil
+import imageio
+from PIL import Image
+import imagehash
+import matplotlib.pyplot as plt
+from llama_index.core.indices import MultiModalVectorStoreIndex
+from llama_index.vector_stores.qdrant import QdrantVectorStore
+from llama_index.core import SimpleDirectoryReader, StorageContext
+import qdrant_client
+from llama_index.core import PromptTemplate
+from llama_index.core.query_engine import SimpleMultiModalQueryEngine
+from llama_index.llms.openai import OpenAI
+from llama_index.core import load_index_from_storage, get_response_synthesizer
+import tempfile
+
+
+def extract_text_from_pdf(pdf_path):
+    reader = PdfReader(pdf_path)
+    full_text = ''
+    for page in reader.pages:
+        text = page.extract_text()
+        full_text += text
+    return full_text
+
+def extract_images_from_pdf(pdf_path, img_save_path):
+    doc = pymupdf.open(pdf_path)
+    for page in doc:
+        img_number = 0
+        for block in page.get_text("dict")["blocks"]:
+            if block['type'] == 1:
+                name = os.path.join(img_save_path, f"img{page.number}-{img_number}.{block['ext']}")
+                out = open(name, "wb")
+                out.write(block["image"])
+                out.close()
+                img_number += 1
+
+def is_empty(img_path):
+    image = cv2.imread(img_path, 0)
+    std_dev = np.std(image)
+    return std_dev < 1
+
+def move_images(source_folder, dest_folder):
+    image_files = [f for f in os.listdir(source_folder) 
+                   if f.lower().endswith(('.jpg', '.jpeg', '.png', '.gif'))]
+    os.makedirs(dest_folder, exist_ok=True)
+    moved_count = 0
+    for file in image_files:
+        src_path = os.path.join(source_folder, file)
+        if not is_empty(src_path):
+            shutil.move(src_path, os.path.join(dest_folder, file))
+            moved_count += 1
+    return moved_count
+
+def remove_low_size_images(data_path):
+    images_list = os.listdir(data_path)
+    low_size_photo_list = []
+    for one_image in images_list:
+        image_path = os.path.join(data_path, one_image)
+        try:
+            pic = imageio.imread(image_path)
+            size = pic.size
+            if size < 100:
+                low_size_photo_list.append(one_image)
+        except:
+            pass
+    for one_image in low_size_photo_list[1:]:
+        os.remove(os.path.join(data_path, one_image))
+
+def initialize_qdrant(temp_dir):
+    try :
+        client = qdrant_client.QdrantClient(path="qdrant_mm_db_pipeline")
+    except :
+        pass
+    if "vectordatabase" not in st.session_state or not st.session_state.vectordatabase:
+        text_store = QdrantVectorStore(client=client, collection_name="text_collection_pipeline")
+        image_store = QdrantVectorStore(client=client, collection_name="image_collection_pipeline")
+        storage_context = StorageContext.from_defaults(vector_store=text_store, image_store=image_store)
+        documents = SimpleDirectoryReader(os.path.join(temp_dir, "my_own_data")).load_data()
+        index = MultiModalVectorStoreIndex.from_documents(documents, storage_context=storage_context)
+        st.session_state.vectordatabase = index
+    else :
+        index = st.session_state.vectordatabase
+    retriever_engine = index.as_retriever(similarity_top_k=1, image_similarity_top_k=1)
+    return retriever_engine
+
+def plot_images(image_paths):
+    images_shown = 0
+    plt.figure(figsize=(16, 9))
+    for img_path in image_paths:
+        if os.path.isfile(img_path):
+            image = Image.open(img_path)
+            plt.subplot(2, 3, images_shown + 1)
+            plt.imshow(image)
+            plt.xticks([])
+            plt.yticks([])
+            images_shown += 1
+            if images_shown >= 6:
+                break
+
+def retrieve_and_query(query, retriever_engine):
+    retrieval_results = retriever_engine.retrieve(query)
+    
+    qa_tmpl_str = (
+        "Context information is below.\n"
+        "---------------------\n"
+        "{context_str}\n"
+        "---------------------\n"
+        "Given the context information , "
+        "answer the query in detail.\n"
+        "Query: {query_str}\n"
+        "Answer: "
+    )
+    qa_tmpl = PromptTemplate(qa_tmpl_str)
+
+    llm = OpenAI(model="gpt-4o", temperature=0)
+    response_synthesizer = get_response_synthesizer(response_mode="refine", text_qa_template=qa_tmpl, llm=llm)
+
+    response = response_synthesizer.synthesize(query, nodes=retrieval_results)
+    
+    retrieved_image_path_list = []
+    for node in retrieval_results:
+        if (node.metadata['file_type'] == 'image/jpeg') or (node.metadata['file_type'] == 'image/png'):
+            if node.score > 0.25:
+                retrieved_image_path_list.append(node.metadata['file_path'])
+    
+    return response, retrieved_image_path_list
+
+def process_pdf(pdf_file):
+    # import pdb; pdb.set_trace()
+    temp_dir = tempfile.TemporaryDirectory()
+    temp_pdf_path = os.path.join(temp_dir.name, pdf_file.name)
+    with open(temp_pdf_path, "wb") as f:
+        f.write(pdf_file.getvalue())
+
+    data_path = os.path.join(temp_dir.name, "my_own_data")
+    os.makedirs(data_path , exist_ok=True) 
+    img_save_path = os.path.join(temp_dir.name, "extracted_images")
+    os.makedirs(img_save_path , exist_ok=True) 
+
+    extracted_text = extract_text_from_pdf(temp_pdf_path)
+    with open(os.path.join(data_path, "content.txt"), "w") as file:
+        file.write(extracted_text)
+
+    extract_images_from_pdf(temp_pdf_path, img_save_path)
+    moved_count = move_images(img_save_path, data_path)
+    remove_low_size_images(data_path)
+
+    retriever_engine = initialize_qdrant(temp_dir.name)
+
+    return temp_dir, retriever_engine
+
+def main():
+    st.title("PDF Vector Database Query Tool")
+    st.markdown("This tool creates a vector database from a PDF and allows you to query it.")
+    
+    if "retriever_engine" not in st.session_state:
+        st.session_state.retriever_engine = None
+    if "vectordatabase" not in st.session_state:
+        st.session_state.vectordatabase = None
+
+    uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
+    if uploaded_file is None:
+        st.info("Please upload a PDF file.")
+    else:
+        st.info(f"Uploaded PDF: {uploaded_file.name}")
+        if st.button("Process PDF"):
+            with st.spinner("Processing PDF..."):
+                temp_dir, st.session_state.retriever_engine = process_pdf(uploaded_file)
+        
+                st.success("PDF processed successfully!")
+
+    
+    query = st.text_input("Enter your question:")
+    
+    
+    if st.button("Ask Question"):
+        print("running")
+        try:
+            import pdb; pdb.set_trace() 
+
+            with st.spinner("Retrieving information..."):
+                import pdb; pdb.set_trace()
+                response, retrieved_image_path_list = retrieve_and_query(query, st.session_state.retriever_engine)
+            
+            st.write("Retrieved Context:")
+            for node in response.source_nodes:
+                st.code(node.node.get_text())
+            
+            st.write("\nRetrieved Images:")
+            plot_images(retrieved_image_path_list)
+            st.pyplot()
+
+            st.write("\nFinal Answer:")
+            st.code(response.response)
+        
+        except Exception as e:
+            st.error(f"An error occurred: {e}")
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,15 @@
+PyPDF2==3.0.1
+PyMuPDF==1.24.9
+numpy==1.26.4
+opencv-python==4.10.0.84
+matplotlib==3.9.2
+llama-index==0.11.2
+llama-index-vector-stores-qdrant==0.3.0
+ipython==8.26.0
+llama-index-embeddings-clip==0.2.0
+imageio==2.35.1
+pillow==10.4.0
+imagehash
+llama-index-embeddings-clip
+git+https://github.com/openai/CLIP.git
+dotenv