app.py

import gradio as gr
import librosa
import torch
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, MBartForConditionalGeneration, MBart50Tokenizer

# Load pre-trained model and processor directly from Hugging Face Hub
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_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")
    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])

demo.launch()