all

app.py

@@ -0,0 +1,311 @@
+import gradio as gr
+import tempfile
+import os
+import fitz  # PyMuPDF
+import uuid
+import shutil
+from pymilvus import MilvusClient
+
+from middleware import Middleware
+from rag import Rag
+from pathlib import Path
+import subprocess
+import getpass
+
+rag = Rag()
+
+
+def generate_uuid(state):
+    # Check if UUID already exists in session state
+    if state["user_uuid"] is None:
+        # Generate a new UUID if not already set
+        state["user_uuid"] = str(uuid.uuid4())
+
+    return state["user_uuid"]
+
+
+class PDFSearchApp:
+    def __init__(self):
+        self.indexed_docs = {}
+        self.current_pdf = None
+        
+    def upload_and_convert(self, state, files, max_pages):
+        #change id
+        #id = generate_uuid(state)
+        
+        
+        pages = 0
+
+        if files is None:
+            return "No file uploaded"
+        try: #if onlyy one file
+            for file in files[:]:  # Iterate over a shallow copy of the list, TEST THIS
+               
+                # Extract the last part of the path (file name)
+                filename = os.path.basename(file.name)
+
+                # Split the base name into name and extension
+                name, ext = os.path.splitext(filename)
+                self.current_pdf = file.name 
+                pdf_path=file.name
+                #if ppt will get replaced with path of ppt!
+
+                #if extension is .ppt or .pptx, convert 
+                if ext == ".ppt" or ext == ".pptx": #need to test with a ppt key...
+                    '''
+                    import comtypes.client
+                    powerpoint = comtypes.client.CreateObject("PowerPoint.Application")
+                    powerpoint.Visible = 1
+                    presentation = powerpoint.Presentations.Open(file)
+                    output_file = os.path.splitext(file)[0] + '.pdf'
+                    output_directory = os.path.dirname(file)
+                    presentation.SaveAs(os.path.join(output_directory, output_file), 32)  # 32 is the formatType for PDF
+                    presentation.Close()
+                    powerpoint.Quit()
+                    file = os.path.join(output_directory, output_file) #swap file to be used to the outputted pdf file instead
+                    # Extract the last part of the path (file name)
+                    name = os.path.basename(file)
+                    # Split the base name into name and extension
+                    name, ext = os.path.splitext(name)
+                    print(name)
+                    self.current_pdf = os.path.join(output_directory, output_file)
+                    pdf_path = os.path.join(output_directory, output_file)'
+                    '''
+                    print("pptx not supported on spaces")
+
+
+                # Replace spaces and hyphens with underscores in the name
+                modified_filename = name.replace(" ", "_").replace("-", "_")
+
+                id = modified_filename #if string cmi then serialize the name, test for later
+
+                print(f"Uploading file: {id}, id: abc")
+                middleware = Middleware(modified_filename, create_collection=True)
+               
+                
+                pages = middleware.index(pdf_path, id=id, max_pages=max_pages)
+                
+
+                self.indexed_docs[id] = True
+                
+            #clear files for next consec upload after loop is complete
+            files = []
+            return f"Uploaded and extracted {len(pages)} pages"
+        except Exception as e:
+            return f"Error processing PDF: {str(e)}"
+    
+
+    def display_file_list(text):
+        try:
+        # Retrieve all entries in the specified directory
+            directory_path = "pages"
+            current_working_directory = os.getcwd()
+            directory_path = os.path.join(current_working_directory, directory_path)
+            entries = os.listdir(directory_path)
+            # Filter out entries that are directories
+            directories = [entry for entry in entries if os.path.isdir(os.path.join(directory_path, entry))]
+            return directories
+        except FileNotFoundError:
+            return f"The directory {directory_path} does not exist."
+        except PermissionError:
+            return f"Permission denied to access {directory_path}."
+        except Exception as e:
+            return str(e)
+
+    
+    def search_documents(self, state, query, num_results=1):
+        print(f"Searching for query: {query}")
+        #id = generate_uuid(state)
+        id = "test" # not used anyway
+       
+        """
+        if not self.indexed_docs[id]:
+            print("Please index documents first")
+            return "Please index documents first", "--"
+        """ #edited out to allow direct query on db to test persistency
+        if not query:
+            print("Please enter a search query")
+            return "Please enter a search query", "--"            
+        try:
+
+            middleware = Middleware(id, create_collection=False)
+            
+            search_results = middleware.search([query])[0]
+            #direct retrieve file path rather than rely on page nums!
+            #try to retrieve multiple files rather than a single page (TBD)
+
+            page_num = search_results[0][1] + 1 # final return value is a list of tuples, each tuple being: (score, doc_id, collection_name), so use [0][2] to get collection name of first ranked item
+            coll_num = search_results[0][2]
+
+            print(f"Retrieved page number: {page_num}")
+
+            img_path = f"pages/{coll_num}/page_{page_num}.png"
+            path = f"pages/{coll_num}/page_{page_num}"
+
+            print(f"Retrieved image path: {img_path}")
+
+            rag_response = rag.get_answer_from_gemini(query, [img_path])
+
+            return path,img_path, rag_response
+            
+        except Exception as e:
+            return f"Error during search: {str(e)}", "--"
+        
+    def delete(state,choice):
+        #delete file in pages, then use middleware to delete collection 
+        # 1. Create a milvus client
+
+        client = MilvusClient(uri="http://localhost:19530")
+        #client = MilvusClient(
+        #        uri="http://localhost:19530",
+        #        token="root:Milvus"
+        #    )
+        path = f"pages/{choice}"
+        if os.path.exists(path):
+            shutil.rmtree(path)
+            #call milvus manager to delete collection
+            client.drop_collection(collection_name=choice)
+            return f"Deleted {choice}"
+        else:
+            return "Directory not found"
+        
+    def list_downloaded_hf_models(state):
+        # Determine the cache directory
+        hf_cache_dir = Path(os.getenv('HF_HOME', Path.home() / '.cache/huggingface/hub'))
+
+        # Initialize a list to store model names
+        model_names = []
+
+        # Traverse the cache directory
+        for repo_dir in hf_cache_dir.glob('models--*'):
+            # Extract the model name from the directory structure
+            model_name = repo_dir.name.split('--', 1)[-1].replace('-', '/')
+            model_names.append(model_name)
+
+        return model_names
+
+
+    def list_downloaded_ollama_models(state,):
+        # Retrieve the current user's name
+        username = getpass.getuser()
+        
+        # Construct the target directory path
+        base_path = f"C:\\Users\\{username}\\NEW_PATH\\manifests\\registry.ollama.ai\\library"
+        
+        try:
+            # List all entries in the directory
+            with os.scandir(base_path) as entries:
+                # Filter and print only directories
+                directories = [entry.name for entry in entries if entry.is_dir()]
+                
+            return directories
+        except FileNotFoundError:
+            print(f"The directory {base_path} does not exist.")
+        except PermissionError:
+            print(f"Permission denied to access {base_path}.")
+        except Exception as e:
+            print(f"An error occurred: {e}")
+            
+    def model_settings(state,hfchoice, ollamachoice,tokensize):
+        os.environ['colpali'] = hfchoice
+        os.environ['ollama'] = ollamachoice
+        os.environ['tokens'] = tokensize
+        return "abc"
+
+
+
+def create_ui():
+    app = PDFSearchApp()
+    
+    with gr.Blocks(css="footer{display:none !important}") as demo:
+        state = gr.State(value={"user_uuid": None})
+        
+  
+        gr.Markdown("# Collar Multimodal RAG Demo")
+        gr.Markdown("Made by Collar")   
+        
+        with gr.Tab("Upload PDF"):
+            with gr.Column():
+                max_pages_input = gr.Slider(
+                    minimum=1,
+                    maximum=10000,
+                    value=20,
+                    step=10,
+                    label="Max pages to extract and index per document"
+                )
+                file_input = gr.Files(label="Upload PDFs")
+                file_list = gr.Textbox(label="Uploaded Files", interactive=False, value=app.display_file_list())
+                status = gr.Textbox(label="Indexing Status", interactive=False)
+                
+        
+        with gr.Tab("Query"):
+            with gr.Column():
+                query_input = gr.Textbox(label="Enter query")
+                #num_results = gr.Slider(
+                #    minimum=1,
+                #    maximum=10,
+                #    value=5,
+                #    step=1,
+                #    label="Number of results"
+                #)
+                search_btn = gr.Button("Query")
+                llm_answer = gr.Textbox(label="RAG Response", interactive=False)
+                path = gr.Textbox(label="Link To Document Page", interactive=False)
+                images = gr.Image(label="Top page matching query")
+        with gr.Tab("Data Settings"): #deletion of collections, changing of model parameters etc
+            with gr.Column():
+                 # Button to delete (TBD)
+                choice = gr.Dropdown(list(app.display_file_list()),label="Choice")
+                delete_button = gr.Button("Delete Document From DB")
+                status1 = gr.Textbox(label="Deletion Status", interactive=False)
+
+        with gr.Tab("AI Model Settings"): #deletion of collections, changing of model parameters etc
+            with gr.Column():
+                 # Button to delete (TBD)
+                hfchoice = gr.Dropdown(app.list_downloaded_hf_models(),label="Visual Document Retrieval (VDR) Model")
+                ollamachoice = gr.Dropdown(app.list_downloaded_ollama_models(),label="Secondary Visual Retrieval-Augmented Generation (RAG) Model")
+                tokensize = gr.Slider(
+                    minimum=256,
+                    maximum=4096,
+                    value=20,
+                    step=10,
+                    label="Max tokens per response (Reply Length)"
+                )
+                model_button = gr.Button("Update Settings")
+                status2 = gr.Textbox(label="Update Status", interactive=False)
+
+        
+
+        
+        # Event handlers
+        file_input.change(
+            fn=app.upload_and_convert,
+            inputs=[state, file_input, max_pages_input],
+            outputs=[status]
+        )
+        
+        search_btn.click(
+         #try to query without uploading first
+            fn= app.search_documents,
+            inputs=[state, query_input],
+            outputs=[path,images, llm_answer]
+        )
+    
+        delete_button.click(
+            fn=app.delete,
+            inputs=[choice],
+            outputs=[status1]
+        )
+        
+        model_button.click(
+            fn=app.model_settings,
+            inputs=[hfchoice, ollamachoice,tokensize],
+            outputs=[status2]
+        )
+        
+    return demo
+
+if __name__ == "__main__":
+    demo = create_ui()
+    demo.launch()
+

