Update app.py

app.py

@@ -9,34 +9,29 @@ from PIL import Image
 from io import BytesIO
 from sentence_transformers import SentenceTransformer
 
-# dataset = load_dataset("nlphuji/flickr30k", streaming=True)
-# df = pd.DataFrame.from_dict(dataset["train"])
+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/paraphrase-multilingual-mpnet-base-v2"
+model_name = "sentence-transformers/all-distilroberta-v1"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = SentenceTransformer(model_name)
 
-# # Load the pre-trained image model
-# image_model_name = 'image_model.ckpt'
-# image_model_url = 'https://huggingface.co/models/flax-community/deit-tiny-random/images/vqvae.png'
-# wget.download(image_model_url, image_model_name)
-# image_model = torch.load(image_model_name, map_location=torch.device('cpu'))
-# image_model.eval()
-
 # Load the FAISS index
 index_name = 'index.faiss'
-index_url = 'https://huggingface.co/models/flax-community/deit-tiny-random/faiss_files/faiss.index'
+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)
 
 # Map the image ids to the corresponding image URLs
-image_map_name = 'image_map.json'
-image_map_url = 'https://huggingface.co/models/flax-community/deit-tiny-random/faiss_files/image_map.json'
+image_map_name = 'captions.json'
+image_map_url = 'https://huggingface.co/spaces/shivangibithel/Text2ImageRetrieval/blob/main/captions.json'
 wget.download(image_map_url, image_map_name)
-image_map = {}
+
 with open(image_map_name, 'r') as f:
-    image_map = json.load(f)
+    caption_dict = json.load(f)
+
+image_list = list(caption_dict.keys())
+caption_list = list(caption_dict.values())
 
 def search(query, k=5):
     # Encode the query
@@ -49,8 +44,9 @@ def search(query, k=5):
     # Map the image ids to the corresponding image URLs
     image_urls = []
     for i in I[0]:
-        image_id = str(i)
-        image_url = image_map[image_id]
+        text_id = i
+        image_id = str(image_list[i])
+        image_url = "https://huggingface.co/spaces/shivangibithel/Text2ImageRetrieval/blob/main/Images/" + image_id
         image_urls.append(image_url)
 
     return image_urls