Update app.py

app.py

@@ -2,12 +2,13 @@ import gradio as gr
 import spaces
 import torch
 from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+import languagecodes
 
-favourite_langs = {"German": "de", "Romanian": "ro", "English": "en", "French": "fr", "-----": "-----"}
-all_langs = {"German": "de", "Romanian": "ro", "English": "en", "French": "fr", "Spanish": "es", "Italian": "it", "Hungarian": "hu"}
-langs = {**favourite_langs, **all_langs}
+favourite_langs = {"German": "de", "Romanian": "ro", "English": "en", "-----": "-----"}
+langs = languagecodes.iso_languages
 
-options = list(langs.keys())
+# options = list(langs.keys())
+options = [(k, v) for k,v in favourite_langs.items()].extend([(k, v) for k,v in langs.items()])
 models = ["Helsinki-NLP", "t5-base", "t5-small", "t5-large", "facebook/nllb-200-distilled-600M", "facebook/nllb-200-distilled-1.3B", "facebook/mbart-large-50-many-to-many-mmt"]
 
 def model_to_cuda(model):
@@ -37,24 +38,22 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
                 return f"Error finding model: {model_name}! Try other available language combination.", error
     
     if model_name.startswith('facebook/nllb'):
-        from languagecodes import nllb_language_codes
-        tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang=nllb_language_codes[sselected_language])
+        tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang=languagecodes.nllb_language_codes[sselected_language])
         model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto")
-        translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=nllb_language_codes[sselected_language], tgt_lang=nllb_language_codes[tselected_language])
+        translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=languagecodes.nllb_language_codes[sselected_language], tgt_lang=languagecodes.nllb_language_codes[tselected_language])
         translated_text = translator(input_text, max_length=512)
         return translated_text[0]['translation_text'], message_text
     
-    if model_name.startswith('facebook/mbart-large'):
-        from languagecodes import mbart_large_languages
+    if model_name.startswith('facebook/mbart-large
         from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
         model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
         tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
         # translate source to target
-        tokenizer.src_lang = mbart_large_languages[sselected_language]
+        tokenizer.src_lang = languagecodes.mbart_large_languages[sselected_language]
         encoded = tokenizer(input_text, return_tensors="pt")
         generated_tokens = model.generate(
             **encoded,
-            forced_bos_token_id=tokenizer.lang_code_to_id[mbart_large_languages[tselected_language]]
+            forced_bos_token_id=tokenizer.lang_code_to_id[languagecodes.mbart_large_languages[tselected_language]]
         )
         return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0], message_text