Update app.py

app.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
 import gradio as gr
 import librosa
 
@@ -9,26 +9,29 @@ BATCH_SIZE = 8
 device = 0 if torch.cuda.is_available() else "cpu"
 
 # Load model and processor
-model = AutoModelForSpeechSeq2Seq.from_pretrained(MODEL_NAME)
-processor = AutoProcessor.from_pretrained(MODEL_NAME)
+model_name = "EwoutLagendijk/whisper-small-indonesian"
+
+model = AutoModelForSpeechSeq2Seq.from_pretrained(model_name)
+processor = AutoProcessor.from_pretrained(model_name)
 
 # Update the generation config for transcription
 model.config.forced_decoder_ids = processor.get_decoder_prompt_ids(language="id", task="transcribe")
-model.config.no_repeat_ngram_size = 3
 
-def transcribe_speech_with_timestamps(filepath):
+# Initialize the translation pipeline (using a model like `Helsinki-NLP/opus-mt-id-en` for Indonesian to English)
+translation_pipeline = pipeline("translation", model="Helsinki-NLP/opus-mt-id-en")
+
+def transcribe_speech(filepath):
     # Load the audio
     audio, sampling_rate = librosa.load(filepath, sr=16000)
 
     # Define chunk size (e.g., 30 seconds)
-    chunk_duration = 30  # in seconds
+    chunk_duration = 5  # in seconds
     chunk_samples = chunk_duration * sampling_rate
 
     # Process audio in chunks
     transcription = []
     for i in range(0, len(audio), chunk_samples):
         chunk = audio[i:i + chunk_samples]
-        chunk_start_time = i / sampling_rate  # Calculate chunk start time in seconds
 
         # Convert the chunk into input features
         inputs = processor(audio=chunk, sampling_rate=16000, return_tensors="pt").input_features
@@ -36,25 +39,22 @@ def transcribe_speech_with_timestamps(filepath):
         # Generate transcription for the chunk
         generated_ids = model.generate(
             inputs,
-            max_new_tokens=444,
-            return_dict_in_generate=True,
-            output_scores=False,
-            output_attentions=False,
-            output_hidden_states=False,
+            max_new_tokens=444,  # Max allowed by Whisper
             forced_decoder_ids=processor.get_decoder_prompt_ids(language="id", task="transcribe")
         )
 
-        # Decode the tokens into text and timestamps
-        token_transcriptions = processor.batch_decode(generated_ids["sequences"], skip_special_tokens=False)[0]
-        decoded_with_timestamps = processor.decode_with_timestamps(generated_ids["sequences"][0])
-        
-        # Parse timestamps and transcription
-        for segment in decoded_with_timestamps:
-            start_time = chunk_start_time + segment['start']
-            end_time = chunk_start_time + segment['end']
-            text = segment['text']
-            transcription.append(f"[{start_time:.2f}s - {end_time:.2f}s]: {text}")
-    
+        # Decode and append the transcription
+        chunk_transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+        # Translate the transcription to English (or another language of choice)
+        chunk_translation = translation_pipeline(chunk_transcription)[0]['translation_text']
+
+        # Append both transcription and translation
+        transcription.append(f"Chunk {i//chunk_samples + 1}:\n")
+        transcription.append(f"Transcription: {chunk_transcription}\n")
+        transcription.append(f"Translation: {chunk_translation}\n\n")
+
+    # Combine all chunk transcriptions and translations into a single string
     return "\n".join(transcription)
 
 demo = gr.Blocks()
@@ -72,6 +72,6 @@ file_transcribe = gr.Interface(
 )
 
 with demo:
-    gr.TabbedInterface([mic_transcribe, file_transcribe], ["Transcribe Microphone", "Transcribe Audio File"])
+    gr.TabbedInterface([mic_transcribe, file_transcribe], ["Transcribe and translate Microphone", "Transcribe and translate Audio File"])
 
 demo.launch(debug=True)