Update app.py

app.py

@@ -1,59 +1,66 @@
 import os
 import torch
 import gradio as gr
-import mimetypes
 from pydub import AudioSegment
-from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
+import tempfile
+import math
 
-# Set device and precision
+from datasets import load_dataset, Audio
+import numpy as np
+import torchaudio
+
+# Set up model
 device = "cpu"
 torch_dtype = torch.float32
 
-# Load KB-Whisper model
 model_id = "KBLab/kb-whisper-large"
 
-model = AutoModelForSpeechSeq2Seq.from_pretrained(
-    model_id, torch_dtype=torch_dtype
-).to(device)
-
 processor = AutoProcessor.from_pretrained(model_id)
+model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype).to(device)
+
+# Helper: Split audio into chunks
+def split_audio(audio_path, chunk_duration_ms=10000):
+    audio = AudioSegment.from_file(audio_path)
+    chunks = [audio[i:i + chunk_duration_ms] for i in range(0, len(audio), chunk_duration_ms)]
+    return chunks
+
+# Helper: Transcribe a single chunk
+def transcribe_chunk(chunk):
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmpfile:
+        chunk.export(tmpfile.name, format="wav")
+        input_audio, _ = torchaudio.load(tmpfile.name)
+        input_features = processor(input_audio.squeeze(), sampling_rate=16000, return_tensors="pt").input_features
+        input_features = input_features.to(device)
+        predicted_ids = model.generate(input_features, task="transcribe", language="sv")
+        transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
+        os.remove(tmpfile.name)
+        return transcription
+
+# Full transcription function with progress
+def transcribe_with_progress(audio_path, progress=gr.Progress()):
+    ext = os.path.splitext(audio_path)[1].lower()
+    if ext != ".wav":
+        sound = AudioSegment.from_file(audio_path)
+        audio_path = audio_path.replace(ext, ".converted.wav")
+        sound.export(audio_path, format="wav")
+
+    chunks = split_audio(audio_path, chunk_duration_ms=8000)
+    full_transcript = ""
+    total_chunks = len(chunks)
+
+    for i, chunk in enumerate(chunks):
+        partial_text = transcribe_chunk(chunk)
+        full_transcript += partial_text + " "
+        progress(i + 1, total_chunks)  # Update progress bar
+        yield full_transcript.strip()  # Stream updated text to UI
 
-pipe = pipeline(
-    "automatic-speech-recognition",
-    model=model,
-    tokenizer=processor.tokenizer,
-    feature_extractor=processor.feature_extractor,
-    device=device,
-    torch_dtype=torch_dtype,
-)
-
-def transcribe(audio_path):
-    try:
-        # Get file extension
-        ext = os.path.splitext(audio_path)[1].lower()
-
-        # Convert to WAV if not already
-        if ext != ".wav":
-            try:
-                sound = AudioSegment.from_file(audio_path)
-                converted_path = audio_path.replace(ext, ".converted.wav")
-                sound.export(converted_path, format="wav")
-                audio_path = converted_path
-            except Exception as e:
-                return f"Error converting audio to WAV: {str(e)}"
-
-        # Transcribe
-        result = pipe(audio_path, chunk_length_s=30, generate_kwargs={"task": "transcribe", "language": "sv"})
-        return result["text"]
-
-    except Exception as e:
-        return f"Transcription failed: {str(e)}"
-
-# Gradio UI
+# UI
 gr.Interface(
-    fn=transcribe,
-    inputs=gr.Audio(type="filepath", label="Upload Audio (.m4a, .mp3, .wav)"),
-    outputs=gr.Textbox(label="Swedish Transcript"),
-    title="Swedish Speech Transcriber with KB-Whisper",
-    description="Supports .m4a, .mp3, .wav files. Transcribes spoken Swedish using KBLab's Whisper Large model. May take time on CPU.",
-).launch()
+    fn=transcribe_with_progress,
+    inputs=gr.Audio(type="filepath", label="Upload Swedish Audio"),
+    outputs=gr.Textbox(label="Live Transcript (Swedish)"),
+    title="Live Swedish Transcriber (KB-Whisper)",
+    description="Streams transcription word-by-word with visual progress. Supports .m4a, .mp3, .wav. May be slow on CPU.",
+    live=True
+).launch()