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)