Update space

app.py

@@ -7,7 +7,7 @@ from datetime import datetime
 import json
 import torch
 from tqdm import tqdm
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, as_completed
 import smtplib
 from email.mime.multipart import MIMEMultipart
 from email.mime.text import MIMEText
@@ -41,50 +41,55 @@ except Exception as e:
 
 
 def process_single_instance(args):
-    """Process a single instance with improved prediction handling"""
+    """Process a single instance with improved validation and error handling"""
     idx, eval_csv, qa_dataset, evaluator, eval_metrics, max_candidate_id = args
     try:
+        # Get query data
         query, query_id, answer_ids, meta_info = qa_dataset[idx]
         
-        # Get predictions with better error handling
+        # Get predictions
         matching_preds = eval_csv[eval_csv['query_id'] == query_id]['pred_rank']
         if len(matching_preds) == 0:
             print(f"Warning: No prediction found for query_id {query_id}")
             return None
         elif len(matching_preds) > 1:
             print(f"Warning: Multiple predictions found for query_id {query_id}, using first one")
-            pred_rank = matching_preds.iloc[0]
-        else:
-            pred_rank = matching_preds.iloc[0]
+        
+        pred_rank = matching_preds.iloc[0]
         
         # Parse prediction
-        try:
-            if isinstance(pred_rank, str):
+        if isinstance(pred_rank, str):
+            try:
                 pred_rank = eval(pred_rank)
-            elif isinstance(pred_rank, list):
-                pass
-            else:
-                print(f"Warning: Unexpected pred_rank type for query_id {query_id}: {type(pred_rank)}")
+            except Exception as e:
+                print(f"Error parsing pred_rank for query_id {query_id}: {str(e)}")
                 return None
-        except Exception as e:
-            print(f"Error parsing pred_rank for query_id {query_id}: {str(e)}")
-            return None
-            
-        # Validate and filter predictions
+        
+        # Validate prediction format
         if not isinstance(pred_rank, list):
             print(f"Warning: pred_rank is not a list for query_id {query_id}")
             return None
             
-        # valid_ranks = [rank for rank in pred_rank if isinstance(rank, (int, np.integer)) and 0 <= rank < max_candidate_id]
-        # if len(valid_ranks) == 0:
-        #     print(f"Warning: No valid predictions for query_id {query_id}")
-        #     return None
+        # Validate and filter prediction values
+        valid_pred_rank = []
+        for rank in pred_rank[:100]:  # Only use top 100 predictions
+            if isinstance(rank, (int, np.integer)) and 0 <= rank < max_candidate_id:
+                valid_pred_rank.append(rank)
+            else:
+                print(f"Warning: Invalid prediction {rank} for query_id {query_id}")
+                
+        if not valid_pred_rank:
+            print(f"Warning: No valid predictions for query_id {query_id}")
+            return None
             
-        # Use only valid predictions
-        pred_dict = {pred_rank[i]: -i for i in range(min(100, len(pred_rank)))}
+        # Create prediction dictionary with valid predictions only
+        pred_dict = {rank: -i for i, rank in enumerate(valid_pred_rank)}
+        
+        # Convert answer_ids to tensor
         answer_ids = torch.LongTensor(answer_ids)
-        result = evaluator.evaluate(pred_dict, answer_ids, metrics=eval_metrics)
         
+        # Evaluate
+        result = evaluator.evaluate(pred_dict, answer_ids, metrics=eval_metrics)
         result["idx"], result["query_id"] = idx, query_id
         return result
         
@@ -93,116 +98,81 @@ def process_single_instance(args):
         return None
 
 def compute_metrics(csv_path: str, dataset: str, split: str, num_workers: int = 4):
-    """Compute metrics with improved prediction handling"""
+    """Compute metrics with improved thread safety and error handling"""
     start_time = time.time()
+    
+    # Dataset configuration
     candidate_ids_dict = {
         'amazon': [i for i in range(957192)],
         'mag': [i for i in range(1172724, 1872968)],
         'prime': [i for i in range(129375)]
     }
+    
+    candidate_size_dict = {
+        'amazon': 957192,
+        'mag': 700244,  # 1872968 - 1172724
+        'prime': 129375
+    }
+    
     try:
