new portuguese model

translation.py

@@ -1,21 +1,19 @@
 from transformers import MarianMTModel, MarianTokenizer
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
 def get_model_name(languageCode: str) -> str:
-    model_name = "Helsinki-NLP/opus-mt-pt-en"
+    match languageCode:
+        case "pt":
+            model_name = "Helsinki-NLP/opus-mt-pt-en"
+    
     return model_name
 
-def init_translation_model():
-    model_name = get_model_name("pt")
-    print(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    return model, tokenizer
-
 def get_translation(text_to_translate: str, languageCode: str) -> str:
-    model, tokenizer = init_translation_model()
-    inputs = tokenizer(text_to_translate, return_tensors="pt", truncation=True, padding=True)
-    translated_ids = model.generate(**inputs)
-    translated_text = tokenizer.decode(translated_ids[0], skip_special_tokens=True)
+        
+    tokenizer = AutoTokenizer.from_pretrained("unicamp-dl/translation-pt-en-t5")
+    model = AutoModelForSeq2SeqLM.from_pretrained("unicamp-dl/translation-pt-en-t5")
+    pten_pipeline = pipeline('text2text-generation', model=model, tokenizer=tokenizer)
+    translated_text = pten_pipeline(text_to_translate)
 
     print("Translated text:", translated_text)
     return translated_text