Bugfix, comma

src/app.py

@@ -49,14 +49,14 @@ def check_index(ds):
 
     else:
         return index_dataset(ds)
-    
+
 
 dataset_with_embeddings = check_index(candidate_subset)
 
 # Main function, to find similar images
 # TODO: implement different distance measures
 
-def get_neighbors(query_image, selected_descriptor, selected_distance top_k=5):
+def get_neighbors(query_image, selected_descriptor, selected_distance, top_k=5):
     """Returns the top k nearest examples to the query image.
 
     Args:
@@ -75,10 +75,10 @@ def get_neighbors(query_image, selected_descriptor, selected_distance top_k=5):
                 'color_embeddings', qi_np, k=top_k)
         elif selected_distance == "Chi-squared":
             tmp_dataset = dataset_with_embeddings.map(lambda row: {'distance': chi2_distance(histA=query_vector, histB=row['color_embeddings'])})
-            retrieved_examples = tmp_dataset.sort("distance")
-        else: 
+            retrieved_examples = tmp_dataset.sort("distance")[:5]
+        else:
             tmp_dataset = dataset_with_embeddings.map(lambda row: {'distance': euclidian_distance(histA=query_vector, histB=row['color_embeddings'])})
-            retrieved_examples = tmp_dataset.sort("distance")
+            retrieved_examples = tmp_dataset.sort("distance")[:5]
         images = retrieved_examples['image'] #retrieved images is a dict, with images and embeddings
         return images
     if "CLIP" == selected_descriptor: