| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
| import torch | |
| from functions.utils import DEVICE, TORCH_DTYPE | |
| MODEL_NAME = "facebook/nllb-200-distilled-600M" # facebook/nllb-200-3.3B | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
| model = AutoModelForSeq2SeqLM.from_pretrained( | |
| MODEL_NAME, | |
| torch_dtype=TORCH_DTYPE, | |
| ).to(DEVICE) | |
| model.eval() | |
| def translateText(text: str, sourceLang: str, targetLang: str) -> str: | |
| tokenizer.src_lang = sourceLang | |
| inputs = tokenizer(text, return_tensors="pt").to(DEVICE) | |
| with torch.no_grad(): | |
| tokens = model.generate( | |
| **inputs, | |
| forced_bos_token_id=tokenizer.convert_tokens_to_ids(targetLang), | |
| max_new_tokens=512, | |
| num_beams=4, | |
| early_stopping=True, | |
| ) | |
| translatedText = tokenizer.batch_decode(tokens, skip_special_tokens=True)[0] | |
| return translatedText | |