Update app.py

app.py

@@ -1,6 +1,5 @@
 from sentence_transformers import SentenceTransformer, util
 from huggingface_hub import hf_hub_download
-import os
 import pickle
 import pandas as pd
 from PIL import Image
@@ -8,33 +7,69 @@ import requests
 from io import BytesIO
 import gradio as gr
 
-pd.options.mode.chained_assignment = None  # Turn off SettingWithCopyWarning
+# Silence SettingWithCopyWarning from pandas
+pd.options.mode.chained_assignment = None  
 
-embeddings = pickle.load(open(hf_hub_download("bhavyagiri/semantic-memes", repo_type="dataset", filename="meme-embeddings.pkl"), "rb"))
+# Load meme embeddings (pre-computed)
+embeddings = pickle.load(open(
+    hf_hub_download("bhavyagiri/semantic-memes", repo_type="dataset", filename="meme-embeddings.pkl"), "rb"))
+
+# Load meme metadata (with 'id' and 'url' columns)
 df = pd.read_csv(hf_hub_download("bhavyagiri/semantic-memes", repo_type="dataset", filename="input.csv"))
 
+# Load sentence transformer model
 model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')
 
+# Meme search function
 def generate_memes(prompt):
+    # Encode user prompt into embedding
     prompt_embedding = model.encode(prompt, convert_to_tensor=True)
+    
+    # Perform semantic search
     hits = util.semantic_search(prompt_embedding, embeddings, top_k=6)
-    hits = pd.DataFrame(hits[0], columns=['corpus_id', 'score'])
-    desired_ids = hits["corpus_id"]
-    filtered_df = df.loc[df['id'].isin(desired_ids)]
-    filtered_list = list(filtered_df["url"]) 
-    images = [Image.open(BytesIO(requests.get(img).content)) for img in filtered_list]
-    return (
-       images
-    )
-input_textbox = gr.inputs.Textbox(lines=1, label="Search something cool")
-output_gallery = gr.Gallery(
-            label="Retrieved Memes", show_label=False, elem_id="gallery"
-        ).style(columns=[3], rows=[2], object_fit="contain", height="auto")
+    hits_df = pd.DataFrame(hits[0], columns=['corpus_id', 'score'])
+    
+    # Get matching meme URLs from original DataFrame
+    matched_ids = hits_df['corpus_id']
+    matched_memes = df[df['id'].isin(matched_ids)]
+    
+    # Download and display meme images
+    images = []
+    for url in matched_memes["url"]:
+        try:
+            response = requests.get(url)
+            image = Image.open(BytesIO(response.content))
+            images.append(image)
+        except Exception as e:
+            print(f"Error loading image from {url}: {e}")
+            continue
+    
+    return images
+
+# Gradio UI setup
+input_textbox = gr.Textbox(lines=1, label="Search something cool")
+
+output_gallery = gr.Gallery(label="Retrieved Memes").style(
+    columns=[3], rows=[2], object_fit="contain", height="auto"
+)
+
+# App info
 title = "Semantic Search for Memes"
-description = "Search Memes from small dataset of 6k memes. Check out [GitHub Repo](https://github.com/bhavya-giri/retrieving-memes)"
-examples = ["Get Shreked","Going Crazy","Spiderman is my teacher"]
-interpretation='default'
-enable_queue=True
+description = "Search memes from a dataset of ~6k memes using semantic similarity. [GitHub Repo](https://github.com/bhavya-giri/retrieving-memes)"
+examples = ["Get Shreked", "Going Crazy", "Spiderman is my teacher"]
+
+# Gradio interface
+iface = gr.Interface(
+    fn=generate_memes,
+    inputs=input_textbox,
+    outputs=output_gallery,
+    examples=examples,
+    cache_examples=True,
+    title=title,
+    description=description,
+    interpretation='default',
+    enable_queue=True
+)
 
-iface = gr.Interface(fn=generate_memes, inputs=input_textbox, outputs=output_gallery,examples=examples,cache_examples=True,title=title,description=description,interpretation=interpretation,enable_queue=enable_queue)
-iface.launch(inline=False)
+# Launch the app
+iface.launch()