change

.env

@@ -0,0 +1,8 @@
+colpali='vidore/colpali-v1.3'
+ollama='minicpm-v'
+flashattn='1'
+metrictype='IP'
+mnum='16'
+efnum='500'
+topk='50'
+temperature='0.8'

app.py

@@ -11,10 +11,23 @@ from rag import Rag
 from pathlib import Path
 import subprocess
 import getpass
+# importing necessary functions from dotenv library
+from dotenv import load_dotenv, dotenv_values 
+import dotenv
+import platform
+import time
+
+# loading variables from .env file
+dotenv_file = dotenv.find_dotenv()
+dotenv.load_dotenv(dotenv_file)
+
+#kickstart docker and ollama servers
+
 
 rag = Rag()
 
 
+
 def generate_uuid(state):
     # Check if UUID already exists in session state
     if state["user_uuid"] is None:
@@ -50,29 +63,6 @@ class PDFSearchApp:
                 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("-", "_")
@@ -154,12 +144,9 @@ class PDFSearchApp:
     def delete(state,choice):
         #delete file in pages, then use middleware to delete collection 
         # 1. Create a milvus client
-
-        client = MilvusClient(uri="localhost")
-        #client = MilvusClient(
-        #        uri="http://localhost:19530",
-        #        token="root:Milvus"
-        #    )
+        client = MilvusClient(
+                uri="localhost"
+            )
         path = f"pages/{choice}"
         if os.path.exists(path):
             shutil.rmtree(path)
@@ -168,6 +155,18 @@ class PDFSearchApp:
             return f"Deleted {choice}"
         else:
             return "Directory not found"
+    def dbupdate(state,metric_type,m_num,ef_num,topk):
+        os.environ['metrictype'] = metric_type
+        # Update the .env file with the new value   
+        dotenv.set_key(dotenv_file, 'metrictype', metric_type)
+        os.environ['mnum'] = str(m_num)
+        dotenv.set_key(dotenv_file, 'mnum', str(m_num))
+        os.environ['efnum'] = str(ef_num)
+        dotenv.set_key(dotenv_file, 'efnum', str(ef_num))
+        os.environ['topk'] = str(topk)
+        dotenv.set_key(dotenv_file, 'topk', str(topk))
+
+        return "DB Settings Updated, Restart App To Load"
         
     def list_downloaded_hf_models(state):
         # Determine the cache directory
@@ -179,18 +178,19 @@ class PDFSearchApp:
         # 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_name = repo_dir.name.split('--', 1)[-1].replace('--', '/')
             model_names.append(model_name)
 
         return model_names
 
 
-    def list_downloaded_ollama_models(state,):
+    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"
+        #base_path = f"C:\\Users\\{username}\\NEW_PATH\\manifests\\registry.ollama.ai\\library" #this is for if ollama pull is called from C://, if ollama pulls are called from the proj dir, use the NEW_PATH in the proj dir!
+        base_path = f"NEW_PATH\\manifests\\registry.ollama.ai\\library" #relative to proj dir! (IMPT: OLLAMA PULL COMMAND IN PROJ DIR!!!)
         
         try:
             # List all entries in the directory
@@ -206,18 +206,29 @@ class PDFSearchApp:
         except Exception as e:
             print(f"An error occurred: {e}")
             
-    def model_settings(state,hfchoice, ollamachoice,tokensize):
+    def model_settings(state,hfchoice, ollamachoice,flash, temp):
         os.environ['colpali'] = hfchoice
+        # Update the .env file with the new value   
+        dotenv.set_key(dotenv_file, 'colpali', hfchoice)
         os.environ['ollama'] = ollamachoice
-        os.environ['tokens'] = tokensize
-        return "abc"
+        dotenv.set_key(dotenv_file, 'ollama', ollamachoice)
+        if flash == "Enabled":
+            os.environ['flashattn'] = "1"
+            dotenv.set_key(dotenv_file, 'flashattn', "1")
+        else:
+             os.environ['flashattn'] = "0"
+             dotenv.set_key(dotenv_file, 'flashattn', "0")
+        os.environ['temperature'] = str(temp)
+        dotenv.set_key(dotenv_file, 'temperature', str(temp))
+        
+        return "Models Updated, Restart App To Use New Settings"
 
 
 
 def create_ui():
     app = PDFSearchApp()
     
-    with gr.Blocks(css="footer{display:none !important}") as demo:
+    with gr.Blocks(theme=gr.themes.Ocean(),css ="footer{display:none !important}") as demo:
         state = gr.State(value={"user_uuid": None})
         
   
@@ -256,26 +267,47 @@ def create_ui():
             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)
