# apne imp libraries
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 load kiye dataset repo se 
embeddings = pickle.load(open(
    hf_hub_download("Go-Raw/semantic-memes", repo_type="dataset", filename="meme-embeddings.pkl"), "rb"))

# apne meme ka metadata load kiya
df = pd.read_csv(
    hf_hub_download("Go-Raw/semantic-memes", repo_type="dataset", filename="semantic_memes.csv"))

# ye apna model hai
model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')

# iss func se meme search hota hai
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

# Gradio ka UI
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"
]

# app launch karne ke liye
iface = gr.Interface(
    fn=generate_memes,
    inputs=input_textbox,
    outputs=output_gallery,
    examples=examples,
    cache_examples=True,
    title=title,
    description=description
)

iface.launch()