translation.py

from transformers import MarianMTModel, MarianTokenizer
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

def get_model_name(languageCode: str) -> str:
    match languageCode:
        case "pt":
            model_name = "Helsinki-NLP/opus-mt-pt-en"
    
    return model_name

def get_translation(text_to_translate: str, languageCode: str) -> str:
        
    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