import gradio as gr
import threading
from translation import Translator, LANGUAGES, MODEL_URL

translator = Translator()
lang_list = list(LANGUAGES.keys())

@gr.cache
def translate(text, src, trg, **kwargs):
    src_lang, tgt_lang = LANGUAGES.get(src), LANGUAGES.get(trg)
    return translator.translate(text, src_lang, tgt_lang, **kwargs)

interface = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(label="Text", lines=2, description="Enter the text you want to translate."),
        gr.Dropdown(lang_list, type="value", label='Source Language', value=lang_list[0], description="Select the language of the input text."),
        gr.Dropdown(lang_list, type="value", label='Target Language', value=lang_list[1], description="Select the language you want to translate to."),
        gr.Checkbox(label="By Sentence", value=True, description="Translate the text sentence by sentence."),
        gr.Checkbox(label="Preprocess", value=True, description="Preprocess the text before translation."),
        gr.Checkbox(label="Randomize", value=False, description="Randomize the translation output."),
        gr.Slider(value=4, label="Number of Beams", minimum=1, maximum=5, step=1, description="Select the number of beams for the translation process.")
    ],
    outputs="text",
    title='Erzya-Russian Translation',
    article=f"Model: https://huggingface.co/{MODEL_URL}"
)

if __name__ == '__main__':
    threading.Thread(target=interface.launch, kwargs=dict(share=True, server_name="0.0.0.0", server_port=7860)).start()