Add check for df_all.task

utils.py

@@ -136,7 +136,8 @@ def load_downloaded_runs(time, cols=KEYS):
         'date-based question answering': 'date_qa',
         'question-answering': 'qa',
     }
-    df_all.task = df_all.task.apply(lambda x: task_mapping.get(x, x))
+    if 'task' in df_all.columns:
+        df_all.task = df_all.task.apply(lambda x: task_mapping.get(x, x))
 
     # Runs which do not have a turn field are imputed to be turn zero (single turn)
     df_all.turn.fillna(0, inplace=True)
@@ -223,7 +224,7 @@ def download_runs(time, df_vali):
         save_path = f'data/wandb/{row.run_id}.parquet'
         # Create the directory if it does not exist
         os.makedirs(os.path.dirname(save_path), exist_ok=True)
-        
+
         if os.path.exists(save_path):
             pbar.set_description(f'>> Skipping {row.run_id!r} because file {save_path!r} already exists')
             continue