Update app.py

app.py

@@ -4,45 +4,55 @@ from pydub import AudioSegment, silence
 import tempfile
 import torch
 import torchaudio
+import os
 
+# ---------------- Config ---------------- #
 MODEL_NAME = "mrmuminov/whisper-small-uz"
+SAMPLE_RATE = 16000
+MIN_LEN_MS = 15000
+MAX_LEN_MS = 25000
+SILENCE_THRESH = -40  # in dBFS
 
+# ---------------- Load Model ---------------- #
 processor = WhisperProcessor.from_pretrained(MODEL_NAME)
 model = WhisperForConditionalGeneration.from_pretrained(MODEL_NAME)
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model = model.to(device)
+model = model.to(device).eval()  # set to eval mode
 
-def split_on_silence_with_duration_control(audio, min_len, max_len, silence_thresh=-40):
+# ---------------- Chunking Logic ---------------- #
+def split_on_silence_with_duration_control(audio, min_len, max_len, silence_thresh):
     silences = silence.detect_silence(audio, min_silence_len=500, silence_thresh=silence_thresh)
-    silences = [((start + end) // 2) for start, end in silences]
+    silence_midpoints = [((start + end) // 2) for start, end in silences]
 
     chunks = []
     start = 0
-    while start < len(audio):
-        end = min(start + max_len, len(audio))
-        candidates = [s for s in silences if start + min_len <= s <= end]
-        split_point = candidates[-1] if candidates else end
-        chunks.append(audio[start:split_point])
+    duration = len(audio)
+
+    while start < duration:
+        end = min(start + max_len, duration)
+        valid_splits = [s for s in silence_midpoints if start + min_len <= s <= end]
+        split_point = valid_splits[-1] if valid_splits else end
+        chunk = audio[start:split_point]
+
+        # Avoid zero-length chunks
+        if len(chunk) > 0:
+            chunks.append(chunk)
+
         start = split_point
+
     return chunks
 
-def transcribe(audio_file):
-    # Load audio using pydub
-    audio = AudioSegment.from_file(audio_file)
-    
-    # Convert to mono and 16kHz if needed
-    if audio.channels > 1:
-        audio = audio.set_channels(1)
-    if audio.frame_rate != 16000:
-        audio = audio.set_frame_rate(16000)
-
-    # Detect silent chunks
+# ---------------- Transcription ---------------- #
+def transcribe(audio_file_path):
+    audio = AudioSegment.from_file(audio_file_path)
+
+    # Ensure mono and target sample rate
+    audio = audio.set_channels(1).set_frame_rate(SAMPLE_RATE)
+
     chunks = split_on_silence_with_duration_control(
-        audio, min_len=15000, max_len=25000, silence_thresh=-40
+        audio, min_len=MIN_LEN_MS, max_len=MAX_LEN_MS, silence_thresh=SILENCE_THRESH
     )
 
-    # Transcribe each chunk
     results = []
     for chunk in chunks:
         with tempfile.NamedTemporaryFile(suffix=".wav", delete=True) as tmpfile:
@@ -50,29 +60,28 @@ def transcribe(audio_file):
             waveform, _ = torchaudio.load(tmpfile.name)
             input_features = processor(
                 waveform.squeeze().numpy(),
-                sampling_rate=16000,
+                sampling_rate=SAMPLE_RATE,
                 return_tensors="pt",
                 language="uz"
             ).input_features.to(device)
 
             with torch.no_grad():
                 predicted_ids = model.generate(input_features)
-                transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
-                results.append(transcription)
+                text = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
+                results.append(text)
 
     return " ".join(results)
 
-demo = gr.Blocks()
+# ---------------- Gradio UI ---------------- #
+with gr.Blocks() as demo:
+    gr.Markdown("### " + MODEL_NAME + " Transcribe Uzbek Audio")
 
-file_transcribe = gr.Interface(
-    fn=transcribe,
-    inputs=gr.Audio(type="filepath", label="Audio file"),
-    outputs="text",
-    title="mrmuminov/whisper-small-uz: Transcribe Audio",
-    description="mrmuminov/whisper-small-uz fine-tuned for Uzbek language",
-)
+    file_transcribe = gr.Interface(
+        fn=transcribe,
+        inputs=gr.Audio(type="filepath", label="Upload Audio"),
+        outputs=gr.Textbox(label="Transcription"),
+    )
 
-with demo:
-    gr.TabbedInterface([file_transcribe], ["Audio file"])
+    gr.TabbedInterface([file_transcribe], ["Audio File"])
 
 demo.launch()