Update space

app.py

@@ -15,6 +15,9 @@ from huggingface_hub import HfApi
 import shutil
 import tempfile
 import time
+from concurrent.futures import ThreadPoolExecutor
+from queue import Queue
+import threading
 
 from stark_qa import load_qa
 from stark_qa.evaluator import Evaluator
@@ -80,16 +83,16 @@ 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...")
+        # print("Initializing evaluator...")
         evaluator = Evaluator(candidate_ids_dict[dataset])
         eval_metrics = ['hit@1', 'hit@5', 'recall@20', 'mrr']
-        print("Loading QA dataset...")
+        # 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 = []
+        # results_list = []
         # query_ids = []
 
         # # Prepare args for each worker
@@ -108,34 +111,56 @@ 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 = 50  # Smaller batch size for more frequent updates
+
+        batch_size = 100
+        results_list = []
+        progress_queue = Queue()
+        
+        def process_batch(batch_indices):
+            batch_results = []
+            with ThreadPoolExecutor(max_workers=num_workers) as executor:
+                futures = [
+                    executor.submit(process_single_instance, 
+                                 (idx, eval_csv, qa_dataset, evaluator, eval_metrics))
+                    for idx in batch_indices
+                ]
+                for future in futures:
+                    result = future.result()
+                    if result is not None:
+                        batch_results.append(result)
+                    progress_queue.put(1)
+            return batch_results
+
+        # Process batches
         total_batches = (len(all_indices) + batch_size - 1) // batch_size
+        remaining_indices = len(all_indices)
         
-        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]
+        def update_progress():
+            with tqdm(total=len(all_indices), desc="Processing instances") as pbar:
+                completed = 0
+                while completed < len(all_indices):
+                    progress_queue.get()
+                    completed += 1
+                    pbar.update(1)
+
+        # Start progress monitoring thread
+        progress_thread = threading.Thread(target=update_progress)
+        progress_thread.start()
+
+        # Process batches
+        for i in range(0, len(all_indices), batch_size):
+            batch_indices = all_indices[i:min(i + batch_size, len(all_indices))]
+            batch_results = process_batch(batch_indices)
+            results_list.extend(batch_results)
+            remaining_indices -= len(batch_indices)
+            print(f"\rBatch {i//batch_size + 1}/{total_batches} completed. Remaining: {remaining_indices}")
+
+        progress_thread.join()
+
+        # Compute final metrics
+        if not results_list:
+            raise ValueError("No valid results were produced")
             
-            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 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() 
@@ -151,13 +176,7 @@ def compute_metrics(csv_path: str, dataset: str, split: str, num_workers: int =
         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:
-        return f"Error: The file {csv_path} could not be found. Please check the file path and try again."
-    except Exception as error:
-        return f"{error}"
+        
 
 
 # Data dictionaries for leaderboard
@@ -666,6 +685,7 @@ def process_submission(
     code_repo, csv_file, model_description, hardware, paper_link, model_type
 ):
     """Process submission with progress updates"""
+    temp_files = []
     try:
         # 1. Initial validation
         yield "Validating submission details..."
@@ -680,6 +700,7 @@ def process_submission(
         else:
             try:
                 temp_fd, temp_csv_path = tempfile.mkstemp(suffix='.csv')
+                temp_files.append(temp_csv_path)
                 os.close(temp_fd)
                 shutil.copy2(csv_file.name, temp_csv_path)
             except Exception as e:
@@ -700,7 +721,7 @@ def process_submission(
             csv_path=temp_csv_path,
             dataset=dataset.lower(),
             split=split,
-            num_workers=2  # Reduced from 4 to 2
+            num_workers=4
         )
         
         if isinstance(results, str):
@@ -762,11 +783,15 @@ def process_submission(
         """
 
     except Exception as e:
-        return f"Error: {str(e)}"
+        total_time = time.time() - start_time
+        return f"Error ({total_time:.1f}s): {str(e)}"
     finally:
-        # Cleanup
-        if temp_csv_path and os.path.exists(temp_csv_path):
-            os.unlink(temp_csv_path)
+        for temp_file in temp_files:
+            try:
+                if os.path.exists(temp_file):
+                    os.unlink(temp_file)
+            except Exception as e:
+                print(f"Warning: Failed to delete {temp_file}: {str(e)}")
     
 def filter_by_model_type(df, selected_types):
     """