updates

app.py

@@ -25,14 +25,15 @@ try:
     dataset = load_dataset(dataset_name, split="train")
     print("Loaded existing dataset:", dataset)
     print("Dataset features:", dataset.features)  # Check if 'split' is part of features
-    
+
+
     # Check if the 'split' column exists; if not, add it
     if 'split' not in dataset.column_names:
-        # Define the 'split' values based on `image_id`
-        print(results)
-        print(example["image_id"] in results)
+        match = re.search(r'_(\d+)\.', example["image_id"])
+        image_id = match.group(1).lstrip('0') 
+
         split_values = [
-            "dev" if example["image_id"] in results else "train" 
+            "dev" if image_id in results else "train" 
             for example in dataset
         ]