Update app.py

app.py

@@ -2,34 +2,53 @@ import joblib
 import torch
 import numpy as np
 import soundfile as sf
+import torchaudio
 from transformers import pipeline
 import gradio as gr
 
 # Load your text classifier
 text_clf = joblib.load("text_pipeline_balanced.joblib")
 
-# Initialize Whisper ASR
+# Choose GPU if available
 device = 0 if torch.cuda.is_available() else -1
+
+# Initialize Whisper-Large ASR with beam search
 asr = pipeline(
     "automatic-speech-recognition",
-    model="openai/whisper-base",
-    chunk_length_s=30,
+    model="openai/whisper-large-v2",
+    chunk_length_s=10,
     device=device,
-    generate_kwargs={"language": "en", "task": "transcribe"}
+    generate_kwargs={
+        "language": "en",
+        "task": "transcribe",
+        "num_beams": 5,
+        "best_of": 5,
+    },
+    ignore_warning=True
 )
 
 def classify(audio_path):
     """
-    audio_path: str → path to uploaded/recorded file
+    audio_path: str – path to the uploaded/recorded file
     returns: transcript (str), safety probabilities (dict), unsafe probability (str)
     """
-    # Read file from disk
+    # Read & (re)sample
     audio, sr = sf.read(audio_path, dtype="float32")
-    # If stereo, convert to mono
+    if sr != 16000:
+        audio = torchaudio.functional.resample(
+            torch.from_numpy(audio), sr, 16000
+        ).numpy()
+        sr = 16000
+
+    # Stereo → mono
     if audio.ndim > 1:
         audio = audio.mean(axis=1)
 
-    # ASR
+    # Normalize peak amplitude
+    peak = np.abs(audio).max() or 1.0
+    audio = audio / peak
+
+    # ASR transcription
     result = asr({"array": audio, "sampling_rate": sr})
     txt = result["text"].strip()
 
@@ -37,18 +56,19 @@ def classify(audio_path):
     proba = float(text_clf.predict_proba([txt])[0][1])
     label_probs = {"safe": 1 - proba, "unsafe": proba}
     unsafe_str = f"{proba:.2f}"
+
     return txt, label_probs, unsafe_str
 
-# Use filepath-based Audio component
-audio_input     = gr.components.Audio(label="Upload or record audio", type="filepath")
-transcript_out  = gr.components.Textbox(label="Transcript")
-probs_out       = gr.components.Label(num_top_classes=2, label="Safety Probabilities")
-unsafe_out      = gr.components.Textbox(label="Unsafe Probability")
+# Gradio components
+audio_input      = gr.components.Audio(label="Upload or record audio", type="filepath")
+transcript_out   = gr.components.Textbox(label="Transcript")
+probs_out        = gr.components.Label(num_top_classes=2, label="Safety Probabilities")
+unsafe_prob_out  = gr.components.Textbox(label="Unsafe Probability")
 
 iface = gr.Interface(
     fn=classify,
     inputs=audio_input,
-    outputs=[transcript_out, probs_out, unsafe_out],
+    outputs=[transcript_out, probs_out, unsafe_prob_out],
     title="BubbleGuard Audio Safety Checker",
     description="Upload or record audio; get ASR transcript plus safe/unsafe probabilities."
 )