colpali_manager.py

@@ -0,0 +1,141 @@
+from colpali_engine.models import ColPali
+from colpali_engine.models.paligemma.colpali.processing_colpali import ColPaliProcessor
+from colpali_engine.utils.processing_utils import BaseVisualRetrieverProcessor
+from colpali_engine.utils.torch_utils import ListDataset, get_torch_device
+from torch.utils.data import DataLoader
+import torch
+from typing import List, cast
+
+#from colpali_engine.models import ColQwen2_5, ColQwen2_5_Processor
+from colpali_engine.models import ColIdefics3, ColIdefics3Processor
+
+from tqdm import tqdm
+from PIL import Image
+import os
+
+import spaces
+
+
+#this part is for local runs
+torch.cpu.empty_cache()
+
+model_name = "vidore/colSmol-256M"
+device = get_torch_device("cpu") #try using cpu instead of cpu?
+
+#switch to locally downloading models & loading locally rather than from hf
+#
+
+current_working_directory = os.getcwd()
+save_directory = model_name  # Directory to save the specific model name
+save_directory = os.path.join(current_working_directory, save_directory)
+
+processor_directory = 'local_processor'  # Directory to save the processor
+processor_directory = os.path.join(current_working_directory, processor_directory)
+
+
+model = ColIdefics3.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16,
+            device_map=device,
+            #attn_implementation="flash_attention_2",
+        ).eval()
+processor = cast(ColIdefics3Processor, ColIdefics3Processor.from_pretrained(model_name))
+
+"""
+if not os.path.exists(save_directory): #download if directory not created/model not loaded
+        # Directory does not exist; create it
+        os.makedirs(save_directory)
+        print(f"Directory '{save_directory}' created.")
+        model = ColIdefics3.from_pretrained(
+            model_name,
+            torch_dtype=torch.bfloat16,
+            device_map=device,
+            attn_implementation="flash_attention_2",
+        ).eval()
+        model.save_pretrained(save_directory)
+        os.makedirs(processor_directory)
+        processor = cast(ColIdefics3Processor, ColIdefics3Processor.from_pretrained(model_name))
+        
+        processor.save_pretrained(processor_directory)
+
+else:
+    model = ColIdefics3.from_pretrained(save_directory)
+    processor = ColIdefics3.from_pretrained(processor_directory, use_fast=True)
+"""
+
+
+class ColpaliManager:
+
+    
+    def __init__(self, device = "cpu", model_name = "vidore/colSmol-256M"): #need to hot potato/use diff gpus between colpali & ollama
+
+        print(f"Initializing ColpaliManager with device {device} and model {model_name}")
+
+        # self.device = get_torch_device(device)
+
+        # self.model = ColPali.from_pretrained(
+        #     model_name,
+        #     torch_dtype=torch.bfloat16,
+        #     device_map=self.device,
+        # ).eval()
+
+        # self.processor = cast(ColPaliProcessor, ColPaliProcessor.from_pretrained(model_name))
+
+    @spaces.GPU
+    def get_images(self, paths: list[str]) -> List[Image.Image]:
+        model.to("cpu")
+        return [Image.open(path) for path in paths]
+
+    @spaces.GPU
+    def process_images(self, image_paths:list[str], batch_size=5):
+        model.to("cpu")
+        print(f"Processing {len(image_paths)} image_paths")
+        
+        images = self.get_images(image_paths)
+
+        dataloader = DataLoader(
+            dataset=ListDataset[str](images),
+            batch_size=batch_size,
+            shuffle=False,
+            collate_fn=lambda x: processor.process_images(x),
+        )
+
+        ds: List[torch.Tensor] = []
+        for batch_doc in tqdm(dataloader):
+            with torch.no_grad():
+                batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
+                embeddings_doc = model(**batch_doc)
+            ds.extend(list(torch.unbind(embeddings_doc.to(device))))
+                
+        ds_np = [d.float().cpu().numpy() for d in ds]
+
+        return ds_np
+    
+
+    @spaces.GPU
+    def process_text(self, texts: list[str]):
+        model.to("cpu") #ensure this is commented out so ollama/multimodal llm can use gpu! (nah wrong, need to enable so that it can process multiple)
+        print(f"Processing {len(texts)} texts")
+
+        dataloader = DataLoader(
+            dataset=ListDataset[str](texts),
+            batch_size=5,
+            shuffle=False,
+            collate_fn=lambda x: processor.process_queries(x),
+        )
+
+        qs: List[torch.Tensor] = []
+        for batch_query in dataloader:
+            with torch.no_grad():
+                batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
+                embeddings_query = model(**batch_query)
+
+            qs.extend(list(torch.unbind(embeddings_query.to(device))))
+
+        qs_np = [q.float().cpu().numpy() for q in qs]
+        model.to("cpu")  # Moves all model parameters and buffers to the CPU, freeing up gpu for ollama call after this process text call! (THIS WORKS!)
+
+        return qs_np
+    
+
+

