Spaces:
Sleeping
Sleeping
File size: 2,008 Bytes
10af882 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
from transformers import MarianMTModel, MarianTokenizer
from translations.model_name_mapping import MODEL_NAME_MAPPING
def load_model_and_tokenizer(from_lang: str, to_lang: str):
print(f"load_model_and_tokenizer from: {from_lang}, to: {to_lang}")
model_name = MODEL_NAME_MAPPING.get((from_lang, to_lang))
print(f"model_name: {model_name}")
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)
return model, tokenizer
def translate(text, model, tokenizer):
translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True))
translated_text = [
tokenizer.decode(t, skip_special_tokens=True) for t in translated
]
return translated_text[0]
def test_translations(from_lang, to_lang, text_to_translate):
model, tokenizer = load_model_and_tokenizer(from_lang=from_lang, to_lang=to_lang)
translation = translate(text_to_translate, model, tokenizer)
print(f"Translated text from: {from_lang}, to: {to_lang}, translation: {translation}")
if __name__ == "__main__":
# text_to_translate = "hola amigos, tengo hambre"
# test_translations(from_lang="es", to_lang="en", text_to_translate=text_to_translate)
# test_translations(from_lang="es", to_lang="ca", text_to_translate=text_to_translate)
# test_translations(from_lang="es", to_lang="po", text_to_translate=text_to_translate)
text_to_translate = "hello friends, who's hungry?"
test_translations(from_lang="en", to_lang="es", text_to_translate=text_to_translate)
# test_translations(from_lang="en", to_lang="ca", text_to_translate=text_to_translate)
test_translations(from_lang="en", to_lang="pt", text_to_translate=text_to_translate)
# test_translations(from_lang="es", to_lang="po", text_to_translate=text_to_translate)
text_to_translate = "un cafè sense sucre i amb llet, si us plau"
test_translations(from_lang="ca", to_lang="es", text_to_translate=text_to_translate) |