Update utils using prompting-api

utils.py

@@ -26,8 +26,6 @@ USERNAME = 'login19861986'
 
 # Initialize wandb with anonymous login
 wandb.login(anonymous='must')
-
-# Your existing code
 api = wandb.Api(timeout=600)
 
 IDENTITIES = {
@@ -138,16 +136,12 @@ 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))
 
-    # Fill the missing values in the turn column with 0
-    df_all['turn'] = df_all['turn'].fillna(0)
+    # Runs which do not have a turn field are imputed to be turn zero (single turn)
+    df_all.turn.fillna(0, inplace=True)
 
     df_all.sort_values(by=['_timestamp'], inplace=True)
-    # Check if df_all has the task columns
-    if 'task' in df_all.columns:
-        df_all.task = df_all.task.apply(lambda x: task_mapping.get(x, x))
-    else:
-        df_all['task'] = "Task was not found"
 
     return df_all
 
@@ -229,13 +223,13 @@ 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
 
         try:
-            pbar.set_description(f'* Downloading run {row.run_id!r}')
+            pbar.set_description(f'* Downloading run {row.run_id!r}', flush=True)
             run = api.run(row.run_path)
 
             # By default we just download a subset of events (500 most recent)
@@ -418,4 +412,4 @@ def load_state_vars(username=USERNAME, percentile=0.95):
 
 if __name__ == '__main__':
 
-    pass
+    pass