ArToEngModel2

app.py

@@ -1,42 +1,50 @@
 import gradio as gr
 import librosa
 import torch
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, MBartForConditionalGeneration, MBart50Tokenizer, MarianMTModel, MarianTokenizer
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, MarianMTModel, MarianTokenizer
 
-# Load pre-trained models
+# Charger le modèle de transcription pour le Darija
 model = Wav2Vec2ForCTC.from_pretrained("boumehdi/wav2vec2-large-xlsr-moroccan-darija")
 processor = Wav2Vec2Processor.from_pretrained("boumehdi/wav2vec2-large-xlsr-moroccan-darija")
 
-#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")
-
 # Charger le modèle de traduction Arabe -> Anglais
 translation_model_name = "Helsinki-NLP/opus-mt-ar-en"
 translation_model = MarianMTModel.from_pretrained(translation_model_name)
 translation_tokenizer = MarianTokenizer.from_pretrained(translation_model_name)
 
-
 def transcribe_audio(audio):
+    """Convertir l'audio en texte et le traduire en anglais"""
+    # Charger et prétraiter l'audio
     audio_array, sr = librosa.load(audio, sr=16000)
     input_values = processor(audio_array, return_tensors="pt", padding=True).input_values
+
+    # Obtenir les prédictions du modèle
     logits = model(input_values).logits
     tokens = torch.argmax(logits, axis=-1)
+
+    # Décoder la transcription en Darija
     transcription = processor.decode(tokens[0])
+
+    # Traduire en anglais
     translation = translate_text(transcription)
+
     return transcription, translation
 
 def translate_text(text):
-    inputs = translation_tokenizer(text, return_tensors="pt")
-    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]
+    """Traduire le texte de l'arabe vers l'anglais"""
+    inputs = translation_tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    translated_tokens = translation_model.generate(**inputs)
+    translated_text = translation_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
     return translated_text
 
+# Interface utilisateur avec Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("# Speech-to-Text and Translation")
-    audio_input = gr.Audio(type="filepath")
-    submit_button = gr.Button("Submit")
-    transcription_output = gr.Textbox(label="Transcription")
-    translation_output = gr.Textbox(label="Translation")
+    gr.Markdown("# 🎙️ Speech-to-Text & Translation")
+    
+    audio_input = gr.Audio(type="filepath", label="Upload Audio or Record")
+    submit_button = gr.Button("Transcribe & Translate")
+    transcription_output = gr.Textbox(label="Transcription (Darija)")
+    translation_output = gr.Textbox(label="Translation (English)")
 
     submit_button.click(transcribe_audio, inputs=[audio_input], outputs=[transcription_output, translation_output])