Update app.py

app.py

@@ -10,7 +10,6 @@ import random
 st.title("Meme Battle AI")
 st.write("Stream memes directly and let AI determine the winner!")
 
-# Load the CLIP model
 @st.cache_resource
 def load_model():
     model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
@@ -19,65 +18,55 @@ def load_model():
 
 model, processor = load_model()
 
-# Stream the dataset
 @st.cache_resource
 def load_streamed_dataset():
     return load_dataset("Dhruv-goyal/memes_with_captions", split="train", streaming=True)
 
 dataset = load_streamed_dataset()
 
-# Randomly fetch two memes
 def fetch_random_memes():
-    samples = list(dataset.take(10))  # Fetch 10 samples from the streamed dataset
-    meme1, meme2 = random.sample(samples, 2)
+    sample_size = 100
+    dataset_samples = list(dataset.shuffle(seed=random.randint(0, 1000)).take(sample_size))
+    meme1, meme2 = random.sample(dataset_samples, 2)
     return meme1, meme2
 
-# Score the memes
+def parse_meme(meme):
+    caption = meme["answers"][0] if meme.get("answers") else "No caption available"
+    image_url = meme["image"]
+    return caption, image_url
+
 def score_meme(image_url, caption):
     try:
-        # Load the image from the URL
-        image = Image.open(BytesIO(requests.get(image_url).content))
-        # Preprocess image and caption
+        response = requests.get(image_url)
+        image = Image.open(BytesIO(response.content)).convert("RGB")
         inputs = processor(text=[caption], images=[image], return_tensors="pt", padding=True)
-        # Get the compatibility score
         outputs = model(**inputs)
         logits_per_text = outputs.logits_per_text
         return logits_per_text.item()
     except Exception as e:
+        st.error(f"Error scoring meme: {e}")
         return 0
 
-# Streamlit Interface
-st.write("### Select two memes and let the AI determine the winner!")
-
 if st.button("Start Meme Battle"):
     meme1, meme2 = fetch_random_memes()
+    caption1, image_url1 = parse_meme(meme1)
+    caption2, image_url2 = parse_meme(meme2)
 
-    # Extract caption and image for Meme 1
-    caption1 = meme1["answers"][0] if meme1["answers"] else "No caption available"
-    image_url1 = meme1["image"]
-
-    # Extract caption and image for Meme 2
-    caption2 = meme2["answers"][0] if meme2["answers"] else "No caption available"
-    image_url2 = meme2["image"]
-
-    # Fetch scores
     score1 = score_meme(image_url1, caption1)
     score2 = score_meme(image_url2, caption2)
 
-    # Display Meme 1
     st.write("#### Meme 1")
     st.image(image_url1, caption=f"Caption: {caption1}")
     st.write(f"AI Score: {score1:.2f}")
 
-    # Display Meme 2
     st.write("#### Meme 2")
     st.image(image_url2, caption=f"Caption: {caption2}")
     st.write(f"AI Score: {score2:.2f}")
 
-    # Determine the winner
     if score1 > score2:
         st.write("🎉 **Meme 1 Wins!**")
     elif score2 > score1:
         st.write("🎉 **Meme 2 Wins!**")
     else:
         st.write("🤝 **It's a tie!**")
+