src/load_from_hub.py

import json
import math

import pandas as pd

from src.assets.symbols import UP_ARROW, DOWN_ARROW

from src.tasks import TASKS, TASK_CODES, TASK_TO_METRIC


def load_from_hub(fs, repo_path):
    files = fs.glob(f"{repo_path}/**/*.json")

    set_organization_models = {}
    tasks = {}

    for file in files:
        organization, model, task = file.split("/")[-3:]

        organization_model = f"{organization}/{model}"
        task_code = task.replace(".json", "")

        if task_code not in TASK_CODES:
            continue

        set_organization_models[organization_model] = 1
        tasks[task_code] = 1

    table = pd.DataFrame(
        index=list(set_organization_models.keys()),
        columns=["Organization", "Model"] + list(tasks.keys()),
        data=None,
    )

    for file in files:
        organization, model, task = file.split("/")[-3:]

        organization_model = f"{organization}/{model}"
        task_code = task.replace(".json", "")

        if task_code not in TASK_CODES:
            continue

        data = json.loads(fs.open(file, "r").read())

        result = round(data["results"][task_code][TASK_TO_METRIC[task_code]], 4)

        table.loc[organization_model, task_code] = result
        table.loc[organization_model, "Organization"] = organization
        table.loc[organization_model, "Model"] = model

    table.rename(columns={
        task.code: f"{task.name} {UP_ARROW if task.higher_is_better else DOWN_ARROW}"
        for task in TASKS}, inplace=True)

    return table