rerank model

RAG/colpali.py

@@ -17,13 +17,12 @@ import matplotlib.pyplot as plt
 import requests
 import boto3
 import streamlit as st
-from IPython.display import display, Markdown
 import base64
-# from colpali_engine.interpretability import (
-#     get_similarity_maps_from_embeddings,
-#     plot_all_similarity_maps,
-#     plot_similarity_map,
-# )
+from colpali_engine.interpretability import (
+    get_similarity_maps_from_embeddings,
+    plot_all_similarity_maps,
+    plot_similarity_map,
+)
 import torch
 # from colpali_engine.models import ColPali, ColPaliProcessor
 # from colpali_engine.utils.torch_utils import get_torch_device
@@ -286,37 +285,37 @@ def img_highlight(img,batch_queries,query_tokens):
     print(f"Number of image patches: {n_patches}")
 
     # # Generate the similarity maps
-    # batched_similarity_maps = get_similarity_maps_from_embeddings(
-    #     image_embeddings=image_embeddings,
-    #     query_embeddings=query_embeddings,
-    #     n_patches=n_patches,
-    #     image_mask = image_mask
-    # )
+    batched_similarity_maps = get_similarity_maps_from_embeddings(
+        image_embeddings=image_embeddings,
+        query_embeddings=query_embeddings,
+        n_patches=n_patches,
+        image_mask = image_mask
+    )
 
-    # # # Get the similarity map for our (only) input image
-    # similarity_maps = batched_similarity_maps[0]  # (query_length, n_patches_x, n_patches_y)
+    # # Get the similarity map for our (only) input image
+    similarity_maps = batched_similarity_maps[0]  # (query_length, n_patches_x, n_patches_y)
 
-    # query_tokens_from_model = query_tokens[0]['tokens']
+    query_tokens_from_model = query_tokens[0]['tokens']
     
-    # plots = plot_all_similarity_maps(
-    # image=image,
-    # query_tokens=query_tokens_from_model,
-    # similarity_maps=similarity_maps,
-    # figsize=(8, 8),
-    # show_colorbar=False,
-    # add_title=True,
-    # )
-    # map_images = []
-    # for idx, (fig, ax) in enumerate(plots):
-    #     if(idx<3):
-    #         continue
-    #     savepath = "/home/user/app/similarity_maps/similarity_map_"+(img.split("/"))[-1]+"_token_"+str(idx)+"_"+query_tokens_from_model[idx]+".png"
-    #     fig.savefig(savepath, bbox_inches="tight")
-    #     map_images.append({'file':savepath})
-    #     print(f"Similarity map for token `{query_tokens_from_model[idx]}` saved at `{savepath}`")
+    plots = plot_all_similarity_maps(
+    image=image,
+    query_tokens=query_tokens_from_model,
+    similarity_maps=similarity_maps,
+    figsize=(8, 8),
+    show_colorbar=False,
+    add_title=True,
+    )
+    map_images = []
+    for idx, (fig, ax) in enumerate(plots):
+        if(idx<3):
+            continue
+        savepath = "/home/user/app/similarity_maps/similarity_map_"+(img.split("/"))[-1]+"_token_"+str(idx)+"_"+query_tokens_from_model[idx]+".png"
+        fig.savefig(savepath, bbox_inches="tight")
+        map_images.append({'file':savepath})
+        print(f"Similarity map for token `{query_tokens_from_model[idx]}` saved at `{savepath}`")
         
-    # plt.close("all")   
-    return ""#map_images 
+    plt.close("all")   
+    return map_images 
 
         
         

pages/Multimodal_Conversational_Search.py

@@ -13,7 +13,7 @@ import botocore.session
 import json
 import random
 import string
-import rag_DocumentLoader 
+#import rag_DocumentLoader 
 import rag_DocumentSearcher
 import pandas as pd
 from PIL import Image