src/populate.py

import json
import os

import pandas as pd

from display.formatting import make_clickable_model
from display.utils_old import EvalQueueColumn


def get_leaderboard_df(results_path: str) -> pd.DataFrame:
    model_result_filepaths = []
    for root, _, files in os.walk(results_path):
        if len(files) == 0 or not all(f.endswith(".json") for f in files):
            continue
        for file in files:
            model_result_filepaths.append(os.path.join(root, file))

    eval_results = {"model": [], "buzz_accuracy": [], "win_rate_human": [], "win_rate_model": []}
    for model_result_filepath in model_result_filepaths:
        with open(model_result_filepath, "r") as fin:
            model_result = json.load(fin)
            model_id = model_result["model_id"]
            buzz_accuracy = model_result["buzz_accuracy"]
            win_rate_human = model_result["win_rate_human"]
            win_rate_model = model_result["win_rate_model"]
            eval_results["model"].append(model_id)
            eval_results["buzz_accuracy"].append(buzz_accuracy)
            eval_results["win_rate_human"].append(win_rate_human)
            eval_results["win_rate_model"].append(win_rate_model)
    return pd.DataFrame(eval_results)


def get_evaluation_queue_df(save_path: str, cols: list) -> list[pd.DataFrame]:
    # TODO: This function is stale, but might be a good reference point for new implementation
    """Creates the different dataframes for the evaluation queues requestes"""
    entries = [entry for entry in os.listdir(save_path) if not entry.startswith(".")]
    all_evals = []

    for entry in entries:
        if ".json" in entry:
            file_path = os.path.join(save_path, entry)
            with open(file_path) as fp:
                data = json.load(fp)

            data[EvalQueueColumn.model.name] = make_clickable_model(data["model"])
            data[EvalQueueColumn.revision.name] = data.get("revision", "main")

            all_evals.append(data)
        elif ".md" not in entry:
            # this is a folder
            sub_entries = [
                e for e in os.listdir(f"{save_path}/{entry}") if os.path.isfile(e) and not e.startswith(".")
            ]
            for sub_entry in sub_entries:
                file_path = os.path.join(save_path, entry, sub_entry)
                with open(file_path) as fp:
                    data = json.load(fp)

                data[EvalQueueColumn.model.name] = make_clickable_model(data["model"])
                data[EvalQueueColumn.revision.name] = data.get("revision", "main")
                all_evals.append(data)

    pending_list = [e for e in all_evals if e["status"] in ["PENDING", "RERUN"]]
    running_list = [e for e in all_evals if e["status"] == "RUNNING"]
    finished_list = [e for e in all_evals if e["status"].startswith("FINISHED") or e["status"] == "PENDING_NEW_EVAL"]
    df_pending = pd.DataFrame.from_records(pending_list, columns=cols)
    df_running = pd.DataFrame.from_records(running_list, columns=cols)
    df_finished = pd.DataFrame.from_records(finished_list, columns=cols)
    return df_finished[cols], df_running[cols], df_pending[cols]