rerank model

RAG/colpali.py

@@ -19,11 +19,11 @@ 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 +286,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