Update app.py

app.py

@@ -1,60 +1,32 @@
-import pickle
 import gradio as gr
 from datasets import load_dataset
-from transformers import AutoModel, AutoFeatureExtractor
-import wikipedia
+from sentence_transformers import SentenceTransformer
 
-
-# Only runs once when the script is first run.
-with open("index_768_cosine.pickle", "rb") as handle:
-    index = pickle.load(handle)
-
-# Load model for computing embeddings.
-feature_extractor = AutoFeatureExtractor.from_pretrained("sasha/autotrain-butterfly-similarity-2490576840")
-model = AutoModel.from_pretrained("sasha/autotrain-butterfly-similarity-2490576840")
+model = SentenceTransformer('clip-ViT-B-32')
 
 # Candidate images.
-dataset = load_dataset("sasha/butterflies_10k_names_multiple")
+dataset = load_dataset("sasha/pedro-embeddings")
 ds = dataset["train"]
+ds.add_faiss_index(column='embeddings')
 
 
-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, k=top_k)
-    inx = results[0][0].tolist()
-    logits = results[1][0].tolist()
-    images = ds.select(inx)["image"]
-    captions = ds.select(inx)["name"]
-    images_with_captions = [(i, c) for i, c in zip(images,captions)]
-    labels_with_probs = dict(zip(captions,logits))
-    labels_with_probs = {k: 1- v for k, v in labels_with_probs.items()}
-    try:
-    	description = wikipedia.summary(captions[0], sentences = 1)
-    	description = "### " + description
-    	url = wikipedia.page(captions[0]).url
-    	url = " You can learn more about your butterfly [here](" + str(url) + ")!"
-    	description = description + url
-    except:
-    	description = "### Butterflies are insects in the order Lepidoptera, which also includes moths. Adult butterflies have large, often brightly coloured wings."
-    	url = "https://en.wikipedia.org/wiki/Butterfly"
-    	url = " You can learn more about butterflies [here](" + str(url) + ")!"
-    	description = description + url
-    return images_with_captions, labels_with_probs, description
+def query(image, number_to_retrieve=1):
+    prompt = model.encode(image)
+    scores, retrieved_examples = ds.get_nearest_examples('embeddings', prompt, k=number_to_retrieve)
+    return retrieved_examples
 
 
 with gr.Blocks() as demo:
-	gr.Markdown("# Find my Butterfly 🦋")
-	gr.Markdown("## Use this Space to find your butterfly, based on the [iNaturalist butterfly dataset](https://huggingface.co/datasets/huggan/inat_butterflies_top10k)!")
+	gr.Markdown("# Find my Pedro Pascal")
+	gr.Markdown("## Use this Space to find the Pedro Pascal most similar to your input image!")
 	with gr.Row():
 		with gr.Column(min_width= 900):
 			inputs = gr.Image(shape=(800, 1600))
-			btn = gr.Button("Find my butterfly!")
+			btn = gr.Button("Find my Pedro!")
 			description = gr.Markdown()
 			
 		with gr.Column():
-			outputs=gr.Gallery().style(grid=[2], height="auto")
+			outputs=gr.Gallery().style(grid=[1], height="auto")
 			labels = gr.Label()
 			
 	gr.Markdown("### Image Examples")