src/details.py

import asyncio

import gradio as gr
import pandas as pd
from huggingface_hub import HfFileSystem

from src.constants import SUBTASKS, DETAILS_DATASET_ID, DETAILS_FILENAME
from src.hub import load_details_file


def update_subtasks_component(task):
    return gr.Radio(
        SUBTASKS.get(task),
        info="Evaluation subtasks to be displayed",
        value=None,
    )


def update_load_details_component(model_id_1, model_id_2, subtask):
    if (model_id_1 or model_id_2) and subtask:
        return gr.Button("Load Details", interactive=True)
    else:
        return gr.Button("Load Details", interactive=False)


async def load_details_dataframe(model_id, subtask):
    fs = HfFileSystem()
    if not model_id or not subtask:
        return
    model_name_sanitized = model_id.replace("/", "__")
    paths = fs.glob(
        f"{DETAILS_DATASET_ID}/**/{DETAILS_FILENAME}".format(
            model_name_sanitized=model_name_sanitized, subtask=subtask
        )
    )
    if not paths:
        return
    path = max(paths)
    data = await load_details_file(path)
    df = pd.json_normalize(data)
    # df = df.rename_axis("Parameters", axis="columns")
    df["model_name"] = model_id  # Keep model_name
    return df
    # return df.set_index(pd.Index([model_id])).reset_index()


async def load_details_dataframes(subtask, *model_ids):
    result = await asyncio.gather(*[load_details_dataframe(model_id, subtask) for model_id in model_ids])
    return result


def display_details(sample_idx, *dfs):
    rows = [df.iloc[sample_idx] for df in dfs if "model_name" in df.columns and sample_idx < len(df)]
    if not rows:
        return
    # Pop model_name and add it to the column name
    df = pd.concat([row.rename(row.pop("model_name")) for row in rows], axis="columns")
    # Wrap long strings to avoid overflow; e.g. URLs in "doc.Websites visited_NEV_2"
    df = df.apply(lambda x: x.str.wrap(140) if x.dtype == "object" else x)
    return (
        df.style
        .format(na_rep="")
        # .hide(axis="index")
        .to_html()
    )


def update_sample_idx_component(*dfs):
    maximum = max([len(df) - 1 for df in dfs])
    return gr.Number(
        label="Sample Index",
        info="Index of the sample to be displayed",
        value=0,
        minimum=0,
        maximum=maximum,
        visible=True,
    )


def clear_details():
    # model_id_1, model_id_2, details_dataframe_1, details_dataframe_2, details_task, subtask, load_details_btn, sample_idx
    return (
        None, None, None, None, None, None,
        gr.Button("Load Details", interactive=False),
        gr.Number(label="Sample Index", info="Index of the sample to be displayed", value=0, minimum=0,visible=False),
    )