app.py

#!/usr/bin/env python3
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.repocard import metadata_load

import pandas as pd
import gradio as gr

METRICS_TO_NOT_DISPLAY = set(["ser"])
NO_LANGUAGE_MODELS = []

api = HfApi()
models = api.list_models(filter="robust-speech-event")
model_ids = [x.modelId for x in models]


def get_metadatas(model_ids):
    metadatas = {}
    for model_id in model_ids:
        readme_path = hf_hub_download(model_id, filename="README.md")
        metadatas[model_id] = metadata_load(readme_path)
    return metadatas


def get_model_results_and_language_map(metadatas):
    all_model_results = {}
    # model_id
    #  - dataset
    #     - metric
    model_language_map = {}
    # model_id: lang
    for model_id, metadata in metadatas.items():
        if "language" not in metadata:
            NO_LANGUAGE_MODELS.append(model_id)
            continue
        lang = metadata["language"]
        model_language_map[model_id] = lang if isinstance(lang, list) else [lang]
        if "model-index" not in metadata:
            all_model_results[model_id] = None
        else:
            result_dict = {}
            for result in metadata["model-index"][0]["results"]:
                dataset = result["dataset"]["type"]
                metrics = [x["type"] for x in result["metrics"]]
                values = [
                    x["value"] if "value" in x else None for x in result["metrics"]
                ]
                result_dict[dataset] = {k: v for k, v in zip(metrics, values)}

        all_model_results[model_id] = result_dict
    return all_model_results, model_language_map


def get_datasets_metrics_langs(all_model_results, model_language_map):
    # get all datasets
    all_datasets = set(
        sum([list(x.keys()) for x in all_model_results.values() if x is not None], [])
    )
    all_langs = set(sum(list(model_language_map.values()), []))

    # get all metrics
    all_metrics = []
    for metric_result in all_model_results.values():
        if metric_result is not None:
            all_metrics += sum([list(x.keys()) for x in metric_result.values()], [])

    all_metrics = set(all_metrics) - METRICS_TO_NOT_DISPLAY
    return all_datasets, all_langs, all_metrics


# get results table (one table for each dataset, metric)
def retrieve_dataframes(
    all_model_results, model_language_map, all_datasets, all_langs, all_metrics
):
    all_datasets_results = {}
    pandas_datasets = {}
    for dataset in all_datasets:
        all_datasets_results[dataset] = {}
        pandas_datasets[dataset] = {}
        for metric in all_metrics:
            all_datasets_results[dataset][metric] = {}
            pandas_datasets[dataset][metric] = {}
            for lang in all_langs:
                all_datasets_results[dataset][metric][lang] = {}
                results = {}
                for model_id, model_result in all_model_results.items():
                    is_relevant = (
                        lang in model_language_map[model_id]
                        and model_result is not None
                        and dataset in model_result
                        and metric in model_result[dataset]
                    )
                    if not is_relevant:
                        continue

                    result = model_result[dataset][metric]
                    if isinstance(result, str):
                        "".join(result.split("%"))
                        try:
                            result = float(result)
                        except:  # noqa: E722
                            result = None
                    elif isinstance(result, float) and result < 1.0:
                        # assuming that WER is given in 0.13 format
                        result = 100 * result
                    results[model_id] = round(result, 2) if result is not None else None

                results = dict(
                    sorted(results.items(), key=lambda item: (item[1] is None, item[1]))
                )
                all_datasets_results[dataset][metric][lang] = [
                    f"{v} : {k}" for k, v in results.items()
                ]

            data = all_datasets_results[dataset][metric]
            data_frame = pd.DataFrame.from_dict(data, orient="index")
            data_frame.fillna("", inplace=True)
            data_frame = data_frame.sort_index().transpose()
            pandas_datasets[dataset][metric] = data_frame
    return pandas_datasets


# 1. Retrieve metadatas
metadatas = get_metadatas(model_ids)

# 2. Parse to results
all_model_results, model_language_map = get_model_results_and_language_map(metadatas)

# 3. Get datasets and langs
all_datasets, all_langs, all_metrics = get_datasets_metrics_langs(
    all_model_results, model_language_map
)

# 4. Get dataframes
all_dataframes = retrieve_dataframes(
    all_model_results, model_language_map, all_datasets, all_langs, all_metrics
)


def select(dataset, metric):
    return all_dataframes[dataset][metric]


iface = gr.Interface(
    select,
    [
        gr.inputs.Dropdown(
            list(all_datasets),
            type="value",
            default="mozilla-foundation/common_voice_7_0",
            label="dataset",
        ),
        gr.inputs.Dropdown(
            list(all_metrics), type="value", default="wer", label="metric"
        ),
    ],
    gr.outputs.Dataframe(
        headers=None,
        max_rows=3,
        max_cols=10,
        overflow_row_behaviour="paginate",
        type="pandas",
        label=None,
    ),
    examples=[
        ["mozilla-foundation/common_voice_7_0", "wer"],
        ["mozilla-foundation/common_voice_7_0", "cer"],
    ],
    layout="vertical",
)

iface.launch()