refactor: streamline dataset and model handling with helper classes

app.py

@@ -1,18 +1,14 @@
 import gradio as gr
-import pandas as pd
-import random
+from typing import List, Tuple
 import plotly.express as px
 from huggingface_hub import snapshot_download
 import os
+import pdb
 import logging
+import pandas as pd
 
-from config import (
-    SETUPS,
-    LOCAL_RESULTS_DIR,
-    CITATION_BUTTON_TEXT,
-    CITATION_BUTTON_LABEL,
-)
-from parsing import read_all_configs, get_common_langs
+from config import LOCAL_RESULTS_DIR, CITATION_BUTTON_TEXT, DatasetHelper, ModelHelper
+from parsing import read_all_configs
 
 # Set up logging
 logging.basicConfig(
@@ -57,27 +53,30 @@ We are currently hiding the results of {', '.join(model_markups)} because they d
 """
 
 
-def build_components(show_common_langs):
-    aggregated_df, lang_df, barplot_fig, models_with_nan = _populate_components(
-        show_common_langs
+def build_components(show_common_langs, selected_datasets: List[str]):
+    aggregated_df, lang_dfs, barplot_figs, models_with_nan = _populate_components(
+        show_common_langs, selected_datasets
     )
     models_with_nan_md = _build_models_with_nan_md(models_with_nan)
 
     return (
         gr.DataFrame(format_dataframe(aggregated_df)),
-        gr.DataFrame(format_dataframe(lang_df, times_100=True)),
-        gr.Plot(barplot_fig),
+        gr.DataFrame(format_dataframe(lang_dfs[0], times_100=True)),
+        gr.DataFrame(format_dataframe(lang_dfs[1], times_100=True)),
+        gr.Plot(barplot_figs[0]),
+        gr.Plot(barplot_figs[1]),
         gr.Markdown(models_with_nan_md, visible=len(models_with_nan) > 0),
     )
 
 
-def _populate_components(show_common_langs):
-    fm = SETUPS[0]
-    setup = fm["majority_group"] + "_" + fm["minority_group"]
-    results = read_all_configs(setup)
+def _populate_components(
+    show_common_langs: bool, selected_datasets: List[str], contrast_type: str = "F-M"
+) -> Tuple[pd.DataFrame, List[pd.DataFrame], List[px.bar], List[str]]:
+
+    results = read_all_configs(contrast_type)
 
     if show_common_langs:
-        common_langs = get_common_langs()
+        common_langs = model_h.get_common_langs()
         logger.info(f"Common langs: {common_langs}")
         results = results[results["Language"].isin(common_langs)]
 
@@ -96,64 +95,116 @@ def _populate_components(show_common_langs):
     logger.info(f"Models with NaN values: {models_with_nan}")
     results = results[~results["Model"].isin(models_with_nan)]
 
-    aggregated_df = (
-        results.pivot_table(
-            index="Model", values="Gap", aggfunc=lambda x: 100 * x.abs().sum()
-        )
-        .reset_index()
-        .sort_values("Gap")
-    )
-    best_model = aggregated_df.iloc[0]["Model"]
-    top_3_models = aggregated_df["Model"].head(3).tolist()
-    # main_df = gr.DataFrame(format_dataframe(model_results))
-
-    lang_df = results.pivot_table(
-        index="Model",
-        values="Gap",
-        columns="Language",
-    ).reset_index()
-    # lang_df = gr.DataFrame(format_dataframe(lang_results, times_100=True))
-
-    # gr.Plot(fig1)
-    results["Gap"] = results["Gap"] * 100
-    barplot_fig = px.bar(
-        results.loc[results["Model"].isin(top_3_models)],
-        x="Language",
-        y="Gap",
-        color="Model",
-        title="Gaps by Language and Model (top 3, sorted by the best model)",
-        labels={
-            "Gap": "Sum of Absolute Gaps (%)",
-            "Language": "Language",
-            "Model": "Model",
-        },
-        barmode="group",
-    )
-    lang_order = (
-        lang_df.set_index("Model").loc[best_model].sort_values(ascending=False).index
-    )
-    logger.info(f"Lang order: {lang_order}")
+    type_dfs = list()
+    lang_dfs = list()
+    barplot_figs = list()
+    for type, type_df in results.groupby("Type"):
 
-    barplot_fig.update_layout(
-        xaxis={"categoryorder": "array", "categoryarray": lang_order}
-    )
+        # Aggregate main
+        aggregated_df = type_df.pivot_table(
+            index="Model",
+            values="Gap",
+            aggfunc=lambda x: 100 * x.abs().sum(),
+        )
+        aggregated_df = aggregated_df.rename(columns={"Gap": f"Gap ({type})"})
+        type_dfs.append(aggregated_df)
+
+        best_model = aggregated_df.index[0]
+        top_3_models = aggregated_df.index[:3].tolist()
+
+        # Aggregate by language
+        lang_df = type_df.pivot_table(
+            index="Model",
+            values="Gap",
+            columns="Language",
+        ).reset_index()
+        lang_dfs.append(lang_df)
+
+        # Create plot
+        type_df["Gap"] = type_df["Gap"] * 100
+        barplot_fig = px.bar(
+            type_df.loc[results["Model"].isin(top_3_models)],
+            x="Language",
+            y="Gap",
+            color="Model",
+            title=f"{type}: Gaps by Language and Model (top 3, sorted by the best model)",
+            labels={
+                "Gap": f"{contrast_type} Gap (%)",
+                "Language": "Language",
+                "Model": "Model",
+            },
+            barmode="group",
+        )
 
-    return aggregated_df, lang_df, barplot_fig, models_with_nan
+        lang_order = (
+            lang_df.set_index("Model")
+            .loc[best_model]
+            .sort_values(ascending=False)
+            .index
+        )
+        logger.info(f"Lang order: {lang_order}")
 
+        barplot_fig.update_layout(
+            xaxis={"categoryorder": "array", "categoryarray": lang_order}
+        )
+        barplot_figs.append(barplot_fig)
+
+    # pdb.set_trace()
+    aggregated_df = pd.concat(type_dfs, axis=1, join="inner")
+    aggregated_df["Avg"] = aggregated_df.mean(axis=1)
+    aggregated_df = aggregated_df.sort_values("Avg").reset_index()
+
+    # lang_df = results.pivot_table(
+    #     index="Model",
+    #     values="Gap",
+    #     columns="Language",
+    # ).reset_index()
+
+    # results["Gap"] = results["Gap"] * 100
+    # barplot_fig = px.bar(
+    #     results.loc[results["Model"].isin(top_3_models)],
+    #     x="Language",
+    #     y="Gap",
+    #     color="Model",
+    #     title="Gaps by Language and Model (top 3, sorted by the best model)",
+    #     labels={
+    #         "Gap": "Sum of Absolute Gaps (%)",
+    #         "Language": "Language",
+    #         "Model": "Model",
+    #     },
+    #     barmode="group",
+    # )
+    # lang_order = (
+    #     lang_df.set_index("Model").loc[best_model].sort_values(ascending=False).index
+    # )
+    # logger.info(f"Lang order: {lang_order}")
+
+    # barplot_fig.update_layout(
+    #     xaxis={"categoryorder": "array", "categoryarray": lang_order}
+    # )
+
+    return aggregated_df, lang_dfs, barplot_figs, models_with_nan
+
+
+dataset_h = DatasetHelper()
+model_h = ModelHelper()
 
 with gr.Blocks() as fm_interface:
-    aggregated_df, lang_df, barplot_fig, model_with_nan = _populate_components(
-        show_common_langs=False
+    aggregated_df, lang_dfs, barplot_figs, model_with_nan = _populate_components(
+        show_common_langs=False, selected_datasets=dataset_h.get_dataset_names()
     )
     model_with_nans_md = gr.Markdown(_build_models_with_nan_md(model_with_nan))
 
     gr.Markdown("### Sum of Absolute Gaps ⬇️")
     aggregated_df_comp = gr.DataFrame(format_dataframe(aggregated_df))
 
-    gr.Markdown("#### F-M gaps by language")
-    lang_df_comp = gr.DataFrame(format_dataframe(lang_df, times_100=True))
+    gr.Markdown("#### Read: gaps by language")
+    lang_df_comp_0 = gr.DataFrame(format_dataframe(lang_dfs[0], times_100=True))
+    barplot_fig_comp_0 = gr.Plot(barplot_figs[0])
 
-    barplot_fig_comp = gr.Plot(barplot_fig)
+    gr.Markdown("#### Spontaneous: gaps by language")
+    lang_df_comp_1 = gr.DataFrame(format_dataframe(lang_dfs[1], times_100=True))
+    barplot_fig_comp_1 = gr.Plot(barplot_figs[1])
 
 ###################
 # LIST MAIN TABS
@@ -179,6 +230,7 @@ banner = """
 # MAIN INTERFACE
 ###################
 with gr.Blocks() as demo:
+
     gr.HTML(banner)
 
     with gr.Row() as config_row:
@@ -186,31 +238,40 @@ with gr.Blocks() as demo:
             choices=["Show only common languages"],
             label="Main configuration",
         )
+
+        datasets_names = dataset_h.get_dataset_names()
         include_datasets = gr.CheckboxGroup(
-            choices=["Mozilla CV 17"],
+            choices=datasets_names,
             label="Include datasets",
-            value=["Mozilla CV 17"],
+            value=datasets_names,
             interactive=False,
         )
 
         show_common_langs.input(
             build_components,
-            inputs=[show_common_langs],
+            inputs=[show_common_langs, include_datasets],
             outputs=[
                 aggregated_df_comp,
-                lang_df_comp,
-                barplot_fig_comp,
+                lang_df_comp_0,
+                lang_df_comp_1,
+                barplot_fig_comp_0,
+                barplot_fig_comp_1,
                 model_with_nans_md,
             ],
         )
 
     gr.TabbedInterface(tabs, titles)
 
-    gr.Textbox(
-        value=CITATION_BUTTON_TEXT,
-        label=CITATION_BUTTON_LABEL,
-        max_lines=6,
-        show_copy_button=True,
+    gr.Markdown(
+        """
+    ### Citation
+    If you find these results useful, please cite the following paper:
+    """
+    )
+
+    gr.Markdown(
+        f"""```
+{CITATION_BUTTON_TEXT}"""
     )
 
 if __name__ == "__main__":

config.py

@@ -4,78 +4,43 @@ to use for a particular datasetm or which language a model should be
 evaluated on.
 """
 
+from fair_asr_code.config import ALL_DATASET_CONFIGS, MODEL2LANG_SUPPORT
+
 LOCAL_RESULTS_DIR = "fair-asr-results"
 SETUPS = [{"majority_group": "male_masculine", "minority_group": "female_feminine"}]
 
 
-class CVInfo:
-    dataset_id: str = "cv_17"
-    full_name: str = "Mozilla Common Voice v17"
-
-    # fmt: off
-    langs = [
-        "de", "en", "nl",  # Germanic
-        "ru", "sr", "cs", "sk",  # Slavic
-        "it", "fr", "es", "ca", "pt", "ro",  # Romance
-        "sw",  # Bantu
-        "yo",  # Niger-Congo
-        "ja",  # Japonic
-        "hu", "fi",  # Uralic
-        "ar"  # Semitic
-    ]
-    # fmt: on
-
-
-dataset2info = {"cv_17": CVInfo}
-
+class DatasetHelper:
+    def __init__(self):
+        self.dataset_configs = ALL_DATASET_CONFIGS
 
-class WhisperInfo:
-    # fmt: off
-    langs = [
-        "de", "en", "nl",  # Germanic
-        "ru", "sr", "cs", "sk",  # Slavic
-        "it", "fr", "es", "ca", "pt", "ro",  # Romance
-        "sw",  # Bantu
-        "yo",  # Niger-Congo
-        "ja",  # Japonic
-        "hu", "fi",  # Uralic
-        "ar"  # Semitic
-    ]
-    # fmt: on
+    def get_dataset_names(self):
+        return [config.name for config in self.dataset_configs]
 
+    def get_dataset_ids(self):
+        return [config.dataset_id for config in self.dataset_configs]
 
-class SeamlessInfo:
-    # fmt: off
-    langs = [
-        "de", "en", "nl",  # Germanic
-        "ru", "sr", "cs", "sk",  # Slavic
-        "it", "fr", "es", "ca", "pt", "ro",  # Romance
-        "sw",  # Bantu
-        "yo",  # Niger-Congo
-        "ja",  # Japonic
-        "hu", "fi",  # Uralic
-        "ar"  # Semitic
-    ]
-    # fmt: on
+    @property
+    def sanitized_dataset_ids(self):
+        return [config.sanitized_id() for config in self.dataset_configs]
 
 
-class CanaryInfo:
-    # fmt: off
-    langs = [
-        "en", "es", "de", "fr",
-    ]
-    # fmt: on
+class ModelHelper:
+    def __init__(self):
+        self.models = list(MODEL2LANG_SUPPORT.keys())
 
+    @property
+    def sanitized_model_ids(self):
+        return [model.replace("/", "--") for model in self.models]
 
-model2info = {
-    "openai--whisper-large-v3": WhisperInfo,
-    "openai--whisper-large-v3-turbo": WhisperInfo,
-    "facebook--seamless-m4t-v2-large": SeamlessInfo,
-    "nvidia--canary-1b": CanaryInfo,
-}
+    def get_common_langs(self):
+        common_langs = set(MODEL2LANG_SUPPORT[self.models[0]])
+        for model in self.models:
+            common_langs = common_langs.intersection(set(MODEL2LANG_SUPPORT[model]))
+        return list(common_langs)
 
 
-CITATION_BUTTON_LABEL = "Please use this bibtex to cite these results"
+# CITATION_BUTTON_LABEL = "Please use this bibtex to cite these results"
 CITATION_BUTTON_TEXT = r"""@inproceedings{attanasio-etal-2024-twists,
     title = "Twists, Humps, and Pebbles: Multilingual Speech Recognition Models Exhibit Gender Performance Gaps",
     author = "Attanasio, Giuseppe  and

parsing.py

@@ -2,14 +2,14 @@ import pandas as pd
 from typing import List
 from os.path import join as opj
 import json
-from config import dataset2info, model2info, LOCAL_RESULTS_DIR
 import logging
+from config import DatasetHelper, ModelHelper, LOCAL_RESULTS_DIR
 
 logger = logging.getLogger(__name__)
 
 
 def load_language_results(
-    model_id: str, dataset_id: str, lang_ids: List[str], setup: str
+    model_id: str, dataset_id: str, lang_ids: List[str], contrast_string: str
 ):
     lang_gaps = dict()
     for lang in lang_ids:
@@ -20,7 +20,7 @@ def load_language_results(
                     LOCAL_RESULTS_DIR,
                     "evaluation",
                     dataset_id,
-                    f"results_{model_id}_{dataset_id}_devtest_{lang}_gender_{setup}.json",
+                    f"results_{model_id}_{dataset_id}_devtest_{lang}_gender_{contrast_string}.json",
                 )
             ) as fp:
                 data = json.load(fp)
@@ -34,26 +34,33 @@ def load_language_results(
     return lang_gaps
 
 
-def read_all_configs(setup: str):
-
-    all_datasets = dataset2info.keys()
-    print("Parsing results datasets:", all_datasets)
-    all_models = model2info.keys()
-    print("Parsing results models:", all_models)
+def read_all_configs(contrast_type: str):
+    dataset_h = DatasetHelper()
+    model_h = ModelHelper()
 
     rows = list()
-    for dataset_id in all_datasets:
-        for model_id in all_models:
+    for dataset_config in dataset_h.dataset_configs:
+        for model_id in model_h.sanitized_model_ids:
+
+            contrast_info = dataset_config.group_contrasts[contrast_type]
+            contrast_string = (
+                f"{contrast_info['majority_group']}_{contrast_info['minority_group']}"
+            )
+
             lang_gaps = load_language_results(
-                model_id, dataset_id, dataset2info[dataset_id].langs, setup
+                model_id,
+                dataset_config.sanitized_id(),
+                dataset_config.langs,
+                contrast_string,
             )
 
             rows.extend(
                 [
                     {
                         "Model": model_id,
-                        "Dataset": dataset_id,
+                        "Dataset": dataset_config.sanitized_id(),
                         "Language": lang,
+                        "Type": dataset_config.speaking_condition.capitalize(),
                         "Gap": lang_gaps[lang],
                     }
                     for lang in lang_gaps
@@ -61,16 +68,4 @@ def read_all_configs(setup: str):
             )
 
     results_df = pd.DataFrame(rows)
-    # results_df = results_df.drop(columns=["Dataset"])
-    # results_df = results_df.sort_values(by="Mean Gap", ascending=True)
-
     return results_df
-
-
-def get_common_langs():
-    """Return a list of langs that are support by all models"""
-    common_langs = set(model2info[list(model2info.keys())[0]].langs)
-    for model_id in model2info.keys():
-        common_langs = common_langs.intersection(model2info[model_id].langs)
-
-    return list(common_langs)

requirements.txt

@@ -1,3 +1,4 @@
 gradio
 plotly
-pandas
+pandas
+-e git+https://github.com/g8a9/fair-asr-code#egg=fair-asr-code