middleware.py

@@ -0,0 +1,62 @@
+from colpali_manager import ColpaliManager
+from milvus_manager import MilvusManager
+from pdf_manager import PdfManager
+import hashlib
+
+
+
+pdf_manager = PdfManager()
+colpali_manager = ColpaliManager()
+
+
+
+class Middleware:
+    def __init__(self, id:str, create_collection=True):
+        #hashed_id = hashlib.md5(id.encode()).hexdigest()[:8]
+        hashed_id = 0 #switched to persistent db, shld use diff id for diff accs
+        milvus_db_name = f"milvus_{hashed_id}.db"
+        self.milvus_manager = MilvusManager(milvus_db_name, id, create_collection) #create collections based on id rather than colpali
+
+    def index(self, pdf_path: str, id:str, max_pages: int, pages: list[int] = None):
+
+        if type(pdf_path) == None: #for direct query without any upload to db
+            print("no docs")
+            return
+        
+        print(f"Indexing {pdf_path}, id: {id}, max_pages: {max_pages}")
+
+        image_paths = pdf_manager.save_images(id, pdf_path, max_pages)
+
+        print(f"Saved {len(image_paths)} images")
+
+        colbert_vecs = colpali_manager.process_images(image_paths)
+
+        images_data = [{
+            "colbert_vecs": colbert_vecs[i],
+            "filepath": image_paths[i]
+        } for i in range(len(image_paths))]
+
+        print(f"Inserting {len(images_data)} images data to Milvus")
+
+        self.milvus_manager.insert_images_data(images_data)
+
+        print("Indexing completed")
+
+        return image_paths
+
+
+        
+    def search(self, search_queries: list[str]):
+        print(f"Searching for {len(search_queries)} queries")
+
+        final_res = []
+
+        for query in search_queries:
+            print(f"Searching for query: {query}")
+            query_vec = colpali_manager.process_text([query])[0]
+            search_res = self.milvus_manager.search(query_vec, topk=1)
+            print(f"Search result: {search_res} for query: {query}")
+            final_res.append(search_res)
+
+        return final_res
+

