Alignment-Lab-AI/orcaleaderboard

1import pathlib2import pandas as pd3from datasets import Dataset4from src.display.formatting import has_no_nan_values, make_clickable_model5from src.display.utils import AutoEvalColumn, EvalQueueColumn, baseline_row6from src.leaderboard.filter_models import filter_models_flags7from src.display.utils import load_json_data8 9 10def _process_model_data(entry, model_name_key="model", revision_key="revision"):11    """Enrich model data with clickable links and revisions."""12    entry[EvalQueueColumn.model.name] = make_clickable_model(entry.get(model_name_key, ""))13    entry[EvalQueueColumn.revision.name] = entry.get(revision_key, "main")14    return entry15 16 17def get_evaluation_queue_df(save_path, cols):18    """Generate dataframes for pending, running, and finished evaluation entries."""19    save_path = pathlib.Path(save_path)20    all_evals = []21 22    for path in save_path.rglob("*.json"):23        data = load_json_data(path)24        if data:25            all_evals.append(_process_model_data(data))26 27    # Organizing data by status28    status_map = {29        "PENDING": ["PENDING", "RERUN"],30        "RUNNING": ["RUNNING"],31        "FINISHED": ["FINISHED", "PENDING_NEW_EVAL"],32    }33    status_dfs = {status: [] for status in status_map}34    for eval_data in all_evals:35        for status, extra_statuses in status_map.items():36            if eval_data["status"] in extra_statuses:37                status_dfs[status].append(eval_data)38 39    return tuple(pd.DataFrame(status_dfs[status], columns=cols) for status in ["FINISHED", "RUNNING", "PENDING"])40 41 42def get_leaderboard_df(leaderboard_dataset: Dataset, cols: list, benchmark_cols: list):43    """Retrieve and process leaderboard data."""44    all_data_json = leaderboard_dataset.to_dict()45    num_items = leaderboard_dataset.num_rows46    all_data_json_list = [{k: all_data_json[k][ix] for k in all_data_json.keys()} for ix in range(num_items)]47    filter_models_flags(all_data_json_list)48 49    df = pd.DataFrame.from_records(all_data_json_list)50    df = df.sort_values(by=[AutoEvalColumn.average.name], ascending=False)51    df = df[cols].round(decimals=2)52    df = df[has_no_nan_values(df, benchmark_cols)]53    return df54