Update app.py

app.py

@@ -26,6 +26,7 @@ models = ["Helsinki-NLP", "Helsinki-NLP/opus-mt-tc-bible-big-mul-mul",
           "bigscience/bloomz-560m", "bigscience/bloomz-1b1", "bigscience/bloomz-1b7", "bigscience/bloomz-3b",
           "t5-small", "t5-base", "t5-large",
           "google/flan-t5-small", "google/flan-t5-base", "google/flan-t5-large", "google/flan-t5-xl",
+          "google/madlad400-3b-mt", "jbochi/madlad400-3b-mt", 
           "Argos", "Google",
           "HuggingFaceTB/SmolLM3-3B", "winninghealth/WiNGPT-Babel-2",
           "utter-project/EuroLLM-1.7B", "utter-project/EuroLLM-1.7B-Instruct",
@@ -125,7 +126,20 @@ class Translators:
                 return self.HelsinkiNLP_mulmul()
         except KeyError as error:
             return f"Error: Translation direction {self.sl} to {self.tl} is not supported by Helsinki Translation Models", error
-        
+
+    def madlad(self):
+        model = T5ForConditionalGeneration.from_pretrained(self.model_name, device_map="auto")
+        tokenizer = T5Tokenizer.from_pretrained(self.model_name)
+        text = f"<2{self.tl}> {self.input_text}"
+        # input_ids = tokenizer(text, return_tensors="pt").input_ids.to(model.device)
+        # outputs = model.generate(input_ids=input_ids)    
+        # return tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Use a pipeline as a high-level helper
+        translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=self.sl, tgt_lang=self.tl)
+        translated_text = translator(text, max_length=512)
+        return translated_text[0]['translation_text']
+        # pipe = pipeline("translation", model=)
+    
     def smollm(self):
         tokenizer = AutoTokenizer.from_pretrained(self.model_name)
         model = AutoModelForCausalLM.from_pretrained(self.model_name)
@@ -402,6 +416,9 @@ def translate_text(input_text: str, s_language: str, t_language: str, model_name
             
         elif 'flan' in model_name.lower():
             translated_text = Translators(model_name, s_language, t_language, input_text).flan()
+
+        elif 'madlad' in model_name.lower():
+            translated_text = Translators(model_name, s_language, t_language, input_text).madlad()
             
         elif 'mt0' in model_name.lower():
             translated_text = Translators(model_name, s_language, t_language, input_text).bigscience()