from typing import Dict, List, Any

from transformers import MBart50TokenizerFast, MBartForConditionalGeneration


class EndpointHandler:
    def __init__(self, path=""):
        self.tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
        self.model = MBartForConditionalGeneration.from_pretrained(path)

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        inputs = data.pop("inputs", data)
        source_lang = data.pop("src_lang", None)
        target_lang = data.pop("target_lang", None)

        self.tokenizer.src_lang = source_lang

        encoded = self.tokenizer(inputs, return_tensors='pt')
        generated_tokens = self.model.generate(
            **encoded, forced_bos_token_id=self.tokenizer.lang_code_to_id[target_lang]
        )
        d = self.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
        return d