rerank model

RAG/colpali.py

@@ -66,7 +66,7 @@ runtime = boto3.client("sagemaker-runtime",aws_access_key_id=st.secrets['user_ac
 # Prepare your payload (e.g., text-only input)
 
 
-
+@st.cache_resource
 def call_nova(
     model,
     messages,
@@ -110,13 +110,14 @@ def call_nova(
             modelId=model, body=json.dumps(request_body)
         )
         return response["body"]
+@st.cache_resource        
 def get_base64_encoded_value(media_path):
     with open(media_path, "rb") as media_file:
         binary_data = media_file.read()
         base_64_encoded_data = base64.b64encode(binary_data)
         base64_string = base_64_encoded_data.decode("utf-8")
         return base64_string
-
+@st.cache_resource
 def generate_ans(top_result,query):
     print(query)
     system_message = "given an image of a PDF page, answer the question. Be accurate to the question. If you don't find the answer in the page, please say, I don't know"
@@ -146,7 +147,7 @@ def generate_ans(top_result,query):
     print(content_text)
     return content_text
 
-    
+@st.cache_resource    
 def colpali_search_rerank(query):
     # Convert to JSON string
     payload = {
@@ -228,7 +229,7 @@ def colpali_search_rerank(query):
         return {'text':ans,'source':img,'image':images_highlighted,'table':[]}#[{'file':img}]
     
 
-
+@st.cache_resource
 def img_highlight(img,batch_queries,query_tokens):
     # Reference from : https://github.com/tonywu71/colpali-cookbooks/blob/main/examples/gen_colpali_similarity_maps.ipynb
     with open(img, "rb") as f:

RAG/rag_DocumentSearcher.py

@@ -12,7 +12,7 @@ headers = {"Content-Type": "application/json"}
 host = "https://search-opensearchservi-shjckef2t7wo-iyv6rajdgxg6jas25aupuxev6i.us-west-2.es.amazonaws.com/"
 
 parent_dirname = "/".join((os.path.dirname(__file__)).split("/")[0:-1])
-
+@st.cache_resource
 def query_(awsauth,inputs, session_id,search_types):
 
     print("using index: "+st.session_state.input_index)
@@ -219,49 +219,49 @@ def query_(awsauth,inputs, session_id,search_types):
         hits = response_['hits']['hits']
     
     ##### GET reference tables separately like *_mm index search for images  ######
-    def lazy_get_table():
-        table_ref = []
-        any_table_exists = False
-        for fname in os.listdir(parent_dirname+"/split_pdf_csv"):
-            if fname.startswith(st.session_state.input_index):
-                any_table_exists = True
-                break       
-        if(any_table_exists):
-            #################### Basic Match query #################
-            # payload_tables = {
-            #                     "query": {
-            #                         "bool":{
+    # def lazy_get_table():
+    #     table_ref = []
+    #     any_table_exists = False
+    #     for fname in os.listdir(parent_dirname+"/split_pdf_csv"):
+    #         if fname.startswith(st.session_state.input_index):
+    #             any_table_exists = True
+    #             break       
+    #     if(any_table_exists):
+    #         #################### Basic Match query #################
+    #         # payload_tables = {
+    #         #                     "query": {
+    #         #                         "bool":{
                                 
-            #                         "must":{"match": {
-            #                                         "processed_element": question
+    #         #                         "must":{"match": {
+    #         #                                         "processed_element": question
                                                 
-            #                                     }},
+    #         #                                     }},
                                                 
-            #                             "filter":{"term":{"raw_element_type": "table"}}
+    #         #                             "filter":{"term":{"raw_element_type": "table"}}
                                     
                                 
-            #                     }}}
+    #         #                     }}}
             
-            #################### Neural Sparse query #################
-            payload_tables = {"query":{"neural_sparse": {
-                    "processed_element_embedding_sparse": {
-                        "query_text": question,
-                        "model_id": "fkol-ZMBTp0efWqBcO2P"
-                    }
-                    }  }     }
+    #         #################### Neural Sparse query #################
+    #         payload_tables = {"query":{"neural_sparse": {
+    #                 "processed_element_embedding_sparse": {
+    #                     "query_text": question,
+    #                     "model_id": "fkol-ZMBTp0efWqBcO2P"
+    #                 }
+    #                 }  }     }
             
             
-            r_ = requests.get(url, auth=awsauth, json=payload_tables, headers=headers)
-            r_tables = json.loads(r_.text)
+    #         r_ = requests.get(url, auth=awsauth, json=payload_tables, headers=headers)
+    #         r_tables = json.loads(r_.text)
             
-            for res_ in r_tables['hits']['hits']:
-                if(res_["_source"]['raw_element_type'] == 'table'):
-                    table_ref.append({'name':res_["_source"]['table'],'text':res_["_source"]['processed_element']})
-                if(len(table_ref) == 2):
-                    break
+    #         for res_ in r_tables['hits']['hits']:
+    #             if(res_["_source"]['raw_element_type'] == 'table'):
+    #                 table_ref.append({'name':res_["_source"]['table'],'text':res_["_source"]['processed_element']})
+    #             if(len(table_ref) == 2):
+    #                 break
                     
             
-        return table_ref
+    #     return table_ref
         
         
     ########################### LLM Generation ########################

utilities/invoke_models.py

@@ -11,7 +11,7 @@ import streamlit as st
 #import torch
 
 region = 'us-east-1'
-
+@st.cache_resource
 bedrock_runtime_client = boto3.client(
     'bedrock-runtime',
     aws_access_key_id=st.secrets['user_access_key'],
@@ -30,7 +30,7 @@ bedrock_runtime_client = boto3.client(
 #     max_length = 16
 #     num_beams = 4
 #     gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
-
+@st.cache_resource
 def invoke_model(input):
     response = bedrock_runtime_client.invoke_model(
         body=json.dumps({
@@ -43,7 +43,7 @@ def invoke_model(input):
     
     response_body = json.loads(response.get("body").read())
     return response_body.get("embedding")
-
+@st.cache_resource
 def invoke_model_mm(text,img):
     body_ = {
             "inputText": text,
@@ -64,7 +64,7 @@ def invoke_model_mm(text,img):
     response_body = json.loads(response.get("body").read())
     #print(response_body)
     return response_body.get("embedding")
-
+@st.cache_resource
 def invoke_llm_model(input,is_stream):
     if(is_stream == False):
         response = bedrock_runtime_client.invoke_model( 
@@ -145,7 +145,7 @@ def invoke_llm_model(input,is_stream):
     #     stream = response.get('body')
         
     #     return stream
-        
+@st.cache_resource        
 def read_from_table(file,question):
     print("started table analysis:")
     print("-----------------------")
@@ -181,7 +181,7 @@ def read_from_table(file,question):
              )
     agent_res = agent.invoke(question)['output']
     return agent_res
-    
+@st.cache_resource    
 def generate_image_captions_llm(base64_string,question):
     
     # ant_client = Anthropic()

utilities/re_ranker.py

@@ -46,7 +46,7 @@ from sentence_transformers import CrossEncoder
 # print("Program ends.")
 #########################
 
-
+@st.cache_resource
 def re_rank(self_, rerank_type, search_type, question, answers):
    
     ans = []