updates

app.py

@@ -29,13 +29,15 @@ try:
 
     # Check if the 'split' column exists; if not, add it
     if 'split' not in dataset.column_names:
-        match = re.search(r'_(\d+)\.', example["image_id"])
-        image_id = match.group(1).lstrip('0') 
-
-        split_values = [
-            "dev" if image_id in results else "train" 
-            for example in dataset
-        ]
+        split_values = []
+        for example in dataset:
+            match = re.search(r'_(\d+)\.', example["image_id"])
+            image_id = match.group(1).lstrip('0') 
+
+            if image_id in results:
+                split_values.append("dev")
+            else:
+                split_values.append("train")
 
         # Add 'split' column to the dataset
         dataset = dataset.add_column("split", split_values)