Update app.py

app.py

@@ -3,60 +3,35 @@ import librosa
 import torch
 from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, MBartForConditionalGeneration, MBart50Tokenizer
 
-# Load pre-trained model and processor directly from Hugging Face Hub
+# Load pre-trained models
 model = Wav2Vec2ForCTC.from_pretrained("boumehdi/wav2vec2-large-xlsr-moroccan-darija")
 processor = Wav2Vec2Processor.from_pretrained("boumehdi/wav2vec2-large-xlsr-moroccan-darija")
 
-# Load translation model
 translation_model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
 translation_tokenizer = MBart50Tokenizer.from_pretrained("facebook/mbart-large-50-many-to-many-mmt", src_lang="ar_AR")
 
 def transcribe_audio(audio):
-    if not audio:
-        return "No audio file provided", "No translation available"
-
-    # Load the audio file
     audio_array, sr = librosa.load(audio, sr=16000)
-    
-    # Tokenize the audio file
     input_values = processor(audio_array, return_tensors="pt", padding=True).input_values
-    
-    # Get the model's logits (predicted token scores)
     logits = model(input_values).logits
-    
-    # Get the predicted tokens
     tokens = torch.argmax(logits, axis=-1)
-    
-    # Decode the tokens into text (Darija transcription)
     transcription = processor.decode(tokens[0])
-
-    # Translate the transcription to English
     translation = translate_text(transcription)
-    
     return transcription, translation
 
 def translate_text(text):
-    # Tokenize the text to translate
     inputs = translation_tokenizer(text, return_tensors="pt")
-    
-    # Generate translated tokens (from Darija to English)
-    translated_tokens = translation_model.generate(**inputs, forced_bos_token_id=translation_tokenizer.lang_code_to_id["en"])
-    
-    # Decode the translated tokens into text
+    translated_tokens = translation_model.generate(**inputs, forced_bos_token_id=translation_tokenizer.lang_code_to_id["en_XX"])
     translated_text = translation_tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
-    
     return translated_text
 
-# Create Gradio Blocks for better UI
 with gr.Blocks() as demo:
     gr.Markdown("# Speech-to-Text and Translation")
-    gr.Markdown("Upload an audio file to transcribe and translate it from Darija to English.")
-    
-    audio_input = gr.Audio(type="filepath", label="Upload or Record Audio")
+    audio_input = gr.Audio(type="filepath")
     submit_button = gr.Button("Submit")
-    transcription_output = gr.Textbox(label="Transcription (Darija)")
-    translation_output = gr.Textbox(label="Translation (English)")
-    
-    submit_button.click(fn=transcribe_audio, inputs=audio_input, outputs=[transcription_output, translation_output])
+    transcription_output = gr.Textbox(label="Transcription")
+    translation_output = gr.Textbox(label="Translation")
+
+    submit_button.click(transcribe_audio, inputs=[audio_input], outputs=[transcription_output, translation_output])
 
 demo.launch()