Update space

app.py

@@ -4,6 +4,98 @@ import os
 import re
 from datetime import datetime
 import json
+import pandas as pd
+import torch
+import numpy as np
+from tqdm import tqdm
+from concurrent.futures import ProcessPoolExecutor, as_completed
+
+from stark_qa import load_qa
+from stark_qa.evaluator import Evaluator
+
+
+def process_single_instance(args):
+    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):
+        try:
+            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
+
+
+def compute_metrics(csv_path: str, dataset: str, split: str, num_workers: int = 4):
+    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)]
+    }
+    try:
+        eval_csv = pd.read_csv(csv_path)
+        if 'query_id' not in eval_csv.columns:
+            raise ValueError('No `query_id` column found in the submitted csv.')
+        if 'pred_rank' not in eval_csv.columns:
+            raise ValueError('No `pred_rank` column found in the submitted csv.')
+
+        eval_csv = eval_csv[['query_id', 'pred_rank']]
+
+        if dataset not in candidate_ids_dict:
+            raise ValueError(f"Invalid dataset '{dataset}', expected one of {list(candidate_ids_dict.keys())}.")
+        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'].")
+
+        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()
+
+        results_list = []
+        query_ids = []
+
+        # Prepare args for each worker
+        args = [(idx, eval_csv, qa_dataset, evaluator, eval_metrics) for idx in all_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)):
+                result = future.result()  # This will raise an error if the worker encountered one
+                results_list.append(result)
+                query_ids.append(result['query_id'])
+
+        # Concatenate results and compute final metrics
+        eval_csv = pd.concat([eval_csv, pd.DataFrame(results_list)], ignore_index=True)
+        final_results = {
+            metric: np.mean(eval_csv[eval_csv['query_id'].isin(query_ids)][metric]) for metric in eval_metrics
+        }
+        return final_results
+
+    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}"
+
+
 
 
 # Sample data based on your table (you'll need to update this with the full dataset)
@@ -58,7 +150,6 @@ data_human_generated = {
 df_synthesized_full = pd.DataFrame(data_synthesized_full)
 df_synthesized_10 = pd.DataFrame(data_synthesized_10)
 df_human_generated = pd.DataFrame(data_human_generated)
-
 def validate_email(email_str):
     """Validate email format(s)"""
     emails = [e.strip() for e in email_str.split(';')]
@@ -247,7 +338,6 @@ def add_submission_form(demo):
             ],
             outputs=result
         )
-
 def format_dataframe(df, dataset):
     # Filter the dataframe for the selected dataset
     columns = ['Method'] + [col for col in df.columns if dataset in col]

requirements.txt

@@ -13,4 +13,5 @@ python-dateutil
 tqdm
 transformers
 tokenizers>=0.15.0
-sentencepiece
+sentencepiece
+stark_qa