milvus_manager.py

@@ -0,0 +1,195 @@
+from pymilvus import MilvusClient, DataType
+import numpy as np
+import concurrent.futures
+from pymilvus import Collection
+
+class MilvusManager:
+    def __init__(self, milvus_uri, collection_name, create_collection, dim=128):
+        self.client = MilvusClient(uri=milvus_uri)
+        # self.client = MilvusClient(uri="http://localhost:19530", token="root:Milvus")
+        self.collection_name = collection_name
+        self.dim = dim
+
+        if self.client.has_collection(collection_name=self.collection_name):
+            self.client.load_collection(collection_name=self.collection_name)
+            print("Loaded existing collection.")
+        elif create_collection:
+            self.create_collection()
+            self.create_index()
+
+    def create_collection(self):
+        if self.client.has_collection(collection_name=self.collection_name):
+            print("Collection already exists.")
+            return
+
+        schema = self.client.create_schema(
+            auto_id=True,
+            enable_dynamic_fields=True,
+        )
+        schema.add_field(field_name="pk", datatype=DataType.INT64, is_primary=True)
+        schema.add_field(
+            field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=self.dim
+        )
+        schema.add_field(field_name="seq_id", datatype=DataType.INT16)
+        schema.add_field(field_name="doc_id", datatype=DataType.INT64)
+        schema.add_field(field_name="doc", datatype=DataType.VARCHAR, max_length=65535)
+
+        self.client.create_collection(
+            collection_name=self.collection_name, schema=schema
+        )
+
+    def create_index(self):
+        index_params = self.client.prepare_index_params()
+        index_params.add_index(
+            field_name="vector",
+            index_name="vector_index",
+            index_type="HNSW", #use HNSW option if got more mem, if not use IVF for faster processing
+            metric_type="IP",
+            params={
+                "M": 16, #M:16 for HNSW, capital M
+                "efConstruction": 500, #for HNSW
+            },
+        )
+
+        self.client.create_index(
+            collection_name=self.collection_name, index_params=index_params, sync=True
+        )
+
+    def search(self, data, topk):
+         # Retrieve all collection names from the Milvus client.
+        collections = self.client.list_collections()
+        
+        # Set search parameters (here, using Inner Product metric).
+        search_params = {"metric_type": "IP", "params": {}}
+        
+        # Set to store unique (doc_id, collection_name) pairs across all collections.
+        doc_collection_pairs = set()
+
+        # Query each collection individually
+        for collection in collections:
+            self.client.load_collection(collection_name=collection)
+            print("collection loaded:"+ collection)
+            results = self.client.search(
+                collection,
+                data,
+                limit=50,  # Adjust limit per collection as needed.
+                output_fields=["vector", "seq_id", "doc_id"],
+                search_params=search_params,
+            )
+            # Accumulate document IDs along with their originating collection.
+            for r_id in range(len(results)):
+                for r in range(len(results[r_id])):
+                    doc_id = results[r_id][r]["entity"]["doc_id"]
+                    doc_collection_pairs.add((doc_id, collection))
+
+        scores = []
+
+        def rerank_single_doc(doc_id, data, client, collection_name):
+            # Query for detailed document vectors in the given collection.
+            doc_colbert_vecs = client.query(
+                collection_name=collection_name,
+                filter=f"doc_id in [{doc_id}, {doc_id + 1}]",
+                output_fields=["seq_id", "vector", "doc"],
+                limit=16380,
+            )
+            # Stack the vectors for dot product computation.
+            doc_vecs = np.vstack(
+                [doc_colbert_vecs[i]["vector"] for i in range(len(doc_colbert_vecs))]
+            )
+            # Compute a similarity score via dot product.
+            score = np.dot(data, doc_vecs.T).max(1).sum()
+            return (score, doc_id, collection_name)
+
+        # Use a thread pool to rerank each document concurrently.
+        with concurrent.futures.ThreadPoolExecutor(max_workers=300) as executor:
+            futures = {
+                executor.submit(rerank_single_doc, doc_id, data, self.client, collection): (doc_id, collection)
+                for doc_id, collection in doc_collection_pairs
+            }
+            for future in concurrent.futures.as_completed(futures):
+                score, doc_id, collection = future.result()
+                scores.append((score, doc_id, collection)) 
+                #doc_id is page number!
+        
+        # Sort the reranked results by score in descending order.
+        scores.sort(key=lambda x: x[0], reverse=True)
+        # Unload the collection after search to free memory.
+        self.client.release_collection(collection_name=collection)
+        
+        return scores[:topk] if len(scores) >= topk else scores
+        """
+        search_params = {"metric_type": "IP", "params": {}}
+        results = self.client.search(
+            self.collection_name,
+            data,
+            limit=50,
+            output_fields=["vector", "seq_id", "doc_id"],
+            search_params=search_params,
+        )
+        doc_ids = {result["entity"]["doc_id"] for result in results[0]}
+
+        scores = []
+
+        def rerank_single_doc(doc_id, data, client, collection_name):
+            doc_colbert_vecs = client.query(
+                collection_name=collection_name,
+                filter=f"doc_id in [{doc_id}, {doc_id + 1}]",
+                output_fields=["seq_id", "vector", "doc"],
+                limit=1000,
+            )
+            doc_vecs = np.vstack(
+                [doc["vector"] for doc in doc_colbert_vecs]
+            )
+            score = np.dot(data, doc_vecs.T).max(1).sum()
+            return score, doc_id
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=300) as executor:
+            futures = {
+                executor.submit(
+                    rerank_single_doc, doc_id, data, self.client, self.collection_name
+                ): doc_id
+                for doc_id in doc_ids
+            }
+            for future in concurrent.futures.as_completed(futures):
+                score, doc_id = future.result()
+                scores.append((score, doc_id))
+
+        scores.sort(key=lambda x: x[0], reverse=True)
+        return scores[:topk]
+        """
+
+    def insert(self, data):
+        colbert_vecs = data["colbert_vecs"]
+        seq_length = len(colbert_vecs)
+        doc_ids = [data["doc_id"]] * seq_length
+        seq_ids = list(range(seq_length))
+        docs = [""] * seq_length
+        docs[0] = data["filepath"]
+
+        self.client.insert(
+            self.collection_name,
+            [
+                {
+                    "vector": colbert_vecs[i],
+                    "seq_id": seq_ids[i],
+                    "doc_id": doc_ids[i],
+                    "doc": docs[i],
+                }
+                for i in range(seq_length)
+            ],
+        )
+
+    def get_images_as_doc(self, images_with_vectors):
+        return [
+            {
+                "colbert_vecs": image["colbert_vecs"],
+                "doc_id": idx,
+                "filepath": image["filepath"],
+            }
+            for idx, image in enumerate(images_with_vectors)
+        ]
+
+    def insert_images_data(self, image_data):
+        data = self.get_images_as_doc(image_data)
+        for item in data:
+            self.insert(item)