-        print(f"\nStarting compute_metrics for {dataset} {split}")
-        
-        # Load CSV and validate format
-        print("Loading and validating CSV file...")
-        eval_csv = pd.read_csv(csv_path)
-        if 'query_id' not in eval_csv.columns or 'pred_rank' not in eval_csv.columns:
-            raise ValueError("CSV must contain 'query_id' and 'pred_rank' columns")
-            
-        # Check for duplicate query_ids
-        duplicate_queries = eval_csv['query_id'].duplicated()
-        if duplicate_queries.any():
-            dup_count = duplicate_queries.sum()
-            print(f"Warning: Found {dup_count} duplicate query_ids in CSV")
-        
-        # Keep only necessary columns
-        eval_csv = eval_csv[['query_id', 'pred_rank']]
-        print(f"CSV loaded, shape: {eval_csv.shape}")
-        
-        # Get dataset-specific candidate size
-        candidate_size_dict = {
-            'amazon': 957192,
-            'mag': 700244,  # 1872968 - 1172724
-            'prime': 129375
-        }
-        
-        if dataset not in candidate_size_dict:
+        # Input validation
+        if dataset not in candidate_ids_dict:
             raise ValueError(f"Invalid dataset '{dataset}'")
+        if split not in ['test', 'test-0.1', 'human_generated_eval']:
+            raise ValueError(f"Invalid split '{split}'")
             
+        # Load and validate CSV
+        print(f"\nLoading data for {dataset} {split}")
+        eval_csv = pd.read_csv(csv_path)
+        required_columns = ['query_id', 'pred_rank']
+        if not all(col in eval_csv.columns for col in required_columns):
+            raise ValueError(f"CSV must contain columns: {required_columns}")
+        
+        eval_csv = eval_csv[required_columns]
         max_candidate_id = candidate_size_dict[dataset]
-        print(f"Dataset {dataset} has {max_candidate_id} candidates")
         
+        # Initialize components
         evaluator = Evaluator(candidate_ids_dict[dataset])
         eval_metrics = ['hit@1', 'hit@5', 'recall@20', 'mrr']
-        
         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"Processing {len(all_indices)} instances...")
+        print(f"Processing {len(all_indices)} instances with {num_workers} threads")
         
-        # Process in batches using ThreadPoolExecutor
-        batch_size = 100
+        # Process instances
         results_list = []
-        progress_queue = Queue()
-        valid_results_count = 0
+        valid_count = 0
         error_count = 0
         
-        def process_batch(batch_indices):
-            nonlocal valid_results_count, error_count
-            batch_results = []
-            with ThreadPoolExecutor(max_workers=num_workers) as executor:
-                futures = [
-                    executor.submit(process_single_instance, 
-                                 (idx, eval_csv, qa_dataset, evaluator, eval_metrics, max_candidate_id))
-                    for idx in batch_indices
-                ]
-                for future in futures:
+        with ThreadPoolExecutor(max_workers=num_workers) as executor:
+            futures = [
+                executor.submit(
+                    process_single_instance, 
+                    (idx, eval_csv, qa_dataset, evaluator, eval_metrics, max_candidate_id)
+                ) 
+                for idx in all_indices
+            ]
+            
+            with tqdm(total=len(futures), desc="Processing") as pbar:
+                for future in as_completed(futures):
                     try:
                         result = future.result()
                         if result is not None:
-                            batch_results.append(result)
-                            valid_results_count += 1
+                            results_list.append(result)
+                            valid_count += 1
                         else:
                             error_count += 1
                     except Exception as e:
-                        print(f"Error in batch processing: {str(e)}")
+                        print(f"Error in future: {str(e)}")
                         error_count += 1
-                    progress_queue.put(1)
-            return batch_results
-
-        # Process batches with progress tracking
-        total_batches = (len(all_indices) + batch_size - 1) // batch_size
-        remaining_indices = len(all_indices)
-        
-        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. "
-                  f"Valid: {valid_results_count}, Errors: {error_count}, Remaining: {remaining_indices}")
-
-        progress_thread.join()
-
+        
         # Compute final metrics
         if not results_list:
             raise ValueError("No valid results were produced")
             
-        print(f"\nProcessing complete. Valid results: {valid_results_count}, Errors: {error_count}")
+        print(f"\nProcessing complete. Valid: {valid_count}, Errors: {error_count}")
         
         results_df = pd.DataFrame(results_list)
         final_results = {
@@ -211,9 +181,9 @@ def compute_metrics(csv_path: str, dataset: str, split: str, num_workers: int =
         }
         
         elapsed_time = time.time() - start_time
-        print(f"\nMetrics computation completed in {elapsed_time:.2f} seconds")
+        print(f"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)}"