|
|
|
from sentence_transformers import SentenceTransformer, util |
|
from huggingface_hub import hf_hub_download |
|
import pickle |
|
import pandas as pd |
|
from PIL import Image |
|
import requests |
|
from io import BytesIO |
|
import gradio as gr |
|
|
|
pd.options.mode.chained_assignment = None |
|
|
|
|
|
embeddings = pickle.load(open( |
|
hf_hub_download("Go-Raw/semantic-memes", repo_type="dataset", filename="meme-embeddings.pkl"), "rb")) |
|
|
|
|
|
df = pd.read_csv( |
|
hf_hub_download("Go-Raw/semantic-memes", repo_type="dataset", filename="semantic_memes.csv")) |
|
|
|
|
|
model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2') |
|
|
|
|
|
def generate_memes(prompt): |
|
prompt_embedding = model.encode(prompt, convert_to_tensor=True) |
|
hits = util.semantic_search(prompt_embedding, embeddings, top_k=6) |
|
hits_df = pd.DataFrame(hits[0], columns=["corpus_id", "score"]) |
|
matched_ids = hits_df["corpus_id"] |
|
matched_memes = df[df["id"].isin(matched_ids)] |
|
|
|
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 {url}: {e}") |
|
return images |
|
|
|
|
|
input_textbox = gr.Textbox(lines=1, label="Type your vibe here 🧠") |
|
output_gallery = gr.Gallery(label="Your Meme Results", columns=3, rows=2, height="auto") |
|
|
|
title = "🧠 Meme Lord" |
|
description = ( |
|
"Search memes from a diverse collection using sentence-level semantic similarity. " |
|
"Built with Sentence Transformers and hosted on Hugging Face. " |
|
"[Dataset](https://huggingface.co/datasets/Go-Raw/semantic-memes)" |
|
) |
|
examples = [ |
|
"When you realize it's Monday again", |
|
"Internet explorer in 2024", |
|
"This meeting could’ve been an email" |
|
] |
|
|
|
|
|
iface = gr.Interface( |
|
fn=generate_memes, |
|
inputs=input_textbox, |
|
outputs=output_gallery, |
|
examples=examples, |
|
cache_examples=True, |
|
title=title, |
|
description=description |
|
) |
|
|
|
iface.launch() |
|
|