renamed pkl file and rename in app.py

app.py

@@ -9,8 +9,12 @@ from sklearn.metrics.pairwise import cosine_similarity
 import csv
 from PIL import Image
 
-model_path = "kaveh/rclip"
-embeddings_file = './image_embeddings_8_clip14_cxrbert.pkl' 
+model_path_rclip = "kaveh/rclip"
+embeddings_file_rclip = './image_embeddings_rclip.pkl' 
+
+model_path_pubmedclip = "flaviagiammarino/pubmed-clip-vit-base-patch32"
+embeddings_file_pubmedclip = './image_embeddings_pubmedclip.pkl' 
+
 csv_path = "./captions.txt"
 
 def load_image_ids(csv_file):
@@ -36,13 +40,19 @@ def find_similar_images(query_embedding, image_embeddings, k=2):
     return closest_indices, scores
 
 
-def main(query, k=2):
-    # Load RCLIP model
-    model = VisionTextDualEncoderModel.from_pretrained(model_path)
-    processor = VisionTextDualEncoderProcessor.from_pretrained(model_path)
+def main(query, model_id="rclip", k=2):
+    if model_id=="rclip":
+        # Load RCLIP model
+        model = VisionTextDualEncoderModel.from_pretrained(model_path_rclip)
+        processor = VisionTextDualEncoderProcessor.from_pretrained(model_path_rclip)
+        # Load image embeddings 
+        image_embeddings = load_embeddings(embeddings_file_rclip)
+    elif mode_id=="pubmedclip":
+        model = CLIPModel.from_pretrained(model_path_pubmedclip)
+        processor = CLIPProcessor.from_pretrained(model_path_pubmedclip)
+        # Load image embeddings 
+        image_embeddings = load_embeddings(embeddings_file_pubmedclip)   
 
-    # Load image embeddings 
-    image_embeddings = load_embeddings(embeddings_file)
 
     # Embed the query
     inputs = processor(text=query, images=None, return_tensors="pt", padding=True)