packages.txt

@@ -0,0 +1 @@
+poppler-utils

pdf_manager.py

@@ -0,0 +1,46 @@
+from pdf2image import convert_from_path
+import os
+import shutil
+
+class PdfManager:
+    def __init__(self):
+        pass
+        
+    def clear_and_recreate_dir(self, output_folder):
+        
+        print(f"Clearing output folder {output_folder}")
+
+        if os.path.exists(output_folder):
+            shutil.rmtree(output_folder)
+            #print("Clearing is unused for now for persistency")
+        else:
+            os.makedirs(output_folder)
+        
+        #print("Clearing is unused for now for persistency")
+
+    def save_images(self, id, pdf_path, max_pages, pages: list[int] = None) -> list[str]:
+        output_folder = f"pages/{id}" #remove last backslash to avoid error,test this
+        images = convert_from_path(pdf_path)
+
+        print(f"Saving images from {pdf_path} to {output_folder}. Max pages: {max_pages}")
+
+        self.clear_and_recreate_dir(output_folder) 
+
+        num_page_processed = 0
+
+        for i, image in enumerate(images):
+            if max_pages and num_page_processed >= max_pages:
+                break
+
+            if pages and i not in pages:
+                continue
+
+            full_save_path = f"{output_folder}/page_{i + 1}.png"
+
+            #print(f"Saving image to {full_save_path}")
+
+            image.save(full_save_path, "PNG")
+
+            num_page_processed += 1
+
+        return [f"{output_folder}/page_{i + 1}.png" for i in range(num_page_processed)]

