trying

app.py

@@ -1,10 +1,8 @@
-
-
 import pickle
 
 import gradio as gr
 from datasets import load_dataset
-from transformers import AutoModel
+from transformers import AutoModel, AutoFeatureExtractor
 
 
 seed = 42
@@ -18,38 +16,21 @@ feature_extractor = AutoFeatureExtractor.from_pretrained("abhishek/autotrain-but
 model = AutoModel.from_pretrained("abhishek/autotrain-butterflies-new-17716425")
 
 # Candidate images.
-dataset = load_dataset("huggan/inat_butterflies_top10k")
-candidate_dataset = dataset["train"]
+dataset = load_dataset("sasha/butterflies_names_multiple")
+ds = dataset["train"]
 
 
 def query(image, top_k):
     inputs = feature_extractor(image, return_tensors="pt")
     model_output = model(**inputs)
-    embedding = model_output.pooler_output
+    embedding = model_output.pooler_output.detach()
     results = index.query(embedding)
-
-    # Should be a list of string file paths for gr.Gallery to work
-    images = []
-    # List of labels for each image in the gallery
-    labels = []
-
-    candidates = []
-
-    for idx, r in enumerate(sorted(results, key=results.get, reverse=True)):
-        if idx == top_k:
-            break
-        image_id, label = r.split("_")[0], r.split("_")[1]
-        candidates.append(candidate_dataset[int(image_id)]["image"])
-        labels.append(f"Label: {label}")
-
-    for i, candidate in enumerate(candidates):
-        filename = f"similar_{i}.png"
-        candidate.save(filename)
-        images.append(filename)
-
-    # The gallery component can be a list of tuples, where the first element is a path to a file
-    # and the second element is an optional caption for that image
-    return list(zip(images, labels))
+    images=[]
+    for i in results[0].tolist():
+    	print(i)
+    	print(type(i))
+    	images.append(ds.select(i)["image"])
+    return images
 
 
 title = "Find my Butterfly 🦋"