colpali fix

RAG/colpali.py

@@ -156,88 +156,6 @@ 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 = {
-    "queries": [query]
-    }
-    body = json.dumps(payload)
-
-    # Call the endpoint
-    response = runtime.invoke_endpoint(
-        EndpointName=endpoint_name,
-        ContentType="application/json",
-        Body=body
-    )
-
-    # Read and print the response
-    result = json.loads(response["Body"].read().decode())
-    #print(len(result['query_embeddings'][0]))
-    
-    final_docs_sorted_20 = []
-    for i in result['query_embeddings']:
-        batch_embeddings = i
-        a = np.array(batch_embeddings)
-        vec = a.mean(axis=0)
-        #print(vec)
-        hits = []
-        #for v in batch_embeddings:
-        query_ = {
-            "size": 200,
-            "query": {
-                "nested": {
-                "path": "page_sub_vectors",
-                "query": {
-                    "knn": {
-                    "page_sub_vectors.page_sub_vector": {
-                        "vector": vec.tolist(),
-                        "k": 200
-                    }
-                    }
-                }
-                    }
-                }
-                }
-        response = aos_client.search(
-            body = query_,
-            index = 'colpali-vs'
-        )
-        #print(response)
-        query_token_vectors = batch_embeddings
-        final_docs = []
-        hits += response['hits']['hits']
-        #print(len(hits))
-        for ind,j in enumerate(hits):
-            max_score_dict_list = []
-            doc={"id":j["_id"],"score":j["_score"],"image":j["_source"]["image"]}
-            with_s = j['_source']['page_sub_vectors']
-            add_score = 0
-
-            for index,i in enumerate(query_token_vectors):
-                query_token_vector = np.array(i)
-                scores = []
-                for m in with_s:
-                    doc_token_vector = np.array(m['page_sub_vector'])
-                    score = np.dot(query_token_vector,doc_token_vector)
-                    scores.append(score)
-                    
-                scores.sort(reverse=True)
-                max_score = scores[0]
-                add_score+=max_score
-            doc["total_score"] = add_score
-            final_docs.append(doc)
-        final_docs_sorted = sorted(final_docs, key=lambda d: d['total_score'], reverse=True)
-        final_docs_sorted_20.append(final_docs_sorted[:20])
-        img = "/home/user/app/vs/"+final_docs_sorted_20[0][0]['image']
-        ans = generate_ans(img,query)
-        images_highlighted = [{'file':img}]
-        st.session_state.top_img = img
-        st.session_state.query_token_vectors = query_token_vectors
-        st.session_state.query_tokens = result['query_tokens']
-        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
@@ -331,3 +249,94 @@ def img_highlight(img,batch_queries,query_tokens):
         
 
             
+
+@st.cache_resource    
+def colpali_search_rerank(query):
+    if(st.session_state.show_columns == True):
+        print("show columns activated------------------------")
+        st.session_state.maxSimImages = img_highlight(st.session_state.top_img, st.session_state.query_token_vectors, st.session_state.query_tokens)
+        st.session_state.show_columns = False
+        st.session_state.input_query = st.session_state.questions_[-1]["question"]
+        st.session_state.answers_.pop()
+        st.session_state.questions_.pop()
+        return {'text':st.session_state.answers_[0]['answer'],'source':st.session_state.answers_[0]['source'],'image':st.session_state.maxSimImages,'table':[]}
+    # Convert to JSON string
+    payload = {
+    "queries": [query]
+    }
+    body = json.dumps(payload)
+
+    # Call the endpoint
+    response = runtime.invoke_endpoint(
+        EndpointName=endpoint_name,
+        ContentType="application/json",
+        Body=body
+    )
+
+    # Read and print the response
+    result = json.loads(response["Body"].read().decode())
+    #print(len(result['query_embeddings'][0]))
+    
+    final_docs_sorted_20 = []
+    for i in result['query_embeddings']:
+        batch_embeddings = i
+        a = np.array(batch_embeddings)
+        vec = a.mean(axis=0)
+        #print(vec)
+        hits = []
+        #for v in batch_embeddings:
+        query_ = {
+            "size": 200,
+            "query": {
+                "nested": {
+                "path": "page_sub_vectors",
+                "query": {
+                    "knn": {
+                    "page_sub_vectors.page_sub_vector": {
+                        "vector": vec.tolist(),
+                        "k": 200
+                    }
+                    }
+                }
+                    }
+                }
+                }
+        response = aos_client.search(
+            body = query_,
+            index = 'colpali-vs'
+        )
+        #print(response)
+        query_token_vectors = batch_embeddings
+        final_docs = []
+        hits += response['hits']['hits']
+        #print(len(hits))
+        for ind,j in enumerate(hits):
+            max_score_dict_list = []
+            doc={"id":j["_id"],"score":j["_score"],"image":j["_source"]["image"]}
+            with_s = j['_source']['page_sub_vectors']
+            add_score = 0
+
+            for index,i in enumerate(query_token_vectors):
+                query_token_vector = np.array(i)
+                scores = []
+                for m in with_s:
+                    doc_token_vector = np.array(m['page_sub_vector'])
+                    score = np.dot(query_token_vector,doc_token_vector)
+                    scores.append(score)
+                    
+                scores.sort(reverse=True)
+                max_score = scores[0]
+                add_score+=max_score
+            doc["total_score"] = add_score
+            final_docs.append(doc)
+        final_docs_sorted = sorted(final_docs, key=lambda d: d['total_score'], reverse=True)
+        final_docs_sorted_20.append(final_docs_sorted[:20])
+        img = "/home/user/app/vs/"+final_docs_sorted_20[0][0]['image']
+        ans = generate_ans(img,query)
+            
+        images_highlighted = [{'file':img}]
+        st.session_state.top_img = img
+        st.session_state.query_token_vectors = query_token_vectors
+        st.session_state.query_tokens = result['query_tokens']
+        return {'text':ans,'source':img,'image':images_highlighted,'table':[]}#[{'file':img}]
+    

pages/Multimodal_Conversational_Search.py

@@ -281,18 +281,18 @@ def render_answer(question,answer,index,res_img):
                               for _ in range(10)])
         def show_maxsim():
             st.session_state.show_columns = True
-            st.session_state.maxSimImages = colpali.img_highlight(st.session_state.top_img, st.session_state.query_token_vectors, st.session_state.query_tokens)
-            st.session_state.input_query = st.session_state.questions_[-1]["question"]
-            st.session_state.answers_.pop()
-            st.session_state.questions_.pop()
+            # st.session_state.input_query = st.session_state.questions_[-1]["question"]
+            # st.session_state.answers_.pop()
+            # st.session_state.questions_.pop()
             handle_input()
-            with placeholder.container():
-                render_all() 
             print("*"*20)
             print(st.session_state.input_query) 
             print(st.session_state.answers_)
             print(st.session_state.questions_)
             print("*"*20) 
+            with placeholder.container():
+                render_all() 
+            
         if(st.session_state.input_is_colpali):
             st.button("Show similarity map",key=rdn_key_1,on_click = show_maxsim)
                 
@@ -309,7 +309,7 @@ def render_answer(question,answer,index,res_img):
                     for i in range(0,len(res_img)):
                         
                         if(st.session_state.input_is_colpali):
-                            if(st.session_state.show_columns == True):
+                            if(len(res_img)>1):
                                 cols_per_row = 3
                                 st.session_state.image_placeholder=st.empty()
                                 with st.session_state.image_placeholder.container():
@@ -324,11 +324,6 @@ def render_answer(question,answer,index,res_img):
                                     col3_,col4_,col5_ = st.columns([33,33,33])
                                     with col3_:
                                         st.image(res_img[i]['file'])
-                                
-                                
-
-                            
-                            
                         else:
                             if(res_img[i]['file'].lower()!='none' and idx < 2):
                                 col3,col4,col5 = st.columns([33,33,33])
@@ -398,10 +393,7 @@ def render_answer(question,answer,index,res_img):
      
 #Each answer will have context of the question asked in order to associate the provided feedback with the respective question
 def write_chat_message(md, q,index):
-    if(st.session_state.show_columns):
-        res_img = st.session_state.maxSimImages
-    else:
-        res_img = md['image']
+    res_img = md['image']
     #st.session_state['session_id'] = res['session_id']   to be added in memory
     chat = st.container()
     with chat: