three tabs

app.py

@@ -1,4 +1,5 @@
 import os
+from functools import partial
 
 import numpy as np
 import unicodedata
@@ -46,63 +47,124 @@ dataset = load_dataset(
     "distil-whisper/tedlium-long-form", split="validation", num_proc=os.cpu_count()
 )
 
-csv = pd.read_csv("assets/large-v2.csv")
+csv_v2 = pd.read_csv("assets/large-v2.csv")
 
-norm_target = csv["Norm Target"]
-norm_pred = csv["Norm Pred"]
+norm_target = csv_v2["Norm Target"]
+norm_pred_v2 = csv_v2["Norm Pred"]
 
 norm_target = [norm_target[i] for i in range(len(norm_target))]
-norm_pred = [norm_pred[i] for i in range(len(norm_pred))]
+norm_pred_v2 = [norm_pred_v2[i] for i in range(len(norm_pred_v2))]
+
+csv_v2 = pd.read_csv("assets/large-32-2.csv")
+
+norm_pred_32_2 = csv_v2["Norm Pred"]
+norm_pred_32_2 = [norm_pred_32_2[i] for i in range(len(norm_pred_32_2))]
 
 target_dtype = np.int16
 max_range = np.iinfo(target_dtype).max
 
 
-def get_visualisation(idx):
+def get_visualisation(idx, model="v2"):
     idx -= 1
     audio = dataset[idx]["audio"]
     array = (audio["array"] * max_range).astype(np.int16)
     sampling_rate = audio["sampling_rate"]
 
     text1 = norm_target[idx]
-    text2 = norm_pred[idx]
+    text2 = norm_pred_v2[idx] if model == "v2" else norm_pred_32_2[idx]
 
     wer_output = process_words(text1, text2, wer_default, wer_default)
-    wer_percentage = 100 * wer_output.wer
-    rel_insertions = wer_output.insertions / len(text1.split())
+    wer_percentage = round(100 * wer_output.wer, 2)
+    ier_percentage = round(100 *  wer_output.insertions / len(wer_output.references[0]), 2)
 
-    rel_length = len(text2.split()) / len(text1.split())
+    rel_length = round(len(text2.split()) / len(text1.split()), 2)
 
     diff = compare_string(text1, text2)
     full_text = style_text(diff)
 
-    return (sampling_rate, array), wer_percentage, rel_insertions, rel_length, full_text
+    return (sampling_rate, array), wer_percentage, ier_percentage, rel_length, full_text
+
+def get_side_by_side_visualisation(idx):
+    large_v2 = get_visualisation(idx, model="v2")
+    large_32_2 = get_visualisation(idx, model="32-2")
+    table = [large_v2[1:-1], large_32_2[1:-1]]
+    table[0] = ["large-v2", *table[0]]
+    table[1] = ["large-32-2", *table[1]]
+    return large_v2[0], table, large_v2[-1], large_32_2[-1]
 
 
 if __name__ == "__main__":
-    gr.Markdown(
-        "Analyse the transcriptions generated by the Whisper large-v2 model on the TEDLIUM dev set."
-    )
     with gr.Blocks() as demo:
-        slider = gr.Slider(
-            minimum=1, maximum=len(norm_target), step=1, label="Dataset sample"
-        )
-        btn = gr.Button("Analyse")
-        audio_out = gr.Audio(label="Audio input")
-        with gr.Row():
-            wer = gr.Number(label="WER")
-            relative_insertions = gr.Number(
-                label="Relative insertions (# insertions / target length)"
-            )
-            relative_length = gr.Number(
-                label="Relative length (reference length / target length)"
+        with gr.Tab("large-v2"):
+            gr.Markdown(
+                "Analyse the transcriptions generated by the Whisper large-v2 model on the TEDLIUM dev set."
             )
-        text_out = gr.Markdown(label="Text difference")
-
-        btn.click(
-            fn=get_visualisation,
-            inputs=slider,
-            outputs=[audio_out, wer, relative_insertions, relative_length, text_out],
-        )
 
+            slider = gr.Slider(
+                minimum=1, maximum=len(norm_target), step=1, label="Dataset sample"
+            )
+            btn = gr.Button("Analyse")
+            audio_out = gr.Audio(label="Audio input")
+            with gr.Row():
+                wer = gr.Number(label="Word Error Rate (WER)")
+                ier = gr.Number(
+                    label="Insertion Error Rate (IER)"
+                )
+                relative_length = gr.Number(
+                    label="Relative length (reference length / target length)"
+                )
+            text_out = gr.Markdown(label="Text difference")
+
+            btn.click(
+                fn=partial(get_visualisation, model="v2"),
+                inputs=slider,
+                outputs=[audio_out, wer, ier, relative_length, text_out],
+            )
+        with gr.Tab("large-32-2"):
+            gr.Markdown(
+                "Analyse the transcriptions generated by the Whisper large-32-2 model on the TEDLIUM dev set."
+            )
+            slider = gr.Slider(
+                minimum=1, maximum=len(norm_target), step=1, label="Dataset sample"
+            )
+            btn = gr.Button("Analyse")
+            audio_out = gr.Audio(label="Audio input")
+            with gr.Row():
+                wer = gr.Number(label="Word Error Rate (WER)")
+                ier = gr.Number(
+                    label="Insertion Error Rate (IER)"
+                )
+                relative_length = gr.Number(
+                    label="Relative length (reference length / target length)"
+                )
+            text_out = gr.Markdown(label="Text difference")
+
+            btn.click(
+                fn=partial(get_visualisation, model="32-2"),
+                inputs=slider,
+                outputs=[audio_out, wer, ier, relative_length, text_out],
+            )
+        with gr.Tab("side-by-side"):
+            gr.Markdown(
+                "Analyse the transcriptions generated by the Whisper large-32-2 model on the TEDLIUM dev set."
+            )
+            slider = gr.Slider(
+                minimum=1, maximum=len(norm_target), step=1, label="Dataset sample"
+            )
+            btn = gr.Button("Analyse")
+            audio_out = gr.Audio(label="Audio input")
+            with gr.Column():
+                table = gr.Dataframe(headers=["Model", "Word Error Rate (WER)", "Insertion Error Rate (IER)", "Rel length (ref length / tgt length)"], height=1000)
+                with gr.Row():
+                    gr.Markdown("large-v2 text diff")
+                    gr.Markdown("large-32-2 text diff")
+                with gr.Row():
+                    text_out_v2 = gr.Markdown(label="Text difference")
+                    text_out_32_2 = gr.Markdown(label="Text difference")
+
+            btn.click(
+                fn=get_side_by_side_visualisation,
+                inputs=slider,
+                outputs=[audio_out, table, text_out_v2, text_out_32_2],
+            )
     demo.launch()