translation.py

from transformers import MarianMTModel, MarianTokenizer

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

def init_translation_model():
    model_name = get_model_name("pt")
    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)

    print("Translated text:", translated_text)
    return translated_text