from transformers import pipeline
import torch
import gradio as gr

# Translation pipeline
translator = pipeline(task="translation",
                      model="facebook/nllb-200-1.3B",
                      torch_dtype=torch.bfloat16)

# list of EU languages and their FLoRes-200 code
eu_languages = {
'Bulgarian':'bul_Cyrl',
'Croatian':'hrv_Latn',
'Czech':'ces_Latn',
'Danish':'dan_Latn',
'Dutch':'nld_Latn',
'English':'eng_Latn',
'Estonian':'est_Latn',
'Finnish':'fin_Latn',
'French':'fra_Latn',
'German':'deu_Latn',
'Greek':'ell_Grek',
'Hungarian':'hun_Latn',
'Irish':'gle_Latn',
'Italian':'ita_Latn',
'Latvian':'lvs_Latn',
'Lithuanian':'lit_Latn',
'Maltese':'mlt_Latn',
'Polish':'pol_Latn',
'Portuguese':'por_Latn',
'Romanian':'ron_Latn',
'Slovak':'slk_Latn',
'Slovenian':'slv_Latn',
'Spanish':'spa_Latn',
'Swedish':'swe_Latn'
}

# Translate function
def translate(input, src, tgt):
    src_lang = eu_languages[src]
    tgt_lang = eu_languages[tgt]
    output = translator(input, src_lang=src_lang, tgt_lang=tgt_lang, max_length=400)
    return output[0]['translation_text']


# Gradio Interface
gr.close_all()
demo = gr.Interface(fn=translate, 
                    inputs=[gr.Textbox(label="Text to translate", lines=6),
                            gr.Dropdown(eu_languages.keys(), label="Source Language"),
                            gr.Dropdown(eu_languages.keys(), label="Target Language")],
                    outputs=[gr.Textbox(label="Result", lines=10)],
                    examples=[["Jokainen on oman onnensa seppä.", "Finnish","English"]],
                    title="NLLB Translator between EU Languages",
                    description="Translate texts in EU languages using the `facebook/nllb-200-1.3B` model!"
                   )
demo.launch()