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

en2ko = "KoJLabs/nllb-finetuned-en2ko"
ko2en = "KoJLabs/nllb-finetuned-ko2en"
style = "KoJLabs/bart-speech-style-converter"

en2ko_model = AutoModelForSeq2SeqLM.from_pretrained(en2ko)
ko2en_model = AutoModelForSeq2SeqLM.from_pretrained(ko2en)
style_model = AutoModelForSeq2SeqLM.from_pretrained(style)

en2ko_tokenizer = AutoTokenizer.from_pretrained(en2ko)
ko2en_tokenizer = AutoTokenizer.from_pretrained(ko2en)
style_tokenizer = AutoTokenizer.from_pretrained(style)

def translate(source, target, text):
    formats = {"English":"eng_Latn", "Korean":"kor_Hang"}
    src = formats[source]
    tgt = formats[target]

    if src == "eng_Latn":
        translator = pipeline(
            'translation',
            model=en2ko_model,
            tokenizer=ko2en_tokenizer,
            src_lang=src, 
            tgt_lang=tgt, 
        )

    if src == "kor_Hang":
        translator = pipeline(
            'translation',
            model=ko2en_model,
            tokenizer=en2ko_tokenizer,
            src_lang=src, 
            tgt_lang=tgt
        )

    output = translator(text)
    translated_text = output[0]['translation_text']

    return translated_text

title = 'KoTAN Translator & Speech-style converter'
lang = ['English','Korean']

translator_app = gr.Interface(
    fn=translate,
    inputs=[gr.inputs.Dropdown(choices=lang, label='Source Language'), gr.inputs.Dropdown(choices=lang, label='Target Language'), gr.inputs.Textbox(lines=5, label='Text to Translate')],
    outputs=[gr.outputs.Textbox(label='Translated Text')],
    title=title,
    description = 'KoTAN: Korean Translation and Augmentation with fine-tuned NLLB. If you want to download as pip package, please visit our github. (https://github.com/KoJLabs/KoTAN)',
    article='Jisu, Kim. Juhwan, Lee',
    enable_queue=True,
)

translator_app.launch()