Update app.py

app.py

@@ -14,7 +14,7 @@ extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
 model = AutoModel.from_pretrained(model_ckpt)
 hidden_dim = model.config.hidden_size
 
-
+# Load dataset
 dataset_with_embeddings = load_dataset("LucyintheSky/24-1-8-ds-embeddings", split="train", token=os.environ.get('TOKEN'))
 dataset_with_embeddings.add_faiss_index(column='embeddings')
 
@@ -28,14 +28,14 @@ def get_neighbors(query_image, top_k=8):
     
 
 def search(image_dict):
-    
+
+    # Open query image
     query_image = Image.open(image_dict['composite']).convert(mode='RGB')
 
-    print('search')
+    # Get similar image
     scores, retrieved_examples = get_neighbors(query_image)
 
-    print('return example')
-
+    # Create result diction for gr.Gallery
     result = []
     for i in range(len(retrieved_examples["image"])):
         id = str(retrieved_examples["text"][i]) + ' ' + str(scores[i])