Update app.py

app.py

@@ -89,8 +89,26 @@ class Translators:
             do_sample=True
         ) 
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        print(response)
         return response.split("Translation:")[-1].strip()
 
+    def flan(self):
+        tokenizer = T5Tokenizer.from_pretrained(self.model_name, legacy=False)
+        model = T5ForConditionalGeneration.from_pretrained(self.model_name)
+        prompt = f"translate {self.sl} to {self.tl}: {self.input_text}"
+        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+        outputs = model.generate(input_ids)
+        return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+
+    def tfive(self):
+        tokenizer = T5Tokenizer.from_pretrained(self.model_name)
+        model = T5ForConditionalGeneration.from_pretrained(self.model_name, device_map="auto")
+        prompt = f"translate {self.sl} to {self.tl}: {self.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).strip()
+        return translated_text
+    
 def mtom(model_name, sl, tl, input_text):
     from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
     model = M2M100ForConditionalGeneration.from_pretrained(model_name)
@@ -143,23 +161,6 @@ def HelsinkiNLP(sl, tl, input_text):
     except KeyError as error:
         return f"Error: Translation direction {sl} to {tl} 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)
-    input_text = f"translate {sl} to {tl}: {input_text}"
-    input_ids = tokenizer(input_text, return_tensors="pt").input_ids
-    outputs = model.generate(input_ids)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
-
-def tfive(model_name, sl, tl, input_text):
-    tokenizer = T5Tokenizer.from_pretrained(model_name)
-    model = T5ForConditionalGeneration.from_pretrained(model_name, device_map="auto")
-    prompt = f"translate {sl} to {tl}: {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)
-    return translated_text
-
 def teuken(model_name, sl, tl, input_text):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = AutoModelForCausalLM.from_pretrained(
@@ -336,9 +337,12 @@ def translate_text(input_text: str, sselected_language: str, tselected_language:
         
         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 = flan(model_name, sselected_language, tselected_language, input_text)
+            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)
@@ -364,9 +368,6 @@ def translate_text(input_text: str, sselected_language: str, tselected_language:
         
         elif 'Unbabel' in model_name:   
             translated_text = unbabel(model_name, sselected_language, tselected_language, input_text)
-        
-        elif model_name.startswith('t5'):
-            translated_text = tfive(model_name, sselected_language, tselected_language, input_text)
             
         elif model_name == "HuggingFaceTB/SmolLM3-3B":
             translated_text = Translators(model_name, sselected_language, tselected_language, input_text).smollm()