Update app.py

app.py

@@ -64,6 +64,49 @@ def argos(sl, tl, input_text):
         print(error)
     return translated_text
 
+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: # Standard bilingual model
+        model_name = f"Helsinki-NLP/opus-mt-{sl}-{tl}"
+        pipe = pipeline("translation", model=model_name, device=-1)
+        translation = pipe(input_text)
+        return translation[0]['translation_text'], f'Translated from {sl} to {tl} with {model_name}.'
+    except EnvironmentError:
+        try: # Tatoeba models
+            model_name = f"Helsinki-NLP/opus-tatoeba-{sl}-{tl}"
+            pipe = pipeline("translation", model=model_name, device=-1)
+            translation = pipe(input_text)
+            return translation[0]['translation_text'], f'Translated from {sl} to {tl} with {model_name}.'
+        except EnvironmentError as error:
+            try: # Last resort: multi to multi
+                model_name = "Helsinki-NLP/opus-mt-tc-bible-big-mul-mul"
+                pipe = pipeline("translation", model=model_name)
+                tl = 'deu' # Hard coded for now for testing
+                translation = pipe(f'>>{tl}<< {input_text}')
+                return translation[0]['translation_text'], f'Translated from {sl} to {tl} with {model_name}.'
+            except Exception as error:
+                return f"Error translating with model: {model_name}! Try other available language combination.", error
+    except KeyError as error:
+        return f"Error: Translation direction {sl} to {tl} is not supported by Helsinki Translation Models", error
+        
 class Translators:
     def __init__(self, model_name: str, sl: str, tl: str, input_text: str):
         self.model_name = model_name
@@ -109,57 +152,75 @@ class Translators:
         translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
         return translated_text
     
-def mtom(model_name, sl, tl, input_text):
-    from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
-    model = M2M100ForConditionalGeneration.from_pretrained(model_name)
-    tokenizer = M2M100Tokenizer.from_pretrained(model_name)
-    tokenizer.src_lang = sl
-    encoded = tokenizer(input_text, return_tensors="pt")
-    generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id(tl))
-    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+    def mbart_many_to_many(self):
+        from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+        model = MBartForConditionalGeneration.from_pretrained(self.model_name)
+        tokenizer = MBart50TokenizerFast.from_pretrained(self.model_name)
+        # translate source to target
+        tokenizer.src_lang = languagecodes.mbart_large_languages[self.sl]
+        encoded = tokenizer(self.input_text, return_tensors="pt")
+        generated_tokens = model.generate(
+            **encoded,
+            forced_bos_token_id=tokenizer.lang_code_to_id[languagecodes.mbart_large_languages[self.tl]]
+        )
+        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+    
+    def mbart_one_to_many(self):
+        # translate from English
+        from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+        model = MBartForConditionalGeneration.from_pretrained(self.model_name)
+        tokenizer = MBart50TokenizerFast.from_pretrained(self.model_name, src_lang="en_XX")
+        model_inputs = tokenizer(self.input_text, return_tensors="pt")
+        langid = languagecodes.mbart_large_languages[self.tl]
+        generated_tokens = model.generate(
+            **model_inputs,
+            forced_bos_token_id=tokenizer.lang_code_to_id[langid]
+        )
+        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+    
+    def mbart_many_to_one(self):
+        # translate to English
+        from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+        model = MBartForConditionalGeneration.from_pretrained(self.model_name)
+        tokenizer = MBart50TokenizerFast.from_pretrained(self.model_name)
+        tokenizer.src_lang = languagecodes.mbart_large_languages[self.sl]
+        encoded = tokenizer(self.input_text, return_tensors="pt")
+        generated_tokens = model.generate(**encoded)
+        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+        
+    def mtom(self):
+        from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
+        model = M2M100ForConditionalGeneration.from_pretrained(self.model_name)
+        tokenizer = M2M100Tokenizer.from_pretrained(self.model_name)
+        tokenizer.src_lang = self.sl
+        encoded = tokenizer(self.input_text, return_tensors="pt")
+        generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id(self.tl))
+        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
 
-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: # Standard bilingual model
-        model_name = f"Helsinki-NLP/opus-mt-{sl}-{tl}"
-        pipe = pipeline("translation", model=model_name, device=-1)
-        translation = pipe(input_text)
-        return translation[0]['translation_text'], f'Translated from {sl} to {tl} with {model_name}.'
-    except EnvironmentError:
-        try: # Tatoeba models
-            model_name = f"Helsinki-NLP/opus-tatoeba-{sl}-{tl}"
-            pipe = pipeline("translation", model=model_name, device=-1)
-            translation = pipe(input_text)
-            return translation[0]['translation_text'], f'Translated from {sl} to {tl} with {model_name}.'
-        except EnvironmentError as error:
-            try: # Last resort: multi to multi
-                model_name = "Helsinki-NLP/opus-mt-tc-bible-big-mul-mul"
-                pipe = pipeline("translation", model=model_name)
-                tl = 'deu' # Hard coded for now for testing
-                translation = pipe(f'>>{tl}<< {input_text}')
-                return translation[0]['translation_text'], f'Translated from {sl} to {tl} with {model_name}.'
-            except Exception as error:
-                return f"Error translating with model: {model_name}! Try other available language combination.", error
-    except KeyError as error:
-        return f"Error: Translation direction {sl} to {tl} is not supported by Helsinki Translation Models", error     
+    def bigscience(self):  
+        tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name)
+        inputs = tokenizer.encode(f"Translate to {self.tl}: {self.input_text}.", return_tensors="pt")
+        outputs = model.generate(inputs)
+        translation = tokenizer.decode(outputs[0])
+        translation = translation.replace('<pad> ', '').replace('</s>', '')
+        return translation
+    
+    def bloomz(self):  
+        tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        model = AutoModelForCausalLM.from_pretrained(self.model_name)
+        inputs = tokenizer.encode(f"Translate from {self.sl} to {self.tl}: {self.input_text}. Translation:", return_tensors="pt")
+        outputs = model.generate(inputs)
+        translation = tokenizer.decode(outputs[0])
+        translation = translation.replace('<pad> ', '').replace('</s>', '')
+        return translation
+    
+    def nllb(self):
+        tokenizer = AutoTokenizer.from_pretrained(self.model_name, src_lang=self.sl)
+        model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name, device_map="auto")
+        translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=self.sl, tgt_lang=self.tl)
+        translated_text = translator(self.input_text, max_length=512)
+        return translated_text[0]['translation_text']
 
 def teuken(model_name, sl, tl, input_text):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -189,24 +250,6 @@ def teuken(model_name, sl, tl, input_text):
     translation = tokenizer.decode(prediction[0].tolist())
     return translation
 
-def bigscience(model_name, sl, tl, input_text):  
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-    inputs = tokenizer.encode(f"Translate to {tl}: {input_text}.", return_tensors="pt")
-    outputs = model.generate(inputs)
-    translation = tokenizer.decode(outputs[0])
-    translation = translation.replace('<pad> ', '').replace('</s>', '')
-    return translation
-
-def bloomz(model_name, sl, tl, input_text):  
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    inputs = tokenizer.encode(f"Translate from {sl} to {tl}: {input_text}. Translation:", return_tensors="pt")
-    outputs = model.generate(inputs)
-    translation = tokenizer.decode(outputs[0])
-    translation = translation.replace('<pad> ', '').replace('</s>', '')
-    return translation
-
 def eurollm(model_name, sl, tl, input_text):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)  
@@ -228,13 +271,6 @@ def eurollm_instruct(model_name, sl, tl, input_text):
         output = output.rsplit(f'{tl}:')[-1].strip().replace('assistant\n', '')
     return output
 
-def nllb(model_name, sl, tl, input_text):
-    tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang=sl)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto")
-    translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=sl, tgt_lang=tl)
-    translated_text = translator(input_text, max_length=512)
-    return translated_text[0]['translation_text']
-
 def unbabel(model_name, sl, tl, input_text):
     pipe = pipeline("text-generation", model=model_name, torch_dtype=torch.bfloat16, device_map="auto")
     messages = [{"role": "user",
@@ -256,43 +292,6 @@ def unbabel(model_name, sl, tl, input_text):
     split_translated_text = translated_text.split('\n', translated_text.count('\n'))
     translated_text = '\n'.join(split_translated_text[:input_text.count('\n')+1])
     return translated_text
-    
-def mbart_many_to_many(model_name, sl, tl, input_text):
-    from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
-    model = MBartForConditionalGeneration.from_pretrained(model_name)
-    tokenizer = MBart50TokenizerFast.from_pretrained(model_name)
-    # translate source to target
-    tokenizer.src_lang = languagecodes.mbart_large_languages[sl]
-    encoded = tokenizer(input_text, return_tensors="pt")
-    generated_tokens = model.generate(
-        **encoded,
-        forced_bos_token_id=tokenizer.lang_code_to_id[languagecodes.mbart_large_languages[tl]]
-    )
-    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
-
-def mbart_one_to_many(model_name, sl, tl, input_text):
-    from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
-    article_en = input_text
-    model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-one-to-many-mmt")
-    tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-one-to-many-mmt", src_lang="en_XX")
-    model_inputs = tokenizer(article_en, return_tensors="pt")
-    # translate from English
-    langid = languagecodes.mbart_large_languages[tl]
-    generated_tokens = model.generate(
-        **model_inputs,
-        forced_bos_token_id=tokenizer.lang_code_to_id[langid]
-    )
-    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
-
-def mbart_many_to_one(model_name, sl, tl, input_text):
-    from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
-    model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-one-mmt")
-    tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-one-mmt")
-    # translate to English
-    tokenizer.src_lang = languagecodes.mbart_large_languages[sl]
-    encoded = tokenizer(input_text, return_tensors="pt")
-    generated_tokens = model.generate(**encoded)
-    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
 
 @spaces.GPU
 def translate_text(input_text: str, sselected_language: str, tselected_language: str, model_name: str) -> tuple[str, str]:
@@ -330,42 +329,42 @@ def translate_text(input_text: str, sselected_language: str, tselected_language:
             translated_text = Translators(model_name, sl, tl, input_text).google()
     
         elif "m2m" in model_name.lower():
-            translated_text = mtom(model_name, 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)
+            translated_text = Translators(model_name, sl, tl, input_text).mtom()
         
-        elif model_name == "utter-project/EuroLLM-1.7B":
-            translated_text = eurollm(model_name, sselected_language, tselected_language, input_text)
-                
         elif model_name.startswith('t5'):
             translated_text = Translators(model_name, sselected_language, tselected_language, input_text).tfive()
             
         elif 'flan' in model_name.lower():
             translated_text = Translators(model_name, sselected_language, tselected_language, input_text).flan()
             
-        elif 'teuken' in model_name.lower():
-            translated_text = teuken(model_name, sselected_language, tselected_language, input_text)
-    
         elif 'mt0' in model_name.lower():
-            translated_text = bigscience(model_name, sselected_language, tselected_language, input_text)
+            translated_text = Translators(model_name, sselected_language, tselected_language, input_text).bigscience()
     
         elif 'bloomz' in model_name.lower():
-            translated_text = bloomz(model_name, sselected_language, tselected_language, input_text)
+            translated_text = Translators(model_name, sselected_language, tselected_language, input_text).bloomz()
             
         elif 'nllb' in model_name.lower():
             nnlbsl, nnlbtl = languagecodes.nllb_language_codes[sselected_language], languagecodes.nllb_language_codes[tselected_language]
-            translated_text = nllb(model_name, nnlbsl, nnlbtl, input_text)
+            translated_text = Translators(model_name, nnlbsl, nnlbtl, input_text).nllb()
         
         elif model_name == "facebook/mbart-large-50-many-to-many-mmt":
-            translated_text = mbart_many_to_many(model_name, sselected_language, tselected_language, input_text)
+            translated_text = Translators(model_name, sselected_language, tselected_language, input_text).mbart_many_to_many()
     
         elif model_name == "facebook/mbart-large-50-one-to-many-mmt":
-            translated_text = mbart_one_to_many(model_name, sselected_language, tselected_language, input_text)
+            translated_text = Translators(model_name, sselected_language, tselected_language, input_text).mbart_one_to_many()
     
         elif model_name == "facebook/mbart-large-50-many-to-one-mmt":
-            translated_text = mbart_many_to_one(model_name, sselected_language, tselected_language, input_text)
+            translated_text = Translators(model_name, sselected_language, tselected_language, input_text).mbart_many_to_one()
         
+        elif 'teuken' in model_name.lower():
+            translated_text = teuken(model_name, sselected_language, tselected_language, input_text)
+    
+        elif model_name == "utter-project/EuroLLM-1.7B-Instruct":
+            translated_text = eurollm_instruct(model_name, sselected_language, tselected_language, input_text)
+        
+        elif model_name == "utter-project/EuroLLM-1.7B":
+            translated_text = eurollm(model_name, sselected_language, tselected_language, input_text)
+                
         elif 'Unbabel' in model_name:   
             translated_text = unbabel(model_name, sselected_language, tselected_language, input_text)