scripts/extract_data.py

import os
import json
import glob
from huggingface_hub import snapshot_download
import pandas as pd

# Constants
REPO_ID = "NohTow/BrowseComp-Plus-results-Agent-ModernColBERT"
CACHE_DIR = "hf_cache"
OUTPUT_FILE = os.getenv("LEADERBOARD_DATA_PATH", "public/data/leaderboard.json")

def download_data():
    print(f"Downloading results from {REPO_ID}...")
    token = os.getenv("HF_TOKEN")
    snapshot_download(
        repo_id=REPO_ID,
        repo_type="dataset",
        local_dir=CACHE_DIR,
        allow_patterns="*.csv",
        token=token
    )

def parse_results():
    print("Parsing results from CSV...")
    csv_path = os.path.join(CACHE_DIR, "agent_results.csv")
    if not os.path.exists(csv_path):
        print(f"CSV not found at {csv_path}")
        return []
    
    try:
        # The CSV has some BOM or encoding that requires utf-8-sig
        df = pd.read_csv(csv_path, encoding="utf-8-sig")
        
        # Only keep rows where Accuracy is present
        df = df.dropna(subset=["Accuracy (%)"])
        
        # Convert numeric columns to float where possible, others to string
        numeric_cols = ["Accuracy (%)", "Recall (%)", "Search Calls", "Calibration Error (%)"]
        for col in numeric_cols:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce').fillna(0)

        # Fill all other NaN values with empty strings for valid JSON
        df = df.fillna("")

        # Convert entire dataframe to list of dicts (preserves headers as keys)
        results = df.to_dict(orient='records')
            
        print(f"Parsed {len(results)} models.")
        return results
    except Exception as e:
        print(f"Error parsing CSV: {e}")
        return []

def main():
    # Always clear cache to ensure fresh data
    import shutil
    if os.path.exists(CACHE_DIR):
        shutil.rmtree(CACHE_DIR)
    
    download_data()
    
    leaderboard = parse_results()
    
    os.makedirs(os.path.dirname(OUTPUT_FILE), exist_ok=True)
    with open(OUTPUT_FILE, 'w') as f:
        json.dump(leaderboard, f, indent=2)
    
    print(f"Extraction complete. Data saved to {OUTPUT_FILE}")

if __name__ == "__main__":
    main()