+                delete_button = gr.Button("Delete Document From DB")
+                
+                # Create the dropdown component with default value as the first option
+                #Milvusindex = gr.Dropdown(["HNSW","FLAT", "IVF_FLAT", "IVF_SQ8", "IVF_PQ", "RHNSW_FLAT"], value="HNSW", label="Select Vector DB Index Parameter")
+                metric_type = gr.Dropdown(choices=["IP", "L2", "COSINE"],value="IP",label="Metric Type (Mathematical function to measure similarity)")
+                m_num = gr.Dropdown(
+                    choices=["8", "16", "32", "64"], value="16",label="M Vectors (Maximum number of neighbors each node can connect to in the graph)")
+                ef_num =  gr.Slider(
+                    minimum=50,
+                    maximum=1000,
+                    value=500,
+                    step=10,
+                    label="EF Construction (Number of candidate neighbors considered for connection during index construction)"
+                )
+                topk = gr.Slider(
+                    minimum=1,
+                    maximum=100,
+                    value=50,
+                    step=1,
+                    label="Top-K (Maximum number of entities to return in a single search of a document)"
+                )
+                db_button = gr.Button("Update DB Settings")
+                status3 = gr.Textbox(label="DB Update 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)"
+                hfchoice = gr.Dropdown(app.list_downloaded_hf_models(),value=os.environ['colpali'], label="Visual Document Retrieval (VDR) Model")
+                ollamachoice = gr.Dropdown(app.list_downloaded_ollama_models(),value=os.environ['ollama'],label="Secondary Visual Retrieval-Augmented Generation (RAG) Model")
+                flash = gr.Dropdown(["Enabled","Disabled"], value = "Enabled",label ="Flash Attention 2.0 Acceleration")
+                temp =  gr.Slider(
+                    minimum=0.1,
+                    maximum=1,
+                    value=0.8,
+                    step=0.1,
+                    label="RAG Temperature"
                 )
                 model_button = gr.Button("Update Settings")
                 status2 = gr.Textbox(label="Update Status", interactive=False)
-
-        
-
         
         # Event handlers
         file_input.change(
@@ -296,10 +328,16 @@ def create_ui():
             inputs=[choice],
             outputs=[status1]
         )
+
+        db_button.click(
+            fn=app.dbupdate,
+            inputs=[metric_type,m_num,ef_num,topk],
+            outputs=[status3]
+        )
         
         model_button.click(
             fn=app.model_settings,
-            inputs=[hfchoice, ollamachoice,tokensize],
+            inputs=[hfchoice, ollamachoice,flash,temp],
             outputs=[status2]
         )
         
@@ -307,5 +345,6 @@ def create_ui():
 
 if __name__ == "__main__":
     demo = create_ui()
+    #demo.launch(auth=("admin", "pass1234")) for with login page config
     demo.launch()
 

colpali_manager.py

@@ -17,9 +17,16 @@ import spaces
 
 
 #this part is for local runs
+torch.cuda.empty_cache()
 
-model_name = "vidore/colSmol-256M"
-device = get_torch_device("cpu") #try using cpu instead of cpu?
+#get model name from .env variable & set directory & processor dir as the model names!
+import dotenv
+# Load the .env file
+dotenv_file = dotenv.find_dotenv()
+dotenv.load_dotenv(dotenv_file)
+
+model_name = os.environ['colpali']   #"vidore/colSmol-256M"
+device = get_torch_device("cuda") #try using cpu instead of cuda?
 
 #switch to locally downloading models & loading locally rather than from hf
 #
@@ -28,45 +35,48 @@ 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 = model_name+'_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
+        if "colSmol" in model_name: #if  colsmol
+            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))
+        else: #if colpali v1.3 etc
+            model = ColPali.from_pretrained(
+                model_name,
+                torch_dtype=torch.bfloat16,
+                device_map=device,
+                attn_implementation="flash_attention_2",
+            ).eval()
+            processor = cast(ColPaliProcessor, ColPaliProcessor.from_pretrained(model_name))
         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)
