import gradio as gr

from translation import Translator, LANGUAGES, MODEL_URL

LANGUAGES_LIST = list(LANGUAGES.keys())

def translate_wrapper(text, src, trg, by_sentence=True, preprocess=True, random=False, num_beams=4):
    src_lang = LANGUAGES.get(src)
    tgt_lang = LANGUAGES.get(trg)
    result = translator.translate(
        text=text,
        src_lang=src_lang,
        tgt_lang=tgt_lang,
        do_sample=random,
        num_beams=int(num_beams),
        by_sentence=by_sentence,
        preprocess=preprocess,
    )
    return result

article = f"""
This is the demo for a NLLB-200-600M model fine-tuned for a few (mostly new) languages.
The model itself is available at https://huggingface.co/{MODEL_URL}
If you want to host in on your own backend, consider running this dockerized app: https://github.com/slone-nlp/nllb-docker-demo.
"""

interface = gr.Interface(
    translate_wrapper,
    [
        gr.Textbox(label="Text to Translate", lines=2, placeholder='Enter text to translate'),
        gr.Dropdown(LANGUAGES_LIST, type="value", label='Source Language', value=LANGUAGES_LIST[0], description='Select the source language'),
        gr.Dropdown(LANGUAGES_LIST, type="value", label='Target Language', value=LANGUAGES_LIST[1], description='Select the target language'),
        gr.Checkbox(label="Translate by Sentence", value=True, description='If checked, the text will be translated sentence by sentence'),
        gr.Checkbox(label="Apply Text Preprocessing", value=True, description='If checked, the text will be preprocessed before translation'),
        gr.Checkbox(label="Randomize", value=False, description='If checked, the translation will use random sampling'),
        gr.Slider(minimum=1, maximum=5, step=1, label="Number of Beams", value=4, description='Select the number of beams for the translation'),
    ],
    "text",
    title='Erzya-Russian Translation',
    article=article,
)

if __name__ == '__main__':
    translator = Translator()
    interface.launch()