Upload app.py

app.py

@@ -13,6 +13,13 @@ from transformers import (
 from datasets import load_dataset
 
 wikiart_dataset = load_dataset("huggan/wikiart", split="train", streaming=True)
+
+def get_item_streaming(dataset, idx):
+    for i, item in enumerate(dataset):
+        if i == idx:
+            return item
+    raise IndexError("Index out of range")
+
 device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
 
 blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
@@ -52,7 +59,7 @@ def get_results_with_images(embedding, index, top_k=2):
     results = []
     for idx in I[0]:
         try:
-            item = wikiart_dataset[int(idx)]
+            item = get_item_streaming(wikiart_dataset, int(idx))
             img = item["image"]
             title = item.get("title", "Untitled")
             artist = item.get("artist", "Unknown")