import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

languages = {
    'Afrikaans': 'afr_Latn',
    'Albanian': 'als_Latn',
    'Basque': 'eus_Latn',
    'Belarusian': 'bel_Cyrl',
    'Bosnian': 'bos_Latn',
    'Bulgarian': 'bul_Cyrl',
    'Catalan': 'cat_Latn',
    'Croatian': 'hrv_Latn',
    'Czech': 'ces_Latn',
    'Danish': 'dan_Latn',
    'Dutch': 'nld_Latn',
    'English': 'eng_Latn',
    'Estonian': 'est_Latn',
    'Finnish': 'fin_Latn',
    'French': 'fra_Latn',
    'Galician': 'glg_Latn',
    'German': 'deu_Latn',
    'Greek': 'ell_Grek',
    'Hungarian': 'hun_Latn',
    'Icelandic': 'isl_Latn',
    'Irish': 'gle_Latn',
    'Italian': 'ita_Latn',
    'Lithuanian': 'lit_Latn',
    'Luxembourgish': 'ltz_Latn',
    'Macedonian': 'mkd_Cyrl',
    'Maltese': 'mlt_Latn',
    'Norwegian Bokmål': 'nob_Latn',
    'Norwegian Nynorsk': 'nno_Latn',
    'Polish': 'pol_Latn',
    'Portuguese': 'por_Latn',
    'Romanian': 'ron_Latn',
    'Russian': 'rus_Cyrl',
    'Serbian': 'srp_Cyrl',
    'Slovak': 'slk_Latn',
    'Slovenian': 'slv_Latn',
    'Spanish': 'spa_Latn',
    'Swedish': 'swe_Latn',
    'Ukrainian': 'ukr_Cyrl',
    'Welsh': 'cym_Latn'
}

tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")

def translate_article(article, language):
    inputs = tokenizer(article, return_tensors="pt")
    if language=='English':
      lang_code='eng_Latn'
      print("Yes")
    else:
      lang_code = languages[language]
      print("No")
    print(lang_code)
    translated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id[lang_code])
    result = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
    return result

language_choices = list(languages.keys())

iface = gr.Interface(
    fn=translate_article,
    inputs=["text",gr.Dropdown(
            language_choices, value="English", multiselect=False, label="Choose the language.")],
    outputs="text",
    title="Translation Tool"
)
iface.launch(debug=True)