Update app.py

app.py

@@ -7,6 +7,14 @@ langs = {"German": "de", "Romanian": "ro", "English": "en", "French": "fr", "Spa
 options = list(langs.keys())
 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):
+    # Move the model to GPU if available
+    if torch.cuda.is_available():
+        model = model.to('cuda')
+        print("CUDA is available! Using GPU.")
+    else:
+        print("CUDA not available! Using CPU.")
+    return model
 @spaces.GPU
 def translate_text(input_text, sselected_language, tselected_language, model_name):
     sl = langs[sselected_language]
@@ -16,14 +24,15 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
         try:
             model_name_full = f"Helsinki-NLP/opus-mt-{sl}-{tl}"
             tokenizer = AutoTokenizer.from_pretrained(model_name_full)
-            model = AutoModelForSeq2SeqLM.from_pretrained(model_name_full)
+            model = model_to_cuda(AutoModelForSeq2SeqLM.from_pretrained(model_name_full))
         except EnvironmentError:
             try :   
                 model_name_full = f"Helsinki-NLP/opus-tatoeba-{sl}-{tl}"
                 tokenizer = AutoTokenizer.from_pretrained(model_name_full)
-                model = AutoModelForSeq2SeqLM.from_pretrained(model_name_full)
+                model = model_to_cuda(AutoModelForSeq2SeqLM.from_pretrained(model_name_full))
             except EnvironmentError as error:
                 return f"Error finding model: {model_name_full}! 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])
@@ -31,6 +40,7 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
         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
+    
     if model_name.startswith('facebook/mbart-large'):
         from languagecodes import mbart_large_languages
         from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
@@ -43,18 +53,12 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
             **encoded,
             forced_bos_token_id=tokenizer.lang_code_to_id[mbart_large_languages[tselected_language]]
         )
-        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True), message_text
+        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0], message_text
 
     if model_name.startswith('t5'):
         tokenizer = T5Tokenizer.from_pretrained(model_name)
         model = T5ForConditionalGeneration.from_pretrained(model_name, device_map="auto")
 
-    # Move the model to GPU if available
-    if torch.cuda.is_available():
-        model = model.to('cuda')
-        print("CUDA is available! Using GPU.")
-    else:
-        print("CUDA not available! Using CPU.")
     if model_name.startswith("Helsinki-NLP"):
         prompt = input_text
     else: