Update app.py

app.py

@@ -2,14 +2,14 @@ import gradio as gr
 import spaces
 from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-langs = {"German": ("de", "deu_Latn"), "Romanian": ("ro", "ron_Latn"), "English": ("en", "eng_Latn"), "French": ("fr", "fra_Latn"), "Spanish": ("es", "spa_Latn"), "Italian": ("it", "ita_Latn")}
+langs = {"German": "de", "Romanian": "ro", "English": "en", "French": "fr", "Spanish": "es", "Italian": "it"}
 options = list(langs.keys())
 models = ["Helsinki-NLP", "t5-base", "t5-small", "t5-large", "facebook/nllb-200-distilled-600M", "facebook/nllb-200-distilled-1.3B"]
 
 @spaces.GPU
 def translate_text(input_text, sselected_language, tselected_language, model_name):
-    sl = langs[sselected_language][0]
-    tl = langs[tselected_language][0]
+    sl = langs[sselected_language]
+    tl = langs[tselected_language]
     message_text = f'Translated from {sselected_language} to {tselected_language} with {model_name}'
     if model_name == "Helsinki-NLP":
         try:
@@ -24,9 +24,10 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
             except EnvironmentError as error:
                 return f"Error finding model: {model_name_full}! Try other available language combination.", error
     elif model_name.startswith('facebook/nllb'):
+        from languagecode import nllb_language_codes
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-        translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=langs[sselected_language][1], tgt_lang=langs[tselected_language][1])
+        translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=nllb_language_codes[sselected_language], tgt_lang=nllb_language_codes[tselected_language])
         translated_text = translator(input_text, max_length=512)
         return translated_text[0]['translation_text'], message_text
     else: