import gradio as gr
from PIL import Image  
import re
import os
import speech_recognition as sr


r = sr.Recognizer()

def transcribe(audio):
    with sr.AudioFile(audio) as source:
        audio_ = r.listen(source) 
        text = r.recognize_google(audio_)#, language = 'en-IN')# , show_all=True)
    return text


import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

# this model was loaded from https://hf.co/models
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
device = -1 #0 if torch.cuda.is_available() else -1
LANGS = ["eng_Latn", "fra_Latn", "spa_Latn", "hin_Deva"]
langs_mapper = {
    "English": "eng_Latn", 
    "French": "fra_Latn", 
    "Spanish": "spa_Latn", 
    "Hindi": "hin_Deva"
}
def translate(text, src_lang, tgt_lang):
    """
    Translate the text from source lang to target lang
    """
    translation_pipeline = pipeline("translation", model=model, tokenizer=tokenizer, src_lang=src_lang, tgt_lang=tgt_lang, max_length=400, device=device)
    result = translation_pipeline(text)
    return result[0]['translation_text']


with gr.Blocks() as demo:

    gr.Markdown("Translate your voice to any of following language - Speech to text -> Translation - Facebook NLL")
    with gr.Tab("Audio Input"):
        sr_lang1 = gr.Dropdown(label="Source Language", choices=LANGS)
        tr_lang1 = gr.Dropdown(label="Target Language", choices=LANGS)
        audio_input = gr.Audio(source="microphone", type="filepath")
        submit_audio_button = gr.Button("Translate")
        text_output = gr.Textbox(label="You said:")
        
    with gr.Tab("Text Input"):
        sr_lang2 = gr.Dropdown(label="Source Language", choices=LANGS)
        tr_lang2 = gr.Dropdown(label="Target Language", choices=LANGS)
        text_input = gr.Textbox(label="Enter text")
        submit_button_text = gr.Button("Translate")
    
    translated_output = gr.Textbox(label="Output Box")

    submit_audio_button.click(fn=transcribe, inputs=audio_input, outputs=text_output)
    text_output.change(fn=translate, inputs=[text_output, sr_lang1, tr_lang1], outputs=translated_output)
    submit_button_text.click(fn=translate, inputs=[text_input, sr_lang2, tr_lang2], outputs=translated_output)


demo.launch()