Update src/populate.py

src/populate.py

@@ -4,41 +4,69 @@ import numpy as np
 import pandas as pd
 
 from src.display.formatting import has_no_nan_values, make_clickable_model
-from src.display.utils import AutoEvalColumn, EvalQueueColumn, AREA_DEFINITIONS, AREA_AVG_COLUMN_MAP
+from src.display.utils import AutoEvalColumn, EvalQueueColumn, AREA_DEFINITIONS, AREA_AVG_COLUMN_MAP, fields
 from src.leaderboard.read_evals import get_raw_eval_results
 from src.about import Tasks
 
 
-def get_leaderboard_df(results_path: str, requests_path: str, cols: list) -> pd.DataFrame:
-    """Creates a dataframe from all the individual experiment results"""
-    raw_data = get_raw_eval_results(results_path, requests_path)
-    all_data_json = [v.to_dict() for v in raw_data]
+def get_leaderboard_df(results_path: str = None, requests_path: str = None, cols: list = None, initial_df: pd.DataFrame = None) -> pd.DataFrame:
+    """Creates a dataframe from all the individual experiment results or uses a provided initial DataFrame."""
+    
+    if initial_df is not None:
+        df = initial_df.copy() # Use a cópia do DataFrame inicial
+        print("Usando DataFrame inicial fornecido.")
+    elif results_path and requests_path:
+        print(f"Lendo resultados de: {results_path}")
+        raw_data = get_raw_eval_results(results_path, requests_path)
+        all_data_json = [v.to_dict() for v in raw_data]
+        df = pd.DataFrame.from_records(all_data_json)
+    else:
+        print("Erro: Nenhum DataFrame inicial nem caminhos de resultados fornecidos.")
+        return pd.DataFrame() # Retorna DataFrame vazio se não houver dados
 
-    df = pd.DataFrame.from_records(all_data_json)
+    # Garantir que colunas de tasks existem antes de calcular médias
+    # (Opcional: Adicionar lógica para lidar com DFs que já têm médias calculadas)
+    tasks_in_df = [task.name for task in Tasks if task.name in df.columns]
+    print(f"Tasks encontrados no DataFrame: {tasks_in_df}")
 
     # Calcular médias por área
     for area_name, tasks_in_area in AREA_DEFINITIONS.items():
+        # Usar task.name que é a chave interna/coluna no df
         area_cols = [task.name for task in tasks_in_area if task.name in df.columns]
         avg_col_name = AREA_AVG_COLUMN_MAP[area_name]
-        if area_cols: # Calcula a média apenas se houver colunas da área no DataFrame
-            df[avg_col_name] = df[area_cols].mean(axis=1)
+        if area_cols:
+            # Lidar com possíveis NaNs nas colunas antes de calcular a média
+            df[avg_col_name] = df[area_cols].mean(axis=1, skipna=True)
+            print(f"Calculada média para {area_name} usando colunas: {area_cols}")
         else:
-            df[avg_col_name] = np.nan # Define como NaN se nenhuma coluna da área estiver presente
+            df[avg_col_name] = np.nan
+            print(f"Nenhuma coluna encontrada para {area_name}, definindo média como NaN.")
     
     # Calcular Média Geral (agora baseada nas médias das áreas)
-    avg_area_cols = list(AREA_AVG_COLUMN_MAP.values())
-    df[AutoEvalColumn.average.name] = df[avg_area_cols].mean(axis=1)
-    
-    df = df.sort_values(by=[AutoEvalColumn.average.name], ascending=False)
+    avg_area_cols = [col for col in AREA_AVG_COLUMN_MAP.values() if col in df.columns]
+    if avg_area_cols:
+        df[AutoEvalColumn.average.name] = df[avg_area_cols].mean(axis=1, skipna=True)
+        print(f"Calculada Média Geral usando colunas: {avg_area_cols}")
+    else: 
+        df[AutoEvalColumn.average.name] = np.nan
+        print("Nenhuma coluna de média de área encontrada, definindo Média Geral como NaN.")
 
+    # Ordenar pela Média Geral
+    if AutoEvalColumn.average.name in df.columns:
+      df = df.sort_values(by=[AutoEvalColumn.average.name], ascending=False)
+    
     # Selecionar e arredondar colunas
-    all_display_cols = [c.name for c in fields(AutoEvalColumn)] # Obter todas as colunas definidas
-    df = df[[col for col in all_display_cols if col in df.columns]] # Manter apenas colunas existentes no df
-    df = df.round(decimals=2)
+    # Usar as colunas definidas em AutoEvalColumn que existem no DataFrame
+    defined_cols = [c.name for c in fields(AutoEvalColumn)]
+    cols_to_keep = [col for col in defined_cols if col in df.columns]
+    
+    # Adicionar colunas que podem não estar em AutoEvalColumn mas estavam no DF original
+    # (Ex: colunas de tasks individuais que foram usadas para cálculo mas não definidas explicitamente para exibição inicial)
+    original_cols_needed_for_display = set(cols) if cols else set() # cols passados para a função
+    final_cols = sorted(list(set(cols_to_keep) | (original_cols_needed_for_display & set(df.columns))))
 
-    # Filtrar linhas com valores NaN nas colunas de benchmark originais (se necessário)
-    # benchmark_cols = [t.name for t in Tasks] # Descomentar se precisar do filtro original
-    # df = df[has_no_nan_values(df, benchmark_cols)]
+    df = df[final_cols]
+    df = df.round(decimals=2)
 
     return df