Update app.py

app.py

@@ -9,18 +9,22 @@ from PIL import Image
 from io import BytesIO
 from sentence_transformers import SentenceTransformer
 
-# dataset = load_dataset("imagefolder", data_files="https://huggingface.co/datasets/nlphuji/flickr30k/blob/main/flickr30k-images.zip")
-
 # Load the pre-trained sentence encoder
 model_name = "sentence-transformers/all-distilroberta-v1"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = SentenceTransformer(model_name)
 
-# Load the FAISS index
-index_name = 'index.faiss'
-index_url = 'https://huggingface.co/spaces/shivangibithel/Text2ImageRetrieval/blob/main/faiss_flickr8k.index'
-wget.download(index_url, index_name)
-index = faiss.read_index(index_name)
+# # Load the FAISS index
+# index_name = 'index.faiss'
+# index_url = 'https://huggingface.co/spaces/shivangibithel/Text2ImageRetrieval/blob/main/faiss_flickr8k.index'
+# wget.download(index_url, index_name)
+# index = faiss.read_index(index_name)
+
+vectors = np.load("https://huggingface.co/spaces/shivangibithel/Text2ImageRetrieval/blob/main/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'
@@ -35,11 +39,13 @@ caption_list = list(caption_dict.values())
 
 def search(query, k=5):
     # Encode the query
-    query_tokens = tokenizer.encode(query, return_tensors='pt')
-    query_embedding = model.encode(query_tokens).detach().numpy()
+    query_embedding = model.encode(query)
+    query_vector = np.array([query_embedding])
+    faiss.normalize_L2(query_vector)
+    index.nprobe = index.ntotal
 
     # Search for the nearest neighbors in the FAISS index
-    D, I = index.search(query_embedding, k)
+    D, I = index.search(query_vector, k)
 
     # Map the image ids to the corresponding image URLs
     image_urls = []