-"""
+    if "colSmol" in model_name:
+        model = ColIdefics3.from_pretrained(save_directory)
+        processor = ColIdefics3Processor.from_pretrained(processor_directory, use_fast=True)
+    else:
+        model = ColPali.from_pretrained(save_directory)
+        processor = ColPaliProcessor.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
+    def __init__(self, device = "cuda", model_name = model_name): #need to hot potato/use diff gpus between colpali & ollama
 
         print(f"Initializing ColpaliManager with device {device} and model {model_name}")
 
@@ -82,12 +92,12 @@ class ColpaliManager:
 
     @spaces.GPU
     def get_images(self, paths: list[str]) -> List[Image.Image]:
-        model.to("cpu")
+        model.to("cuda")
         return [Image.open(path) for path in paths]
 
     @spaces.GPU
     def process_images(self, image_paths:list[str], batch_size=5):
-        model.to("cpu")
+        model.to("cuda")
         print(f"Processing {len(image_paths)} image_paths")
         
         images = self.get_images(image_paths)
@@ -113,7 +123,7 @@ class ColpaliManager:
 
     @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)
+        model.to("cuda") #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(

milvus_manager.py

@@ -2,11 +2,18 @@ from pymilvus import MilvusClient, DataType
 import numpy as np
 import concurrent.futures
 from pymilvus import Collection
+import os
 
 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")
+  
+        #import environ variables from .env
+        import dotenv
+         # Load the .env file
+        dotenv_file = dotenv.find_dotenv()
+        dotenv.load_dotenv(dotenv_file)
+
+        self.client = MilvusClient(uri="http://localhost:19530", token="root:Milvus")
         self.collection_name = collection_name
         self.dim = dim
 
@@ -40,13 +47,15 @@ class MilvusManager:
 
     def create_index(self):
         index_params = self.client.prepare_index_params()
+
         index_params.add_index(
             field_name="vector",
             index_name="vector_index",
-            index_type="IVF_FLAT", #use HNSW option if got more mem, if not use IVF for faster processing
-            metric_type="IP",
+            index_type="HNSW", #use HNSW option if got more mem, if not use IVF for faster processing
+            metric_type=os.environ["metrictype"], #"IP"
             params={
-                "nlist": 1024
+                "M": int(os.environ["mnum"]), #M:16 for HNSW, capital M
+                "efConstruction": int(os.environ["efnum"]), #500 for HNSW
             },
         )
 
@@ -59,7 +68,7 @@ class MilvusManager:
         collections = self.client.list_collections()
         
         # Set search parameters (here, using Inner Product metric).
-        search_params = {"metric_type": "IP", "params": {}}
+        search_params = {"metric_type": os.environ["metrictype"], "params": {}} #default metric type is "IP"
         
         # Set to store unique (doc_id, collection_name) pairs across all collections.
         doc_collection_pairs = set()
@@ -71,7 +80,7 @@ class MilvusManager:
             results = self.client.search(
                 collection,
                 data,
-                limit=50,  # Adjust limit per collection as needed.
+                limit=int(os.environ["topk"]),  # Adjust limit per collection as needed. (default is 50)
                 output_fields=["vector", "seq_id", "doc_id"],
                 search_params=search_params,
             )

rag.py

@@ -1,10 +1,16 @@
 import requests
 import os
-import google.generativeai as genai
 
 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
 
 
 
@@ -35,8 +41,14 @@ class Rag:
             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):
+        #import environ variables from .env
+        import dotenv
+
+         # Load the .env file
+        dotenv_file = dotenv.find_dotenv()
+        dotenv.load_dotenv(dotenv_file)
         """ #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}")
 
@@ -65,8 +77,12 @@ class Rag:
 
         #ollama method below
         
-           
-        os.environ['OLLAMA_FLASH_ATTENTION'] = '1'
+        torch.cuda.empty_cache() #release cuda so that ollama can use gpu!
+
+    
+        os.environ['OLLAMA_FLASH_ATTENTION'] = os.environ['flashattn'] #int "1"
+        if os.environ['ollama'] == "minicpm-v":
+            os.environ['ollama'] = "minicpm-v:8b-2.6-q8_0" #set to quantized version
         
 
         # Close model thread (colpali)
@@ -74,13 +90,14 @@ class Rag:
 
         try:    
             
-            response = chat(    
-                model='minicpm-v:8b-2.6-q8_0',
-                messages=[
+            response = chat(  
+                    model=os.environ['ollama'],
+                    messages=[
                     {
                     'role': 'user',
                     'content': query,
                     'images': imagesPaths,
+                    "temperature":float(os.environ['temperature']), #test if temp makes a diff
                     }
                 ],
                 )
@@ -136,4 +153,4 @@ class 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)
+#     rag.get_answer_from_gemini(query, imagesPaths)

requirements.txt

@@ -1,9 +1,15 @@
+gradio
+PyMuPDF
+pdf2image
+pymilvus
+tqdm
+pillow
+spaces
+google-generativeai
 git+https://github.com/illuin-tech/colpali
-gradio==4.25.0
-PyMuPDF==1.24.9
-pdf2image==1.17.0
-pymilvus==2.4.9
-tqdm==4.66.5
-pillow==10.4.0
-spaces==0.30.4
-google-generativeai==0.8.3
+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
+colpali-engine[interpretability]