Update app.py

app.py

@@ -1,6 +1,5 @@
 import torch
-from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
-from transformers.pipelines.audio_utils import ffmpeg_read
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
 import gradio as gr
 import librosa
 
@@ -10,15 +9,14 @@ BATCH_SIZE = 8
 device = 0 if torch.cuda.is_available() else "cpu"
 
 # Load model and processor
-model_name = "EwoutLagendijk/whisper-small-indonesian"
-
-model = AutoModelForSpeechSeq2Seq.from_pretrained(model_name)
-processor = AutoProcessor.from_pretrained(model_name)
+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(filepath):
+def transcribe_speech_with_timestamps(filepath):
     # Load the audio
     audio, sampling_rate = librosa.load(filepath, sr=16000)
 
@@ -30,6 +28,7 @@ def transcribe_speech(filepath):
     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
@@ -37,35 +36,42 @@ def transcribe_speech(filepath):
         # Generate transcription for the chunk
         generated_ids = model.generate(
             inputs,
-            max_new_tokens=444,  # Max allowed by Whisper
-            forced_decoder_ids=processor.get_decoder_prompt_ids(language="id", task="transcribe"),
-            return_timestamps = True
+            max_new_tokens=444,
+            return_dict_in_generate=True,
+            output_scores=False,
+            output_attentions=False,
+            output_hidden_states=False,
+            forced_decoder_ids=processor.get_decoder_prompt_ids(language="id", task="transcribe")
         )
 
-        # Decode and append the transcription
-        chunk_transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        transcription.append(chunk_transcription)
-
-    # Combine all chunk transcriptions into a single string
-    return " ".join(transcription)
-
-
+        # 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}")
+    
+    return "\n".join(transcription)
 
 demo = gr.Blocks()
 
 mic_transcribe = gr.Interface(
-    fn=transcribe_speech,
+    fn=transcribe_speech_with_timestamps,
     inputs=gr.Audio(sources="microphone", type="filepath"),
-    outputs=gr.components.Textbox(),
+    outputs=gr.Textbox(lines=10, label="Transcription with Timestamps"),
 )
 
 file_transcribe = gr.Interface(
-    fn=transcribe_speech,
+    fn=transcribe_speech_with_timestamps,
     inputs=gr.Audio(sources="upload", type="filepath"),
-    outputs=gr.components.Textbox(),
+    outputs=gr.Textbox(lines=10, label="Transcription with Timestamps"),
 )
 
 with demo:
     gr.TabbedInterface([mic_transcribe, file_transcribe], ["Transcribe Microphone", "Transcribe Audio File"])
 
-demo.launch(debug=True)
+demo.launch(debug=True)