rag.py

@@ -0,0 +1,147 @@
+import requests
+import os
+
+from typing import List
+from utils import encode_image
+from PIL import Image
+from ollama import chat
+import torch
+import subprocess
+import psutil
+import torch
+from transformers import AutoModel, AutoTokenizer
+import   google.generativeai as genai
+
+
+
+class Rag:
+    
+    def get_answer_from_gemini(self, query, imagePaths):
+
+        print(f"Querying Gemini for query={query}, imagePaths={imagePaths}")
+
+        try:
+            genai.configure(api_key="AIzaSyCwRr9054tCuh2S8yGpwKFvOAxYMT4WNIs")
+            model = genai.GenerativeModel('gemini-1.5-flash')
+            
+            images = [Image.open(path) for path in imagePaths]
+
+            chat = model.start_chat()
+
+            response = chat.send_message([*images, query])
+
+            answer = response.text
+
+            print(answer)
+            
+            return answer
+        
+        except Exception as e:
+            print(f"An error occurred while querying Gemini: {e}")
+            return f"Error: {str(e)}"
+    
+    #os.environ['OPENAI_API_KEY'] = "for the love of Jesus let this work"
+
+    def get_answer_from_openai(self, query, imagesPaths):
+        """ #scuffed local hf inference (transformers incompatible to colpali version req, use ollama, more reliable, easier to use plus web server ready)
+        print(f"Querying for query={query}, imagesPaths={imagesPaths}")
+
+        model = AutoModel.from_pretrained(
+            'openbmb/MiniCPM-o-2_6-int4',
+            trust_remote_code=True,
+            attn_implementation='flash_attention_2', # sdpa or flash_attention_2
+            torch_dtype=torch.bfloat16,
+            init_vision=True,
+        )
+
+
+        model = model.eval().cuda()
+        tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V-2_6-int4', trust_remote_code=True)
+        image = Image.open(imagesPaths[0]).convert('RGB')
+        
+        msgs = [{'role': 'user', 'content': [image, query]}]
+        answer = model.chat(
+            image=None,
+            msgs=msgs,
+            tokenizer=tokenizer
+        )
+        print(answer)
+        return answer
+        """
+
+        #ollama method below
+        
+        torch.cuda.empty_cache() #release cuda so that ollama can use gpu!
+
+    
+        os.environ['OLLAMA_FLASH_ATTENTION'] = '1'
+        
+
+        # Close model thread (colpali)
+        print(f"Querying OpenAI for query={query}, imagesPaths={imagesPaths}")
+
+        try:    
+            
+            response = chat(    
+                model='minicpm-v:8b-2.6-q8_0',
+                messages=[
+                    {
+                    'role': 'user',
+                    'content': query,
+                    'images': imagesPaths,
+                    }
+                ],
+                )
+    
+            answer = response.message.content
+    
+            print(answer)
+    
+            return answer
+    
+        except Exception as e:
+            print(f"An error occurred while querying OpenAI: {e}")
+            return None
+        
+
+
+    def __get_openai_api_payload(self, query:str, imagesPaths:List[str]):
+        image_payload = []
+
+        for imagePath in imagesPaths:
+            base64_image = encode_image(imagePath)
+            image_payload.append({
+                "type": "image_url",
+                "image_url": {
+                    "url": f"data:image/jpeg;base64,{base64_image}"
+                }
+            })
+
+        payload = {
+            "model": "Llama3.2-vision", #change model here as needed
+            "messages": [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": query
+                        },
+                        *image_payload
+                    ]
+                }
+            ],
+            "max_tokens": 1024 #reduce token size to reduce processing time
+        }
+
+        return payload
+    
+
+
+# if __name__ == "__main__":
+#     rag = Rag()
+    
+#     query = "Based on attached images, how many new cases were reported during second wave peak"
+#     imagesPaths = ["covid_slides_page_8.png", "covid_slides_page_8.png"]
+    
+#     rag.get_answer_from_gemini(query, imagesPaths)

