# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/08_translation.ipynb (unless otherwise specified).

__all__ = ['logger', 'TranslationResult', 'TransformersTranslator', 'EasyTranslator']

# Cell
import logging
from typing import List, Dict, Union
from collections import defaultdict, OrderedDict

import torch
from torch.utils.data import TensorDataset, DataLoader

from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    PreTrainedTokenizer,
    PreTrainedModel,
    T5ForConditionalGeneration,
)

from ..model import AdaptiveModel
from ..callback import GeneratorCallback

from fastai.torch_core import apply, to_device

from fastcore.basics import Self
from ..model_hub import HFModelResult, FlairModelResult, HFModelHub, FlairModelHub
from ..result import DetailLevel

# Cell
logger = logging.getLogger(__name__)

# Cell
class TranslationResult:
    "A basic result class for Translation problems"
    def __init__(
        self,
        inputs:List[str], # A list of input string sentences
        input_lang:str, # A input language
        output_lang:str, # An output language
        translations:List[str] # A list of the translated sentences
    ):
        self.inputs = inputs
        self.input_lang = input_lang
        self.output_lang = output_lang
        self.translations = translations

    def to_dict(
        self,
        detail_level:DetailLevel=DetailLevel.Low # A detail level to return
    ):
        "Convert `self` to a filtered dictionary"
        o = OrderedDict()
        o.update(
            {'translations':self.translations,}
        )
        if detail_level == 'medium' or detail_level == 'high':
            o.update({
                'inputs': self.inputs,
                'input_language': self.input_lang,
                'output_language': self.output_lang
            })
        if detail_level == 'high':
            print("Warning: There is no difference between `medium` and `high` for Translation tasks")
        return o

# Cell
class TransformersTranslator(AdaptiveModel):
    "Adaptive model for Transformer's Conditional Generation or Language Models (Transformer's T5 and Bart conditional generation models have a language modeling head)"

    def __init__(
        self,
        tokenizer: PreTrainedTokenizer, # A tokenizer object from Huggingface's transformers (TODO)and tokenizers
        model: PreTrainedModel # A transformers Conditional Generation (Bart or T5) or Language model
    ):
        # Load up model and tokenizer
        self.tokenizer = tokenizer
        super().__init__()

        # Sets internal model
        self.set_model(model)

        # Set inputs to come in as `dict`
        super().set_as_dict(True)

    @classmethod
    def load(
        cls,
        model_name_or_path: str # A key string of one of Transformer's pre-trained translator Model
    ) -> AdaptiveModel:
        "Class method for loading and constructing this classifier"
        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
        translator = cls(tokenizer, model)
        return translator

    def predict(
        self,
        text: Union[List[str], str], # Sentences to run inference on
        t5_prefix: str = 'translate English to German', # The pre-appended prefix for the specificied task. Only in use for T5-type models.
        mini_batch_size: int = 32,  # Mini batch size
        num_beams: int = 1, # Number of beams for beam search. Must be between 1 and infinity. 1 means no beam search
        min_length: int = 0, # The min length of the sequence to be generated
        max_length: int = 128, # The max length of the sequence to be generated. Between min_length and infinity
        early_stopping: bool = True, # If set to `True` beam search is stopped when at least num_beams sentences finished per batch
        detail_level:DetailLevel = DetailLevel.Low, # The level of detail to return
        **kwargs, # Optional arguments for the Transformers `PreTrainedModel.generate()` method
    ) -> List[str]: # A list of translated sentences
        "Predict method for running inference using the pre-trained sequence classifier model. Keyword arguments for parameters of the method `Transformers.PreTrainedModel.generate()` can be used as well"

        # Make all inputs lists
        if isinstance(text, str):
            text = [text]

        # T5 requires 'translate: ' precursor text for pre-trained translator
        if isinstance(self.model, T5ForConditionalGeneration):
            text = [f'{t5_prefix}: {t}' for t in text]

        dataset = self._tokenize(text)
        dl = DataLoader(dataset, batch_size=mini_batch_size)
        translations = []

        logger.info(f'Running translator on {len(dataset)} text sequences')
        logger.info(f'Batch size = {mini_batch_size}')

        cb = GeneratorCallback(num_beams, min_length, max_length, early_stopping, **kwargs)

        preds,_ = super().get_preds(dl=dl, cbs=[cb])

        preds = apply(Self.squeeze(0), preds)

        for o in preds:
            translations.append(
                [
                    self.tokenizer.decode(
                        o,
                        skip_special_tokens=True,
                        clean_up_tokenization_spaces=False,
                    )
                ].pop()
            )

        languages = t5_prefix.strip('translate ').split(' to ')

        res = TranslationResult(text, *languages, translations)

        return res if detail_level is None else res.to_dict(detail_level)

    def _tokenize(self, text: Union[List[str], str]) -> TensorDataset:
        """ Batch tokenizes text and produces a `TensorDataset` with text """

        tokenized_text = self.tokenizer.batch_encode_plus(
            text,
            return_tensors='pt',
            max_length=512,
            padding='max_length',
            add_special_tokens=True,
        )

        # Bart doesn't use `token_type_ids`
        dataset = TensorDataset(
            tokenized_text['input_ids'],
            tokenized_text['attention_mask'],
        )
        return dataset

# Cell
class EasyTranslator:
    "Translation Module"

    def __init__(self):
        self.translators: Dict[AdaptiveModel] = defaultdict(bool)

    def translate(
        self,
        text: Union[List[str], str], # Sentences to run inference on
        model_name_or_path: str = "t5-small", # A model id or path to a pre-trained model repository or custom trained model directory
        t5_prefix: str = "translate English to German", # The pre-appended prefix for the specificied task. Only in use for T5-type models
        detail_level=DetailLevel.Low, # The level of detail to return
        mini_batch_size: int = 32, # Mini batch size
        num_beams: int = 1, # Number of beams for beam search. Must be between 1 and infinity. 1 means no beam search
        min_length: int = 0, # The min length of the sequence to be generated
        max_length: int = 128, # The max length of the sequence to be generated. Between min_length and infinity
        early_stopping: bool = True, # If set to `True` beam search is stopped when at least num_beams sentences finished per batch
        **kwargs, # Optional arguments for the Transformers `PreTrainedModel.generate()` method
    ) -> List[str]:
        "Predict method for running inference using the pre-trained sequence classifier model. Keyword arguments for parameters of the method `Transformers.PreTrainedModel.generate()` can be used as well."
        name = getattr(model_name_or_path, 'name', model_name_or_path)
        if not self.translators[name]:
            self.translators[name] = TransformersTranslator.load(
                name
            )

        translator = self.translators[name]
        return translator.predict(
            text=text,
            t5_prefix=t5_prefix,
            mini_batch_size=mini_batch_size,
            num_beams=num_beams,
            min_length=min_length,
            max_length=max_length,
            early_stopping=early_stopping,
            **kwargs,
        )