Update app.py

app.py

@@ -1,52 +1,63 @@
-from transformers import pipeline
-asr = pipeline(
-    task="automatic-speech-recognition",
-    model="openai/whisper-small",  # multilingual model
-    generate_kwargs={"task": "translate", "language": "es"}
-)
-
-
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
+import torchaudio
+import torch
 import os
 import gradio as gr
 from pydub import AudioSegment
 
-demo = gr.Blocks()
-def transcribe_speech(filepath):
+# Load Whisper model and processor
+model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
+processor = WhisperProcessor.from_pretrained("openai/whisper-small")
+
+# Get decoder prompts for translation to Spanish
+forced_decoder_ids = processor.get_decoder_prompt_ids(language="es", task="translate")
+
+# Function to process and translate audio
+def translate_audio(filepath):
     if filepath is None:
         gr.Warning("No audio found, please retry.")
         return ""
 
-    # Load audio using pydub
+    # Load audio using pydub for chunking
     audio = AudioSegment.from_file(filepath)
     chunk_length_ms = 30 * 1000  # 30 seconds
     chunks = [audio[i:i + chunk_length_ms] for i in range(0, len(audio), chunk_length_ms)]
 
-    full_transcription = ""
+    full_translation = ""
 
     for i, chunk in enumerate(chunks):
-        # Export chunk to temporary wav file
         chunk_path = f"chunk_{i}.wav"
         chunk.export(chunk_path, format="wav")
 
-        # Transcribe the chunk
-        result = asr(chunk_path)
-        full_transcription += result["text"] + " "
+        # Load chunk for model input
+        waveform, sample_rate = torchaudio.load(chunk_path)
+        inputs = processor(waveform[0], sampling_rate=sample_rate, return_tensors="pt")
+
+        # Generate translated output
+        with torch.no_grad():
+            generated_ids = model.generate(
+                inputs["input_features"],
+                forced_decoder_ids=forced_decoder_ids
+            )
+
+        translation = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        full_translation += translation + " "
 
-        # Clean up (optional)
         os.remove(chunk_path)
 
-    return full_transcription.strip()
-    
+    return full_translation.strip()
+
+# Gradio UI components
 mic_transcribe = gr.Interface(
-    fn=transcribe_speech,
+    fn=translate_audio,
     inputs=gr.Audio(sources="microphone",
                     type="filepath"),
-    outputs=gr.Textbox(label="Transcription",
+    outputs=gr.Textbox(label="Translation (English to Spanish)",
                        lines=3),
     allow_flagging="never")
 
 file_transcribe = gr.Interface(
-    fn=transcribe_speech,
+    fn=translate_audio,
     inputs=gr.Audio(sources="upload",
                     type="filepath"),
     outputs=gr.Textbox(label="Translation (English to Spanish)",