File size: 846 Bytes
b7a229d 405b2f1 b7a229d 405b2f1 b7a229d 405b2f1 b7a229d 405b2f1 b7a229d | 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 | import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from functions.utils import DEVICE, TORCH_DTYPE
MODEL_NAME = "code-li/nllb-moore"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(
MODEL_NAME,
torch_dtype=TORCH_DTYPE,
).to(DEVICE)
model.eval()
def translateFRMOS(text: str, sourceLang: str, targetLang: str):
tokenizer.src_lang = sourceLang
TGT_LANG_ID = tokenizer.convert_tokens_to_ids(targetLang)
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128).to(DEVICE)
with torch.no_grad():
tokens = model.generate(
**inputs,
forced_bos_token_id=TGT_LANG_ID,
max_length=128,
num_beams=4,
)
return tokenizer.decode(tokens[0], skip_special_tokens=True)
|