src/data.py

"""Data handling functions for the Safe Scan competition dashboard."""
import streamlit as st
import wandb
import pandas as pd
import time
import datetime
from typing import Dict, Any, Tuple, Optional
from utils import fetch_competition_summary, fetch_models_evaluation

def update_leader_info(
    leader_info: Dict[str, Dict[str, Any]], 
    competition: str, 
    best_model: Dict[str, Any]
) -> Dict[str, Any]:
    """Update the leader information for a competition.
    
    Args:
        leader_info: Dictionary containing leader information for all competitions
        competition: Name of the competition
        best_model: Dictionary containing information about the best model
        
    Returns:
        Updated leader information for the specified competition
    """
    if leader_info.get(competition) is None:
        leader_info[competition] = {
            "Miner hotkey": best_model["Miner hotkey"],
            "Date": time.strftime("%Y-%m-%d"),
        }
    else:
        if leader_info[competition]["Miner hotkey"] != best_model["Miner hotkey"]:
            leader_info[competition]["Miner hotkey"] = best_model["Miner hotkey"]
            leader_info[competition]["Date"] = time.strftime("%Y-%m-%d")
    return leader_info[competition]

@st.cache_data()
def fetch_competition_summaries(
    _wandb_api: wandb.Api,
    entity: str,
    competitions: list
) -> Dict[str, pd.DataFrame]:
    """Fetch and cache competition summaries from Weights & Biases.
    
    Args:
        _wandb_api: Weights & Biases API instance
        entity: W&B entity name
        competitions: List of competition names
        
    Returns:
        Dictionary of competition summaries
    """
    competition_summaries = {}
    for competition in competitions:
        competition_summaries[competition] = fetch_competition_summary(_wandb_api, entity, competition)
    return competition_summaries


def fetch_model_evaluations(
    _wandb_api: wandb.Api,
    entity: str,
    competitions: list
) -> Dict[str, pd.DataFrame]:
    """Fetch model evaluations from Weights & Biases without caching.
    
    Args:
        _wandb_api: Weights & Biases API instance
        entity: W&B entity name
        competitions: List of competition names
        
    Returns:
        Dictionary of model evaluations
    """
    model_evaluations = {}
    for competition in competitions:
        model_evaluations[competition] = fetch_models_evaluation(_wandb_api, entity, competition)
    return model_evaluations


def load_competition_data(
    _wandb_api: wandb.Api,
    entity: str,
    competitions: list,
    last_update_time: Optional[float] = None
) -> Tuple[Dict[str, pd.DataFrame], Dict[str, pd.DataFrame], float]:
    """Load competition data from Weights & Biases.
    
    Args:
        _wandb_api: Weights & Biases API instance
        entity: W&B entity name
        competitions: List of competition names
        last_update_time: Timestamp of last update
        
    Returns:
        Tuple containing:
        - Dictionary of competition summaries
        - Dictionary of model evaluations
        - Current timestamp
    """
    competition_summaries = fetch_competition_summaries(_wandb_api, entity, competitions)
    model_evaluations = fetch_model_evaluations(_wandb_api, entity, competitions)
    return competition_summaries, model_evaluations, time.time()


def process_competition_leaders(
    competitions: list,
    competition_summaries: Dict[str, pd.DataFrame],
    model_evaluations: Dict[str, pd.DataFrame],
    leader_info: Dict[str, Dict[str, Any]]
) -> Dict[str, Dict[str, Any]]:
    """Process and update competition leaders information.
    
    Args:
        competitions: List of competition names
        competition_summaries: Dictionary of competition summaries
        model_evaluations: Dictionary of model evaluations
        leader_info: Current leader information
        
    Returns:
        Updated leader information dictionary
    """
    for competition in competitions:
        if not competition_summaries[competition].empty:
            winning_hotkeys = competition_summaries[competition]["Winning Hotkey"].value_counts()
            
            if not winning_hotkeys.empty:
                best_hotkey = winning_hotkeys.idxmax()
                best_model_filtered = model_evaluations[competition][
                    model_evaluations[competition]["Miner hotkey"] == best_hotkey
                ]
                
                if not best_model_filtered.empty:
                    best_model = best_model_filtered.iloc[0]
                    leader_info[competition] = update_leader_info(leader_info, competition, best_model)
                else:
                    st.warning(f"No model found for the best hotkey: {best_hotkey} in competition {competition}.")
        else:
            leader_info[competition] = {
                "Miner hotkey": "N/A",
                "Date": "N/A",
            }
    
    return leader_info