requirements.txt

@@ -0,0 +1,14 @@
+gradio
+PyMuPDF
+pdf2image
+pymilvus
+tqdm
+pillow
+spaces
+google-generativeai
+git+https://github.com/illuin-tech/colpali
+timm==1.0.13
+transformers
+https://github.com/woct0rdho/triton-windows/releases/download/v3.2.0-windows.post10/triton-3.2.0-cp311-cp311-win_amd64.wh
+comtypes
+python-dotenv

test.py

@@ -0,0 +1,30 @@
+from pymilvus import MilvusClient
+from pymilvus import (
+    connections,
+    utility,
+    FieldSchema, CollectionSchema, DataType,
+    Collection,
+)
+
+# 1. Create a milvus client
+client = MilvusClient(
+    uri="http://localhost:19530",
+    token="root:Milvus"
+)
+
+# 2. Create a collection
+client.drop_collection(collection_name="fy2025_budget_statement")
+
+# 3. List collections
+print(client.list_collections() )
+
+# ['test_collection']
+
+"""
+res = client.get(
+    collection_name="colpali",
+    ids=[0, 1, 2],
+)
+
+print(res)
+"""

uploaded_files.txt

@@ -0,0 +1,3 @@
+EMERGING_MISSILE_THREATS_16382509
+handwriting
+multimediareport

utils.py

@@ -0,0 +1,5 @@
+import base64
+
+def encode_image(image_path):
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode('utf-8')