import torch
import spaces
from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
import os
from huggingface_hub import login

max_length = 512
auth_token = os.getenv('HF_SPACE_TOKEN')
login(token=auth_token)


@spaces.GPU
def goai_traduction(text, src_lang, tgt_lang):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if src_lang == "fra_Latn" and tgt_lang == "mos_Latn":
        model_id = "ArissBandoss/nllb-200-distilled-600M-finetuned-fr-to-mos-V4"
        
    elif src_lang == "mos_Latn" and tgt_lang == "fra_Latn":
        model_id = "ArissBandoss/nllb-200-distilled-600M-finetuned-mos-to-fr-V5"

    else:
        model_id = "ArissBandoss/nllb-200-distilled-600M-finetuned-fr-to-mos-V4"

    tokenizer = AutoTokenizer.from_pretrained(model_id, token=auth_token)
    model     = AutoModelForSeq2SeqLM.from_pretrained(model_id, token=auth_token)
        
    trans_pipe = pipeline("translation", 
                          model=model, tokenizer=tokenizer, 
                          src_lang=src_lang, tgt_lang=tgt_lang, 
                          max_length=max_length,
                          device=device
                         )
    
    return trans_pipe(text)[0]["translation_text"]