update

app.py

@@ -4,6 +4,7 @@ import threading
 import random
 from datasets import load_dataset, Dataset, Features, Value, concatenate_datasets
 from huggingface_hub import login
+
 # Authenticate with Hugging Face
 token = os.getenv("HUGGINGFACE_TOKEN")
 if token:
@@ -29,6 +30,12 @@ except Exception as e:
     dataset["train"].push_to_hub(f"{dataset_name}", split="train")
     dataset["val"].push_to_hub(f"{dataset_name}", split="val")
 
+# Verify that both splits are in the dataset
+if "train" not in dataset:
+    dataset["train"] = Dataset.from_dict({'image_id': [], 'caption': []}, features=features)
+if "val" not in dataset:
+    dataset["val"] = Dataset.from_dict({'image_id': [], 'caption': []}, features=features)
+
 image_folder = "test"
 image_files = [f for f in os.listdir(image_folder) if f.endswith(('.png', '.jpg', '.jpeg'))]
 lock = threading.Lock()
@@ -117,4 +124,4 @@ with gr.Blocks() as demo:
     # Load initial image
     demo.load(fn=initialize_interface, inputs=session_data, outputs=[image, caption])
 
-demo.launch(share=True)
+demo.launch(share=True)