adding new index and changing app a bit

app.py

@@ -8,23 +8,23 @@ from transformers import AutoModel, AutoFeatureExtractor
 seed = 42
 
 # Only runs once when the script is first run.
-with open("index.pickle", "rb") as handle:
+with open("index_768.pickle", "rb") as handle:
     index = pickle.load(handle)
 
 # Load model for computing embeddings.
-feature_extractor = AutoFeatureExtractor.from_pretrained("abhishek/autotrain-butterflies-new-17716425")
-model = AutoModel.from_pretrained("abhishek/autotrain-butterflies-new-17716425")
+feature_extractor = AutoFeatureExtractor.from_pretrained("sasha/autotrain-butterfly-similarity-2490576840")
+model = AutoModel.from_pretrained("sasha/autotrain-butterfly-similarity-2490576840")
 
 # Candidate images.
-dataset = load_dataset("sasha/butterflies_names_multiple")
+dataset = load_dataset("sasha/butterflies_10k_names_multiple")
 ds = dataset["train"]
 
 
-def query(image, top_k):
+def query(image, top_k=4):
     inputs = feature_extractor(image, return_tensors="pt")
     model_output = model(**inputs)
     embedding = model_output.pooler_output.detach()
-    results = index.query(embedding)
+    results = index.query(embedding, k=top_k)
     inx = results[0][0].tolist()
     images = ds.select(inx)["image"]
     return images
@@ -37,8 +37,8 @@ description = "This Space demos an image similarity system. You can refer to [th
 # Not sure what the best for this demo is.
 gr.Interface(
     query,
-    inputs=[gr.Image(type="pil"), gr.Slider(value=5, minimum=1, maximum=10, step=1)],
-    outputs=gr.Gallery().style(grid=[3], height="auto"),
+    inputs=[gr.Image(type="pil")],
+    outputs=gr.Gallery().style(grid=[2], height="auto"),
     # Filenames denote the integer labels. Know here: https://hf.co/datasets/beans
     title=title,
     description=description,

index_768.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eccd83bb743f6de5eaf05886f948d95daceeede60805ad01cdba0baddd1a60cc
+size 53317256

requirements.txt

@@ -2,4 +2,5 @@ transformers==4.25.1
 datasets==2.7.1
 numpy==1.21.6
 torch==1.12.1
-torchvision
+torchvision
+pynndescent