Update space

app.py

@@ -14,6 +14,7 @@ from email.mime.text import MIMEText
 from huggingface_hub import HfApi
 import shutil
 import tempfile
+import time
 
 from stark_qa import load_qa
 from stark_qa.evaluator import Evaluator
@@ -30,31 +31,32 @@ except Exception as e:
 
 
 def process_single_instance(args):
+    """Process a single instance with progress tracking"""
     idx, eval_csv, qa_dataset, evaluator, eval_metrics = args
-    query, query_id, answer_ids, meta_info = qa_dataset[idx]
-
     try:
-        pred_rank = eval_csv[eval_csv['query_id'] == query_id]['pred_rank'].item()
-    except IndexError:
-        raise IndexError(f'Error when processing query_id={query_id}, please make sure the predicted results exist for this query.')
-    except Exception as e:
-        raise RuntimeError(f'Unexpected error occurred while fetching prediction rank for query_id={query_id}: {e}')
-
-    if isinstance(pred_rank, str):
+        query, query_id, answer_ids, meta_info = qa_dataset[idx]
+        
+        # Print progress for debugging
+        print(f"Processing query_id: {query_id}")
+        
         try:
+            pred_rank = eval_csv[eval_csv['query_id'] == query_id]['pred_rank'].item()
+        except Exception as e:
+            print(f"Error getting pred_rank for query_id {query_id}: {str(e)}")
+            raise
+            
+        if isinstance(pred_rank, str):
             pred_rank = eval(pred_rank)
-        except SyntaxError as e:
-            raise ValueError(f'Failed to parse pred_rank as a list for query_id={query_id}: {e}')
-    
-    if not isinstance(pred_rank, list):
-        raise TypeError(f'Error when processing query_id={query_id}, expected pred_rank to be a list but got {type(pred_rank)}.')
-
-    pred_dict = {pred_rank[i]: -i for i in range(min(100, len(pred_rank)))}
-    answer_ids = torch.LongTensor(answer_ids)
-    result = evaluator.evaluate(pred_dict, answer_ids, metrics=eval_metrics)
-
-    result["idx"], result["query_id"] = idx, query_id
-    return result
+        
+        pred_dict = {pred_rank[i]: -i for i in range(min(100, len(pred_rank)))}
+        answer_ids = torch.LongTensor(answer_ids)
+        result = evaluator.evaluate(pred_dict, answer_ids, metrics=eval_metrics)
+        
+        result["idx"], result["query_id"] = idx, query_id
+        return result
+    except Exception as e:
+        print(f"Error in process_single_instance for idx {idx}: {str(e)}")
+        raise
 
 
 def compute_metrics(csv_path: str, dataset: str, split: str, num_workers: int = 4):
@@ -63,6 +65,7 @@ def compute_metrics(csv_path: str, dataset: str, split: str, num_workers: int =
         'mag': [i for i in range(1172724, 1872968)],
         'prime': [i for i in range(129375)]
     }
+    start_time = time.time()
     try:
         eval_csv = pd.read_csv(csv_path)
         if 'query_id' not in eval_csv.columns:
@@ -77,11 +80,14 @@ def compute_metrics(csv_path: str, dataset: str, split: str, num_workers: int =
         if split not in ['test', 'test-0.1', 'human_generated_eval']:
             raise ValueError(f"Invalid split '{split}', expected one of ['test', 'test-0.1', 'human_generated_eval'].")
 
+        print("Initializing evaluator...")
         evaluator = Evaluator(candidate_ids_dict[dataset])
         eval_metrics = ['hit@1', 'hit@5', 'recall@20', 'mrr']
+        print("Loading QA dataset...")
         qa_dataset = load_qa(dataset, human_generated_eval=split == 'human_generated_eval')
         split_idx = qa_dataset.get_idx_split()
         all_indices = split_idx[split].tolist()
+        print(f"Dataset loaded, processing {len(all_indices)} instances")
 
         results_list = []
         # query_ids = []
@@ -102,28 +108,50 @@ def compute_metrics(csv_path: str, dataset: str, split: str, num_workers: int =
         #     metric: np.mean(eval_csv[eval_csv['query_id'].isin(query_ids)][metric]) for metric in eval_metrics
         # }
         # return final_result
-        batch_size = 100
-        for i in range(0, len(all_indices), batch_size):
-            max_ind = min(i+batch_size, len(all_indices))
-            batch_indices = all_indices[i:max_ind]
+        batch_size = 50  # Smaller batch size for more frequent updates
+        total_batches = (len(all_indices) + batch_size - 1) // batch_size
+        
+        for batch_num in range(total_batches):
+            batch_start = batch_num * batch_size
+            batch_end = min((batch_num + 1) * batch_size, len(all_indices))
+            batch_indices = all_indices[batch_start:batch_end]
+            
+            print(f"\nProcessing batch {batch_num + 1}/{total_batches}")
+            print(f"Batch size: {len(batch_indices)}")
+            
             args = [(idx, eval_csv, qa_dataset, evaluator, eval_metrics) 
                    for idx in batch_indices]
             
             with ProcessPoolExecutor(max_workers=num_workers) as executor:
                 futures = [executor.submit(process_single_instance, arg) 
                           for arg in args]
-                for future in as_completed(futures):
-                    results_list.append(future.result())
-
-        # Compute final metrics
+                for future in tqdm(as_completed(futures), 
+                                 total=len(futures),
+                                 desc=f"Batch {batch_num + 1}"):
+                    try:
+                        result = future.result()
+                        results_list.append(result)
+                    except Exception as e:
+                        print(f"Error processing result: {str(e)}")
+                        raise
+        
+        print("\nComputing final metrics...")
         results_df = pd.DataFrame(results_list)
         final_results = {
             metric: results_df[metric].mean() 
             for metric in eval_metrics
         }
         
+        elapsed_time = time.time() - start_time
+        print(f"\nMetrics computation completed in {elapsed_time:.2f} seconds")
         return final_results
 
+    except Exception as error:
+        elapsed_time = time.time() - start_time
+        error_msg = f"Error in compute_metrics ({elapsed_time:.2f}s): {str(error)}"
+        print(error_msg)
+        return error_msg
+
     except pd.errors.EmptyDataError:
         return "Error: The CSV file is empty or could not be read. Please check the file and try again."
     except FileNotFoundError: