Add batch mode

app.py

@@ -33,7 +33,8 @@ Follow them on social networks and **contact** if you need any help or have any
 """.strip()
 
 examples = [
-    ["evaluation_results.jsonl", True, False],
+    ["evaluation_results.jsonl", True, False, False],
+    ["evaluation_results_batch.jsonl", True, False, True],
 ]
 
 description_head = f"""
@@ -100,25 +101,58 @@ def clean_value(x):
     return s
 
 
-def inference(file_name, clear_punctuation, show_chars, progress=gr.Progress()):
+def inference(file_name, _clear_punctuation, _show_chars, _batch_mode):
     if not file_name:
         raise gr.Error("Please paste your JSON file.")
 
-    progress(0, desc="Calculating...")
-
     df = pl.read_ndjson(file_name)
 
     inference_seconds = df["inference_total"].sum()
-    duration_seconds = df["duration"].sum()
 
-    rtf = inference_seconds / duration_seconds
+    if _batch_mode:
+        if "durations" not in df.columns:
+            raise gr.Error(
+                "Please use a JSONL file with 'durations' column for batch mode."
+            )
+
+        duration_seconds = 0
+        for durations in df["durations"]:
+            duration_seconds += durations.sum()
+
+        rtf = inference_seconds / duration_seconds
 
-    references = df["reference"]
+        references_batch = df["references"]
+        predictions_batch = df["predictions"]
 
-    if clear_punctuation:
-        predictions = df["prediction"].map_elements(clean_value, return_dtype=pl.String)
+        predictions = []
+        for prediction in predictions_batch:
+            if _clear_punctuation:
+                prediction = prediction.map_elements(
+                    clean_value, return_dtype=pl.String
+                )
+                predictions.extend(prediction)
+            else:
+                predictions.extend(prediction)
+
+        references = []
+        for reference in references_batch:
+            references.extend(reference)
     else:
-        predictions = df["prediction"]
+        duration_seconds = df["duration"].sum()
+
+        rtf = inference_seconds / duration_seconds
+
+        references = df["reference"]
+
+        if _clear_punctuation:
+            predictions = df["prediction"].map_elements(
+                clean_value, return_dtype=pl.String
+            )
+        else:
+            predictions = df["prediction"]
+
+    n_predictions = len(predictions)
+    n_references = len(references)
 
     # Evaluate
     wer_value = round(wer.compute(predictions=predictions, references=references), 4)
@@ -131,6 +165,10 @@ def inference(file_name, clear_punctuation, show_chars, progress=gr.Progress()):
 
     results = []
 
+    results.append(
+        f"- Number of references / predictions: {n_references} / {n_predictions}"
+    )
+    results.append(f"")
     results.append(f"- WER: {wer_value} metric, {round(wer_value * 100, 4)}%")
     results.append(f"- CER: {cer_value} metric, {round(cer_value * 100, 4)}%")
     results.append("")
@@ -146,7 +184,7 @@ def inference(file_name, clear_punctuation, show_chars, progress=gr.Progress()):
     results.append("")
     results.append(f"- RTF: {round(rtf, 4)}")
 
-    if show_chars:
+    if _show_chars:
         all_chars = set()
         for pred in predictions:
             for c in pred:
@@ -175,12 +213,16 @@ with demo:
     with gr.Row():
         with gr.Column():
             jsonl_file = gr.File(label="A JSONL file")
+
             clear_punctuation = gr.Checkbox(
                 label="Clear punctuation, some chars and convert to lowercase",
             )
             show_chars = gr.Checkbox(
                 label="Show chars in predictions",
             )
+            batch_mode = gr.Checkbox(
+                label="Use batch mode",
+            )
 
         metrics = gr.Textbox(
             label="Metrics",
@@ -191,14 +233,14 @@ with demo:
     gr.Button("Calculate").click(
         inference,
         concurrency_limit=concurrency_limit,
-        inputs=[jsonl_file, clear_punctuation, show_chars],
+        inputs=[jsonl_file, clear_punctuation, show_chars, batch_mode],
         outputs=metrics,
     )
 
     with gr.Row():
         gr.Examples(
             label="Choose an example",
-            inputs=[jsonl_file, clear_punctuation, show_chars],
+            inputs=[jsonl_file, clear_punctuation, show_chars, batch_mode],
             examples=examples,
         )