Update app.py

app.py

@@ -1,87 +1,38 @@
 import gradio as gr
+import torch
 from transformers import pipeline
-from jiwer import wer
-import difflib
-import matplotlib.pyplot as plt
-import numpy as np
 
-# Load Model 1
-model_name_1 = "maliahson/whisper-agri"
-whisper_pipeline_1 = pipeline("automatic-speech-recognition", model=model_name_1)
+# Load models
+whisper_pipeline_1 = pipeline("automatic-speech-recognition", model="maliahson/whisper-agri")
+device = 0 if torch.cuda.is_available() else "cpu"
+whisper_pipeline_2 = pipeline("automatic-speech-recognition", model="openai/whisper-large-v3-turbo", device=device)
 
-# Load Model 2
-model_name_2 = "openai/whisper-large-v3-turbo"
-whisper_pipeline_2 = pipeline("automatic-speech-recognition", model=model_name_2)
+def transcribe_and_compare(audio_path, original_transcription=None):
+    transcription_1 = whisper_pipeline_1(audio_path)["text"]
+    transcription_2 = whisper_pipeline_2(audio_path)["text"]
+    comparison_result = {"Model 1 Output": transcription_1, "Model 2 Output": transcription_2}
 
-def compare_transcriptions(audio, original_transcription=None):
-    """
-    Compares transcriptions from two models and optionally calculates WER.
-    Args:
-        audio: Uploaded audio file.
-        original_transcription: The ground-truth transcription (optional).
-    Returns:
-        Transcriptions from both models, and comparison results.
-    """
-    # Transcribe audio using Model 1
-    transcription_1 = whisper_pipeline_1(audio)["text"]
-    
-    # Transcribe audio using Model 2
-    transcription_2 = whisper_pipeline_2(audio)["text"]
-    
-    result = {
-        "Model 1 Transcription": transcription_1,
-        "Model 2 Transcription": transcription_2
-    }
-    
     if original_transcription:
-        # Calculate WER for both models
+        from jiwer import wer
         wer_1 = wer(original_transcription, transcription_1)
         wer_2 = wer(original_transcription, transcription_2)
-        result["WER (Model 1)"] = f"{wer_1:.2%}"
-        result["WER (Model 2)"] = f"{wer_2:.2%}"
-        
-        # Create a WER comparison plot
-        fig, ax = plt.subplots()
-        bars = ax.bar(["Model 1", "Model 2"], [wer_1, wer_2], color=["blue", "orange"])
-        ax.set_title("Word Error Rate Comparison")
-        ax.set_ylabel("WER")
-        ax.set_ylim(0, 1)
-        ax.bar_label(bars, fmt="%.2f")
-        plt.tight_layout()
-        return result, fig
-    else:
-        # Calculate and visualize differences between transcriptions
-        diff = difflib.unified_diff(
-            transcription_1.split(),
-            transcription_2.split(),
-            lineterm="",
-            n=0
-        )
-        diff_result = "\n".join(diff)
-        result["Differences"] = diff_result
-        return result, None
+        comparison_result["WER Model 1"] = wer_1
+        comparison_result["WER Model 2"] = wer_2
+
+    return comparison_result
 
 # Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Audio Transcription Comparison with Whisper Models")
-    gr.Markdown("Upload an audio file, and optionally provide the original transcription to compare the outputs.")
-    
-    with gr.Row():
-        audio_input = gr.Audio(source="upload", type="filepath", label="Upload or Record Audio")
-        original_transcription = gr.Textbox(
-            placeholder="Enter original transcription here (optional)",
-            label="Original Transcription"
-        )
-    
-    compare_button = gr.Button("Compare Transcriptions")
-    result_output = gr.JSON(label="Transcriptions and Comparison")
-    wer_plot = gr.Plot(label="Word Error Rate Comparison (if applicable)")
-    
-    compare_button.click(
-        compare_transcriptions,
+    gr.Markdown("## Audio Transcription and Comparison")
+    audio_input = gr.Audio(type="filepath", label="Upload or Record Audio (supports upload and recording)")
+    original_transcription = gr.Textbox(lines=2, label="Original Transcription (Optional)")
+    output = gr.JSON(label="Comparison Results")
+    submit_btn = gr.Button("Transcribe and Compare")
+
+    submit_btn.click(
+        transcribe_and_compare,
         inputs=[audio_input, original_transcription],
-        outputs=[result_output, wer_plot]
+        outputs=output
     )
 
-# Launch app
-demo.launch(enable_queue=True)
+demo.launch(debug=True)