Update app.py

app.py

@@ -5,30 +5,11 @@ import torch
 from transformers import AutoTokenizer, AutoModel
 import faiss
 import numpy as np
-import wget
 from PIL import Image
 from sentence_transformers import SentenceTransformer
 import json
-from zipfile import ZipFile
 import zipfile
 
-# Load the pre-trained sentence encoder
-model_name = "sentence-transformers/all-distilroberta-v1"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = SentenceTransformer(model_name)
-
-# Define the path to the zip folder containing the images
-zip_path = "Images.zip"
-
-# Open the zip folder
-zip_file = zipfile.ZipFile(zip_path)
-
-vectors = np.load("./sbert_text_features.npy")
-vector_dimension = vectors.shape[1]
-index = faiss.IndexFlatL2(vector_dimension)
-faiss.normalize_L2(vectors)
-index.add(vectors)
-
 # Map the image ids to the corresponding image URLs
 image_map_name = 'captions.json'
 
@@ -37,8 +18,20 @@ with open(image_map_name, 'r') as f:
 
 image_list = list(caption_dict.keys())
 caption_list = list(caption_dict.values())
+zip_path = "Images.zip"
+zip_file = zipfile.ZipFile(zip_path)
+
+model_name = "sentence-transformers/all-distilroberta-v1"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = SentenceTransformer(model_name)
+vectors = model.encode(caption_list)
+# vectors = np.load("./sbert_text_features.npy")
+vector_dimension = vectors.shape[1]
+index = faiss.IndexFlatL2(vector_dimension)
+faiss.normalize_L2(vectors)
+index.add(vectors)
 
-def search(query, k=5):
+def search(query, k=4):
     # Encode the query
     query_embedding = model.encode(query)
     query_vector = np.array([query_embedding])