Add ability to calculate WER/CER values per each row

app.py

@@ -1,14 +1,17 @@
 import sys
-import re
 
 from importlib.metadata import version
 
+import evaluate
 import polars as pl
 import gradio as gr
+from joblib import Parallel, delayed
 
-# Config
-concurrency_limit = 5
+# Load evaluators
+wer = evaluate.load("wer")
+cer = evaluate.load("cer")
 
+# Config
 title = "See ASR Outputs"
 
 # https://www.tablesgenerator.com/markdown_tables
@@ -27,8 +30,8 @@ Follow them on social networks and **contact** if you need any help or have any
 """.strip()
 
 examples = [
-    ["evaluation_results.jsonl", False],
-    ["evaluation_results_batch.jsonl", True],
+    ["evaluation_results.jsonl", False, True],
+    ["evaluation_results_batch.jsonl", True, True],
 ]
 
 description_head = f"""
@@ -36,7 +39,7 @@ description_head = f"""
 
 ## Overview
 
-See generated JSONL files made by ASR models as a dataframe.
+See generated JSONL files made by ASR models as a dataframe. Also, this app calculates WER and CER metrics for each row.
 """.strip()
 
 description_foot = f"""
@@ -57,17 +60,34 @@ tech_libraries = f"""
 #### Libraries
 
 - gradio: {version("gradio")}
+- jiwer: {version("jiwer")}
+- evaluate: {version("evaluate")}
 - polars: {version("polars")}
 """.strip()
 
 
-def inference(file_name, _batch_mode):
+def compute_wer(prediction, reference):
+    return round(wer.compute(predictions=[prediction], references=[reference]), 4)
+
+
+def compute_cer(prediction, reference):
+    return round(cer.compute(predictions=[prediction], references=[reference]), 4)
+
+
+def compute_batch_wer(predictions, references):
+    return round(wer.compute(predictions=predictions, references=references), 4)
+
+
+def compute_batch_cer(predictions, references):
+    return round(cer.compute(predictions=predictions, references=references), 4)
+
+
+def inference(file_name, _batch_mode, _calculate_metrics):
     if not file_name:
         raise gr.Error("Please paste your JSON file.")
 
     df = pl.read_ndjson(file_name)
 
-
     required_columns = [
         "filename",
         "inference_start",
@@ -105,9 +125,70 @@ def inference(file_name, _batch_mode):
         df = df.drop(["inference_start", "inference_end", "filename"])
 
     # round "inference_total" field to 2 decimal places
-    df = df.with_columns(pl.col("inference_total").round(2))
+    df = df.with_columns(pl.col("inference_total").round(2).alias("elapsed"))
+    df = df.drop(["inference_total"])
 
-    return df
+    # reassign columns
+    if _batch_mode:
+        if _calculate_metrics:
+            wer_values = Parallel(n_jobs=-1)(
+                delayed(compute_batch_wer)(row["predictions"], row["references"])
+                for row in df.iter_rows(named=True)
+            )
+            cer_values = Parallel(n_jobs=-1)(
+                delayed(compute_batch_cer)(row["predictions"], row["references"])
+                for row in df.iter_rows(named=True)
+            )
+
+            df.insert_column(2, pl.Series("wer", wer_values))
+            df.insert_column(3, pl.Series("cer", cer_values))
+
+            fields = [
+                "elapsed",
+                "durations",
+                "wer",
+                "cer",
+                "predictions",
+                "references",
+            ]
+        else:
+            fields = [
+                "elapsed",
+                "durations",
+                "predictions",
+                "references",
+            ]
+    else:
+        if _calculate_metrics:
+            wer_values = Parallel(n_jobs=-1)(
+                delayed(compute_wer)(row["prediction"], row["reference"])
+                for row in df.iter_rows(named=True)
+            )
+            cer_values = Parallel(n_jobs=-1)(
+                delayed(compute_cer)(row["prediction"], row["reference"])
+                for row in df.iter_rows(named=True)
+            )
+
+            df.insert_column(2, pl.Series("wer", wer_values))
+            df.insert_column(3, pl.Series("cer", cer_values))
+
+            fields = [
+                "elapsed",
+                "duration",
+                "wer",
+                "cer",
+                "prediction",
+                "reference",
+            ]
+        else:
+            fields = [
+                "elapsed",
+                "duration",
+                "prediction",
+                "reference",
+            ]
+
+    return df.select(fields)
 
 
 demo = gr.Blocks(
@@ -134,18 +215,21 @@ with demo:
                 label="Use batch mode",
             )
 
+            calculate_metrics = gr.Checkbox(
+                label="Calculate WER/CER metrics",
+                value=True,
+            )
 
     gr.Button("Show").click(
         inference,
-        concurrency_limit=concurrency_limit,
-        inputs=[jsonl_file, batch_mode],
+        inputs=[jsonl_file, batch_mode, calculate_metrics],
         outputs=df,
     )
 
     with gr.Row():
         gr.Examples(
             label="Choose an example",
-            inputs=[jsonl_file, batch_mode],
+            inputs=[jsonl_file, batch_mode, calculate_metrics],
             examples=examples,
         )
 

justfile

@@ -0,0 +1,5 @@
+check:
+    ruff check
+
+fmt: check
+    ruff format

requirements.txt

@@ -1,3 +1,7 @@
 gradio==5.23.0
 
 polars==1.26.0
+evaluate==0.4.3
+jiwer==3.1.0
+
+joblib==1.4.2

ruff.toml

@@ -0,0 +1,2 @@
+[lint]
+ignore = ["F403"]