Update app.py

app.py

@@ -43,6 +43,42 @@ def model_to_cuda(model):
         print("CUDA not available! Using CPU.")
     return model
 
+def HelsinkiNLPAutoTokenizer(sl, tl, input_text):
+    if model_name == "Helsinki-NLP":
+        message_text = f'Translated from {sl} to {tl} with {model_name}.'
+        try:
+            model_name = f"Helsinki-NLP/opus-mt-{sl}-{tl}"
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            model = model_to_cuda(AutoModelForSeq2SeqLM.from_pretrained(model_name))
+        except EnvironmentError:
+            try:   
+                model_name = f"Helsinki-NLP/opus-tatoeba-{sl}-{tl}"
+                tokenizer = AutoTokenizer.from_pretrained(model_name)
+                model = model_to_cuda(AutoModelForSeq2SeqLM.from_pretrained(model_name))
+                input_ids = tokenizer.encode(prompt, return_tensors="pt")
+                output_ids = model.generate(input_ids, max_length=512)
+                translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+                return translated_text, message_text
+            except EnvironmentError as error:
+                return f"Error finding model: {model_name}! Try other available language combination.", error
+                
+def HelsinkiNLP(sl, tl, input_text):
+    try:
+        model_name = f"Helsinki-NLP/opus-mt-{sl}-{tl}"
+        pipe = pipeline("translation", model=model_name)
+        # translation = pipe(input_text)
+        # return translation[0]['translation_text'], f'Translated from {sl} to {tl} with {model_name}.'
+    except EnvironmentError:
+        try:   
+            model_name = f"Helsinki-NLP/opus-tatoeba-{sl}-{tl}"
+            pipe = pipeline("translation", model=model_name)
+            translation = pipe(input_text)
+            return translation[0]['translation_text'], f'Translated from {sl} to {tl} with {model_name}.'
+        except EnvironmentError as error:
+            return f"Error finding model: {model_name}! Try other available language combination.", error
+    except KeyError as error:
+        return f"Error: Translation direction {source_readable} to {target} is not supported by Helsinki Translation Models", error     
+
 def flan(model_name, sl, tl, input_text):
     tokenizer = T5Tokenizer.from_pretrained(model_name, legacy=False)
     model = T5ForConditionalGeneration.from_pretrained(model_name)
@@ -174,21 +210,12 @@ def translate_text(input_text: str, sselected_language: str, tselected_language:
     
     sl = all_langs[sselected_language]
     tl = all_langs[tselected_language]
-    message_text = f'Translating from {sselected_language} to {tselected_language} with {model_name}'
+    message_text = f'Translated from {sselected_language} to {tselected_language} with {model_name}'
     print(message_text)
-    if model_name == "Helsinki-NLP":
-        try:
-            model_name = f"Helsinki-NLP/opus-mt-{sl}-{tl}"
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model = model_to_cuda(AutoModelForSeq2SeqLM.from_pretrained(model_name))
-        except EnvironmentError:
-            try:   
-                model_name = f"Helsinki-NLP/opus-tatoeba-{sl}-{tl}"
-                tokenizer = AutoTokenizer.from_pretrained(model_name)
-                model = model_to_cuda(AutoModelForSeq2SeqLM.from_pretrained(model_name))
-            except EnvironmentError as error:
-                return f"Error finding model: {model_name}! Try other available language combination.", error
     
+    if model_name.startswith("Helsinki-NLP"):
+        return translated_text, message_text = HelsinkiNLP(sl, tl, input_text)
+        
     elif model_name == "utter-project/EuroLLM-1.7B-Instruct":
         translated_text = eurollm_instruct(model_name, sselected_language, tselected_language, input_text)
         return translated_text, message_text
@@ -226,22 +253,18 @@ def translate_text(input_text: str, sselected_language: str, tselected_language:
         translated_text = unbabel(model_name, sselected_language, tselected_language, input_text)
         return translated_text, message_text
     
-    if model_name.startswith('t5'):
+    elif model_name.startswith('t5'):
         tokenizer = T5Tokenizer.from_pretrained(model_name)
         model = T5ForConditionalGeneration.from_pretrained(model_name, device_map="auto")
-
-    if model_name.startswith("Helsinki-NLP"):
-        prompt = input_text
-    else:
         prompt = f"translate {sselected_language} to {tselected_language}: {input_text}"
+        input_ids = tokenizer.encode(prompt, return_tensors="pt")
+        output_ids = model.generate(input_ids, max_length=512)
+        translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True) 
+        # print(f'Translated from {sselected_language} to {tselected_language} with {model_name}:', f'{input_text} =  {translated_text}', sep='\n')
+        return translated_text, message_text
 
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    output_ids = model.generate(input_ids, max_length=512)
-    translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    
-    message_text = f'Translated from {sselected_language} to {tselected_language} with {model_name}.'
-    print(f'Translated from {sselected_language} to {tselected_language} with {model_name}:', f'{input_text} =  {translated_text}', sep='\n')
-    return translated_text, message_text
+    else:
+        return "No model selected or error in application" message_text
 
 # Define a function to swap dropdown values
 def swap_languages(src_lang, tgt_lang):