Snigdhapaul2003's picture
Update app.py
47a3422 verified
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)