import gradio as gr
import whisper
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast

## Cargar modelos de Whisper
whisper_model = whisper.load_model("base.en")

## Cargar modelos de MBart
translation_model = MBartForConditionalGeneration.from_pretrained("SnypzZz/Llama2-13b-Language-translate")
tokenizer = MBart50TokenizerFast.from_pretrained("SnypzZz/Llama2-13b-Language-translate", src_lang="en_XX")

## Función para transcribir y traducir el audio
def transcribe_translate(audio_file, target_language):
    # Transcribir audio con Whisper (aquí se usa la variable whisper_model)
    transcription = whisper_model.transcribe(audio_file, language="english")["text"]
    
    # Traducir texto a idioma seleccionado (aquí se usa translation_model y tokenizer)
    model_inputs = tokenizer(transcription, return_tensors="pt")
    generated_tokens = translation_model.generate(
        **model_inputs,
        forced_bos_token_id=tokenizer.lang_code_to_id[target_language]
    )
    translated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
    return translated_text.strip("[]' ")

## Interfaz de Gradio
# Está creado en filas para "organizar" la distribución de cada caja
with gr.Blocks(theme="Nymbo/Nymbo_Theme") as app:
    # Título
    gr.Markdown("## Transcripción y Traducción de Audio")
    
    # Primera fila -> Input de audio y elección del idioma
    with gr.Row():
        # Audio
        audio_input = gr.Audio(label="Subir o grabar audio en `inglés` exclusivamente`", sources=["upload", "microphone"], type="filepath")
        
        # Elección de idioma
        language_dropdown = gr.Dropdown(
            ["de_DE", "es_XX", "fr_XX", "sv_SE", "ru_RU"], 
            label="Selecciona el idioma de traducción",
            value="es_XX"
        )
    
    # Segunda fila -> Botón y salida de texto traducido
    with gr.Row():
        # Boton
        translate_button = gr.Button("Transcribir y Traducir")
        
        # Caja de texto (output)
        translation_output = gr.Textbox(label="Texto Traducido")
    
    # Configuración botón
    translate_button.click(
        transcribe_translate,
        inputs=[audio_input, language_dropdown],
        outputs=translation_output
    )

##Iniciar aplicacion
app.queue().launch()