Update app.py

app.py

@@ -101,7 +101,7 @@ def clean_value(x):
     return s
 
 
-def inference(file_name, _clear_punctuation, _show_chars, _batch_mode):
+def inference(file_name, _clear_punctuation_references, _clear_punctuation_predictions, _show_chars, _batch_mode):
     if not file_name:
         raise gr.Error("Please paste your JSON file.")
 
@@ -143,19 +143,25 @@ def inference(file_name, _clear_punctuation, _show_chars, _batch_mode):
         references_batch = df["references"]
         predictions_batch = df["predictions"]
 
+        references = []
+        for reference in references_batch:
+            if _clear_punctuation_references:
+                reference = reference.map_elements(
+                    clean_value, return_dtype=pl.String
+                )
+                references.extend(reference)
+            else:
+                references.extend(reference)
+
         predictions = []
         for prediction in predictions_batch:
-            if _clear_punctuation:
+            if _clear_punctuation_predictions:
                 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:
         if not all(col in df.columns for col in required_columns):
             raise gr.Error(
@@ -166,9 +172,14 @@ def inference(file_name, _clear_punctuation, _show_chars, _batch_mode):
 
         rtf = inference_seconds / duration_seconds
 
-        references = df["reference"]
+        if _clear_punctuation_references:
+            references = df["references"].map_elements(
+                clean_value, return_dtype=pl.String
+            )
+        else:
+            references = df["reference"]
 
-        if _clear_punctuation:
+        if _clear_punctuation_predictions:
             predictions = df["prediction"].map_elements(
                 clean_value, return_dtype=pl.String
             )
@@ -238,8 +249,11 @@ with demo:
         with gr.Column():
             jsonl_file = gr.File(label="A JSONL file")
 
-            clear_punctuation = gr.Checkbox(
-                label="Clear punctuation, some chars and convert to lowercase",
+            clear_punctuation_references = gr.Checkbox(
+                label="Clear punctuation (in references), some chars and convert to lowercase",
+            )
+            clear_punctuation_predictions = gr.Checkbox(
+                label="Clear punctuation (in predictions), some chars and convert to lowercase",
             )
             show_chars = gr.Checkbox(
                 label="Show chars in predictions",
@@ -257,7 +271,7 @@ with demo:
     gr.Button("Calculate").click(
         inference,
         concurrency_limit=concurrency_limit,
-        inputs=[jsonl_file, clear_punctuation, show_chars, batch_mode],
+        inputs=[jsonl_file, clear_punctuation_references, clear_punctuation_predictions, show_chars, batch_mode],
         outputs=metrics,
     )