<gh_stars>10-100
from abc import ABC, abstractmethod
from argparse import Namespace
from typing import Any, Dict, List, Optional, Tuple, Union, cast

import pytorch_lightning as pl
import torch
import torch.nn as nn
from loguru import logger
from omegaconf import DictConfig
from slp.config.omegaconf import OmegaConf
from slp.util.pytorch import pad_mask, subsequent_mask
from slp.util.system import print_separator
from slp.util.types import Configuration, LossType
from torch.optim import Optimizer
from torch.optim.lr_scheduler import _LRScheduler


class _Predictor(ABC):
    """Base predictor class

    Define an interface that can be used to extend the lightning module to new tasks and models

    * parse_batch: Parse input batch and extract necessery masks etc.
    * get_predictions_and_targets: Perform a forward pass through the model to get the logits and return logits and targets
    """

    @abstractmethod
    def parse_batch(self, batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        """Abstract parse_batch method to be implemented by child class

        Args:
            batch (Tuple[torch.Tensor, ...]): A tuple of tensors that contains inputs to the model and targets

        Returns:
            Tuple[torch.Tensor, ...]: The processed inputs, ready to provide to the model
        """
        raise NotImplementedError

    @abstractmethod
    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Abstract get_predictions and targets method to be implemented by child class

        This method gets exposed to the PLModule classes

        Args:
            model (nn.Module): model to use for forward pass
            batch (Tuple[torch.Tensor, torch.Tensor]): A tuple of tensors that contains inputs to the model and targets

        **Note**: Maybe it should be useful to move loss calculation here. Then multitask learning and auxiliary losses should be easier

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, ground_truths), ready to be passed to the loss function
        """
        raise NotImplementedError


class _Classification(_Predictor):
    """Classification task"""

    def parse_batch(self, batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        """Parse incoming batch

        Input batch just contains inputs and targets

        Args:
            batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, labels)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (inputs, labels)
        """
        inputs = batch[0]
        targets = batch[1]

        return inputs, targets

    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return logits and ground truths to be passed in loss function

        Args:
            model (nn.Module): Model to use for prediction
            batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, labels)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, labels)
        """
        inputs, targets = self.parse_batch(batch)
        y_pred = model(inputs)

        return y_pred.squeeze(), targets.squeeze()


class _AutoEncoder(_Predictor):
    """Autoencoder task"""

    def parse_batch(self, batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        """Parse incoming batch

        Input batch just contains inputs. Targets are the same as inputs, because we are doing reconstruction.

        Args:
            batch (Tuple[torch.Tensor]): (inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (inputs, inputs)
        """
        inputs = batch[0]

        return inputs, inputs

    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return logits and ground truths to be passed in loss function

        Args:
            model (nn.Module): Model to use for prediction
            batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, inputs)
        """
        inputs, targets = self.parse_batch(batch)
        y_pred = model(inputs)

        return y_pred.view(y_pred.size(0), -1), targets.view(targets.size(0), -1)


class _RnnClassification(_Predictor):
    """RNN classification task"""

    def parse_batch(self, batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        """Parse incoming batch

        Input batch just contains inputs, targets and lengths.
        Comes from slp.data.collators.SequentialCollator.

        Args:
            batch (Tuple[torch.Tensor]): (inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (inputs, inputs)
        """
        inputs = batch[0]
        targets = batch[1]
        lengths = batch[2]

        return inputs, targets, lengths

    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return logits and ground truths to be passed in loss function

        Args:
            model (nn.Module): Model to use for prediction
            batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, inputs)
        """
        inputs, targets, lengths = self.parse_batch(batch)
        y_pred = model(inputs, lengths)

        return y_pred.squeeze(), targets.squeeze()


class _TransformerClassification(_Predictor):
    """Transformer classification task"""

    def parse_batch(self, batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        """Parse incoming batch

        Input batch just contains inputs, targets and lengths.
        Comes from slp.data.collators.SequentialCollator.
        Create pad masks to be passed to transformer attention

        Args:
            batch (Tuple[torch.Tensor]): (inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (inputs, inputs)
        """
        inputs = batch[0]
        targets = batch[1]
        lengths = batch[2]
        attention_mask = pad_mask(lengths, max_length=inputs.size(1))

        return inputs, targets, attention_mask

    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return logits and ground truths to be passed in loss function

        Args:
            model (nn.Module): Model to use for prediction
            batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, inputs)
        """
        inputs, targets, attention_mask = self.parse_batch(batch)
        y_pred = model(inputs, attention_mask=attention_mask)

        return y_pred.squeeze(), targets.squeeze()


class _MultimodalTransformerClassification(_Predictor):
    """Transformer classification task"""

    def parse_batch(self, batch):
        """Parse incoming batch

        Input batch just contains inputs, targets and lengths.
        Comes from slp.data.collators.SequentialCollator.
        Create pad masks to be passed to transformer attention

        Args:
            batch (Tuple[Dict[str, torch.Tensor], torch.Tensor, Dict[str, torch.Tensor]]): (inputs, targets, lengths)

        Returns:
            Tuple[Dict[str, torch.Tensor], torch.Tensor, Dict[str, torch.Tensor]]: (inputs, targets, attention_masks)
        """
        inputs = batch[0]
        targets = batch[1]
        lengths = batch[2]
        attention_masks = {
            m: pad_mask(lengths[m], max_length=inputs[m].size(1))
            for m in lengths.keys()
        }

        return inputs, targets, attention_masks

    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return logits and ground truths to be passed in loss function

        Args:
            model (nn.Module): Model to use for prediction
            batch (Tuple[Dict[str, torch.Tensor], torch.Tensor, Dict[str, torch.Tensor]]): (inputs, targets, lengths)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, targets)
        """
        inputs, targets, attention_masks = self.parse_batch(batch)
        y_pred = model(inputs, attention_masks=attention_masks)

        return y_pred.squeeze(), targets.squeeze()


class _Transformer(_Predictor):
    """Generic transformer seq2seq task"""

    def parse_batch(self, batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        """Parse incoming batch

        Input batch just contains inputs, targets and lengths.
        Comes from slp.data.collators.SequentialCollator.
        Create pad masks and subsequent_masks to be passed to transformer attention

        Args:
            batch (Tuple[torch.Tensor]): (inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (inputs, targets)
        """
        inputs = batch[0]
        targets = batch[1]
        lengths_inputs = batch[2]
        lengths_targets = batch[3]

        pad_inputs = pad_mask(
            lengths_inputs,
            max_length=inputs.size(1),
        )
        pad_targets = pad_mask(
            lengths_targets,
            max_length=targets.size(1),
        )
        sub_m = subsequent_mask(targets.size(1))  # type: ignore
        pad_targets = pad_targets.unsqueeze(-2) * sub_m.to(pad_targets.device)

        return inputs, targets, pad_inputs, pad_targets

    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return logits and ground truths to be passed in loss function

        Args:
            model (nn.Module): Model to use for prediction
            batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, inputs)
        """
        inputs, targets, source_mask, target_mask = self.parse_batch(batch)
        y_pred = model(
            inputs, targets, source_mask=source_mask, target_mask=target_mask
        )

        y_pred = y_pred.view(-1, y_pred.size(-1))
        targets = targets.view(-1)

        return y_pred, targets


class _BertSequenceClassification(_Predictor):
    """Bert Classification task"""

    def parse_batch(self, batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        """Parse incoming batch

        Input batch just contains inputs, targets and lengths.
        Comes from slp.data.collators.SequentialCollator.
        Create pad masks to be passed to BERT attention

        Args:
            batch (Tuple[torch.Tensor]): (inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (inputs, inputs)
        """
        inputs = batch[0]
        targets = batch[1]
        lengths = batch[2]

        attention_mask = pad_mask(lengths, inputs.size(1))

        return inputs, targets, attention_mask

    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return logits and ground truths to be passed in loss function

        Args:
            model (nn.Module): Model to use for prediction
            batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, inputs)
        """
        inputs, targets, attention_mask = self.parse_batch(batch)
        out = model(
            input_ids=inputs,
            attention_mask=attention_mask,
            labels=None,
            return_dict=False,
        )
        y_pred = out[0].view(-1, out[0].size(-1))
        targets = targets.view(-1)

        return y_pred.squeeze(), targets.squeeze()


class SimplePLModule(pl.LightningModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        predictor_cls=_Classification,
        calculate_perplexity: bool = False,  # for LM. Dirty but much more efficient
    ):
        """Wraps a (model, optimizer, criterion, lr_scheduler) tuple in a LightningModule

        Handles the boilerplate for metrics calculation and logging and defines the train_step / val_step / test_step
        with use of the predictor helper classes (e.g. _Classification, _RnnClassification)

        Args:
            model (nn.Module): Module to use for prediction
            optimizer (Union[Optimizer, List[Optimizer]]): Optimizers to use for training
            criterion (LossType): Task loss
            lr_scheduler (Union[_LRScheduler, List[_LRScheduler]], optional): Learning rate scheduler. Defaults to None.
            hparams (Configuration, optional): Hyperparameter values. This ensures they are logged with trainer.loggers. Defaults to None.
            metrics (Optional[Dict[str, pl.metrics.Metric]], optional): Metrics to track. Defaults to None.
            predictor_cls ([type], optional): Class that defines a parse_batch and a
                    get_predictions_and_targets method. Defaults to _Classification.
            calculate_perplexity (bool, optional): Whether to calculate perplexity.
                    Would be cleaner as a metric, but this is more efficient. Defaults to False.
        """
        super(SimplePLModule, self).__init__()
        self.calculate_perplexity = calculate_perplexity
        self.model = model
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler
        self.criterion = criterion

        if metrics is not None:
            self.train_metrics = nn.ModuleDict(metrics)
            self.val_metrics = nn.ModuleDict({k: v.clone() for k, v in metrics.items()})
            self.test_metrics = nn.ModuleDict(
                {k: v.clone() for k, v in metrics.items()}
            )
        else:
            self.train_metrics = nn.ModuleDict(modules=None)
            self.val_metrics = nn.ModuleDict(modules=None)
            self.test_metrics = nn.ModuleDict(modules=None)
        self.predictor = predictor_cls()

        if hparams is not None:
            if isinstance(hparams, Namespace):
                dict_params = vars(hparams)
            elif isinstance(hparams, DictConfig):
                dict_params = cast(Dict[str, Any], OmegaConf.to_container(hparams))
            else:
                dict_params = hparams
            # self.hparams = dict_params
            self.save_hyperparameters(dict_params)

    def configure_optimizers(self):
        """Return optimizers and learning rate schedulers

        Returns:
            Tuple[List[Optimizer], List[_LRScheduler]]: (optimizers, lr_schedulers)
        """

        if self.lr_scheduler is not None:
            scheduler = {
                "scheduler": self.lr_scheduler,
                "interval": "epoch",
                "monitor": "val_loss",
            }

            return [self.optimizer], [scheduler]

        return self.optimizer

    def forward(self, *args, **kwargs):
        """Call wrapped module forward"""

        return self.model(*args, **kwargs)

    def _compute_metrics(self, metrics, loss, y_hat, targets, mode="train"):
        """Compute all metrics and aggregate in a dict

        Args:
            metrics (Dict[str, pl.metrics.Metric]): metrics to compute
            loss (torch.Tensor): Computed loss
            y_hat (torch.Tensor): Logits
            targets (torch.Tensor): Ground Truths
            mode (str, optional): "train", "val" or "test". Defaults to "train".
        """

        def fmt(name):
            """Format metric name"""

            return f"{mode}_{name}"

        metrics = {f"{mode}_{k}": v(y_hat, targets) for k, v in metrics.items()}

        if mode == "train":
            metrics["loss"] = loss
        else:
            metrics[fmt("loss")] = loss

        if self.calculate_perplexity:
            metrics[fmt("ppl")] = torch.exp(loss)

        return metrics

    def log_to_console(self, metrics, mode="Training"):
        """Log metrics to console

        Args:
            metrics (Dict[str, torch.Tensor]): Computed metrics
            mode (str, optional): "Training", "Validation" or "Testing". Defaults to "Training".
        """
        logger.info("Epoch {} {} results".format(self.current_epoch + 1, mode))
        print_separator(symbol="-", n=50, print_fn=logger.info)

        for name, value in metrics.items():
            if name == "epoch":
                continue
            logger.info("{:<15} {:<15}".format(name, value))

        print_separator(symbol="%", n=50, print_fn=logger.info)

    def aggregate_epoch_metrics(self, outputs, mode="Training"):
        """Aggregate metrics over a whole epoch

        Args:
            outputs (List[Dict[str, torch.Tensor]]): Aggregated outputs from train_step, validation_step or test_step
            mode (str, optional): "Training", "Validation" or "Testing". Defaults to "Training".
        """

        def fmt(name):
            """Format metric name"""

            return f"{name}" if name != "loss" else "train_loss"

        keys = list(outputs[0].keys())
        aggregated = {fmt(k): torch.stack([x[k] for x in outputs]).mean() for k in keys}
        aggregated["epoch"] = self.current_epoch + 1

        self.log_dict(aggregated, logger=True, prog_bar=False, on_epoch=True)

        return aggregated

    def training_step(self, batch, batch_idx):
        """Compute loss for a single training step and log metrics to loggers

        Args:
            batch (Tuple[torch.Tensor, ...]): Input batch
            batch_idx (int): Index of batch

        Returns:
            Dict[str, torch.Tensor]: computed metrics
        """
        y_hat, targets = self.predictor.get_predictions_and_targets(self.model, batch)
        loss = self.criterion(y_hat, targets)
        metrics = self._compute_metrics(
            self.train_metrics, loss, y_hat, targets, mode="train"
        )

        self.log_dict(
            metrics,
            on_step=True,
            on_epoch=False,
            logger=True,
            prog_bar=False,
        )

        metrics["loss"] = loss

        return metrics

    def training_epoch_end(self, outputs):
        """Aggregate metrics of a training epoch

        Args:
            outputs (List[Dict[str, torch.Tensor]]): Aggregated outputs from train_step
        """
        outputs = self.aggregate_epoch_metrics(outputs, mode="Training")
        self.log_to_console(outputs, mode="Training")

    def validation_step(self, batch, batch_idx):
        """Compute loss for a single validation step and log metrics to loggers

        Args:
            batch (Tuple[torch.Tensor, ...]): Input batch
            batch_idx (int): Index of batch

        Returns:
            Dict[str, torch.Tensor]: computed metrics
        """
        y_hat, targets = self.predictor.get_predictions_and_targets(self, batch)
        loss = self.criterion(y_hat, targets)
        metrics = self._compute_metrics(
            self.val_metrics, loss, y_hat, targets, mode="val"
        )

        metrics[
            "best_score"
        ] = self.trainer.early_stopping_callback.best_score.detach().cpu()

        return metrics

    def validation_epoch_end(self, outputs):
        """Aggregate metrics of a validation epoch

        Args:
            outputs (List[Dict[str, torch.Tensor]]): Aggregated outputs from validation_step
        """
        outputs = self.aggregate_epoch_metrics(outputs, mode="Validation")

        if torch.isnan(outputs["val_loss"]) or torch.isinf(outputs["val_loss"]):
            outputs["val_loss"] = 1000000

        outputs["best_score"] = min(
            outputs[self.trainer.early_stopping_callback.monitor].detach().cpu(),
            self.trainer.early_stopping_callback.best_score.detach().cpu(),
        )
        self.log_to_console(outputs, mode="Validation")

    def test_step(self, batch, batch_idx):
        """Compute loss for a single test step and log metrics to loggers

        Args:
            batch (Tuple[torch.Tensor, ...]): Input batch
            batch_idx (int): Index of batch

        Returns:
            Dict[str, torch.Tensor]: computed metrics
        """
        y_hat, targets = self.predictor.get_predictions_and_targets(self, batch)
        loss = self.criterion(y_hat, targets)
        metrics = self._compute_metrics(
            self.test_metrics, loss, y_hat, targets, mode="test"
        )

        return metrics

    def test_epoch_end(self, outputs):
        """Aggregate metrics of a test epoch

        Args:
            outputs (List[Dict[str, torch.Tensor]]): Aggregated outputs from test_step
        """
        outputs = self.aggregate_epoch_metrics(outputs, mode="Test")
        self.log_to_console(outputs, mode="Test")


class PLModule(SimplePLModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        calculate_perplexity=False,
    ):
        """Pass arguments through to base class"""
        super(PLModule, self).__init__(
            model,
            optimizer,
            criterion,
            predictor_cls=_Classification,
            lr_scheduler=lr_scheduler,
            hparams=hparams,
            metrics=metrics,
            calculate_perplexity=calculate_perplexity,
        )


class AutoEncoderPLModule(SimplePLModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        calculate_perplexity=False,
    ):
        """Pass arguments through to base class"""
        super(AutoEncoderPLModule, self).__init__(
            model,
            optimizer,
            criterion,
            predictor_cls=_AutoEncoder,
            lr_scheduler=lr_scheduler,
            hparams=hparams,
            metrics=metrics,
            calculate_perplexity=calculate_perplexity,
        )


class RnnPLModule(SimplePLModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        calculate_perplexity=False,
    ):
        """Pass arguments through to base class"""
        super(RnnPLModule, self).__init__(
            model,
            optimizer,
            criterion,
            predictor_cls=_RnnClassification,
            lr_scheduler=lr_scheduler,
            hparams=hparams,
            metrics=metrics,
            calculate_perplexity=calculate_perplexity,
        )


class TransformerClassificationPLModule(SimplePLModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        calculate_perplexity=False,
    ):
        """Pass arguments through to base class"""
        super(TransformerClassificationPLModule, self).__init__(
            model,
            optimizer,
            criterion,
            predictor_cls=_TransformerClassification,
            lr_scheduler=lr_scheduler,
            hparams=hparams,
            metrics=metrics,
            calculate_perplexity=calculate_perplexity,
        )


class TransformerPLModule(SimplePLModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        calculate_perplexity=False,
    ):
        """Pass arguments through to base class"""
        super(TransformerPLModule, self).__init__(
            model,
            optimizer,
            criterion,
            predictor_cls=_Transformer,
            lr_scheduler=lr_scheduler,
            hparams=hparams,
            metrics=metrics,
            calculate_perplexity=calculate_perplexity,
        )


class BertPLModule(SimplePLModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        calculate_perplexity=False,
    ):
        """Pass arguments through to base class"""
        super(BertPLModule, self).__init__(
            model,
            optimizer,
            criterion,
            predictor_cls=_BertSequenceClassification,
            lr_scheduler=lr_scheduler,
            hparams=hparams,
            metrics=metrics,
            calculate_perplexity=calculate_perplexity,
        )


class MultimodalTransformerClassificationPLModule(SimplePLModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        calculate_perplexity=False,
    ):
        """Pass arguments through to base class"""
        super(MultimodalTransformerClassificationPLModule, self).__init__(
            model,
            optimizer,
            criterion,
            predictor_cls=_MultimodalTransformerClassification,
            lr_scheduler=lr_scheduler,
            hparams=hparams,
            metrics=metrics,
            calculate_perplexity=calculate_perplexity,
        )

from abc import ABC, abstractmethod
from argparse import Namespace
from typing import Any, Dict, List, Optional, Tuple, Union, cast

import pytorch_lightning as pl
import torch
import torch.nn as nn
from loguru import logger
from omegaconf import DictConfig
from slp.config.omegaconf import OmegaConf
from slp.util.pytorch import pad_mask, subsequent_mask
from slp.util.system import print_separator
from slp.util.types import Configuration, LossType
from torch.optim import Optimizer
from torch.optim.lr_scheduler import _LRScheduler


class _Predictor(ABC):
    """Base predictor class

    Define an interface that can be used to extend the lightning module to new tasks and models

    * parse_batch: Parse input batch and extract necessery masks etc.
    * get_predictions_and_targets: Perform a forward pass through the model to get the logits and return logits and targets
    """

    @abstractmethod
    def parse_batch(self, batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        """Abstract parse_batch method to be implemented by child class

        Args:
            batch (Tuple[torch.Tensor, ...]): A tuple of tensors that contains inputs to the model and targets

        Returns:
            Tuple[torch.Tensor, ...]: The processed inputs, ready to provide to the model
        """
        raise NotImplementedError

    @abstractmethod
    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Abstract get_predictions and targets method to be implemented by child class

        This method gets exposed to the PLModule classes

        Args:
            model (nn.Module): model to use for forward pass
            batch (Tuple[torch.Tensor, torch.Tensor]): A tuple of tensors that contains inputs to the model and targets

        **Note**: Maybe it should be useful to move loss calculation here. Then multitask learning and auxiliary losses should be easier

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, ground_truths), ready to be passed to the loss function
        """
        raise NotImplementedError


class _Classification(_Predictor):
    """Classification task"""

    def parse_batch(self, batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        """Parse incoming batch

        Input batch just contains inputs and targets

        Args:
            batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, labels)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (inputs, labels)
        """
        inputs = batch[0]
        targets = batch[1]

        return inputs, targets

    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return logits and ground truths to be passed in loss function

        Args:
            model (nn.Module): Model to use for prediction
            batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, labels)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, labels)
        """
        inputs, targets = self.parse_batch(batch)
        y_pred = model(inputs)

        return y_pred.squeeze(), targets.squeeze()


class _AutoEncoder(_Predictor):
    """Autoencoder task"""

    def parse_batch(self, batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        """Parse incoming batch

        Input batch just contains inputs. Targets are the same as inputs, because we are doing reconstruction.

        Args:
            batch (Tuple[torch.Tensor]): (inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (inputs, inputs)
        """
        inputs = batch[0]

        return inputs, inputs

    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return logits and ground truths to be passed in loss function

        Args:
            model (nn.Module): Model to use for prediction
            batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, inputs)
        """
        inputs, targets = self.parse_batch(batch)
        y_pred = model(inputs)

        return y_pred.view(y_pred.size(0), -1), targets.view(targets.size(0), -1)


class _RnnClassification(_Predictor):
    """RNN classification task"""

    def parse_batch(self, batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        """Parse incoming batch

        Input batch just contains inputs, targets and lengths.
        Comes from slp.data.collators.SequentialCollator.

        Args:
            batch (Tuple[torch.Tensor]): (inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (inputs, inputs)
        """
        inputs = batch[0]
        targets = batch[1]
        lengths = batch[2]

        return inputs, targets, lengths

    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return logits and ground truths to be passed in loss function

        Args:
            model (nn.Module): Model to use for prediction
            batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, inputs)
        """
        inputs, targets, lengths = self.parse_batch(batch)
        y_pred = model(inputs, lengths)

        return y_pred.squeeze(), targets.squeeze()


class _TransformerClassification(_Predictor):
    """Transformer classification task"""

    def parse_batch(self, batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        """Parse incoming batch

        Input batch just contains inputs, targets and lengths.
        Comes from slp.data.collators.SequentialCollator.
        Create pad masks to be passed to transformer attention

        Args:
            batch (Tuple[torch.Tensor]): (inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (inputs, inputs)
        """
        inputs = batch[0]
        targets = batch[1]
        lengths = batch[2]
        attention_mask = pad_mask(lengths, max_length=inputs.size(1))

        return inputs, targets, attention_mask

    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return logits and ground truths to be passed in loss function

        Args:
            model (nn.Module): Model to use for prediction
            batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, inputs)
        """
        inputs, targets, attention_mask = self.parse_batch(batch)
        y_pred = model(inputs, attention_mask=attention_mask)

        return y_pred.squeeze(), targets.squeeze()


class _MultimodalTransformerClassification(_Predictor):
    """Transformer classification task"""

    def parse_batch(self, batch):
        """Parse incoming batch

        Input batch just contains inputs, targets and lengths.
        Comes from slp.data.collators.SequentialCollator.
        Create pad masks to be passed to transformer attention

        Args:
            batch (Tuple[Dict[str, torch.Tensor], torch.Tensor, Dict[str, torch.Tensor]]): (inputs, targets, lengths)

        Returns:
            Tuple[Dict[str, torch.Tensor], torch.Tensor, Dict[str, torch.Tensor]]: (inputs, targets, attention_masks)
        """
        inputs = batch[0]
        targets = batch[1]
        lengths = batch[2]
        attention_masks = {
            m: pad_mask(lengths[m], max_length=inputs[m].size(1))
            for m in lengths.keys()
        }

        return inputs, targets, attention_masks

    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return logits and ground truths to be passed in loss function

        Args:
            model (nn.Module): Model to use for prediction
            batch (Tuple[Dict[str, torch.Tensor], torch.Tensor, Dict[str, torch.Tensor]]): (inputs, targets, lengths)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, targets)
        """
        inputs, targets, attention_masks = self.parse_batch(batch)
        y_pred = model(inputs, attention_masks=attention_masks)

        return y_pred.squeeze(), targets.squeeze()


class _Transformer(_Predictor):
    """Generic transformer seq2seq task"""

    def parse_batch(self, batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        """Parse incoming batch

        Input batch just contains inputs, targets and lengths.
        Comes from slp.data.collators.SequentialCollator.
        Create pad masks and subsequent_masks to be passed to transformer attention

        Args:
            batch (Tuple[torch.Tensor]): (inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (inputs, targets)
        """
        inputs = batch[0]
        targets = batch[1]
        lengths_inputs = batch[2]
        lengths_targets = batch[3]

        pad_inputs = pad_mask(
            lengths_inputs,
            max_length=inputs.size(1),
        )
        pad_targets = pad_mask(
            lengths_targets,
            max_length=targets.size(1),
        )
        sub_m = subsequent_mask(targets.size(1))  # type: ignore
        pad_targets = pad_targets.unsqueeze(-2) * sub_m.to(pad_targets.device)

        return inputs, targets, pad_inputs, pad_targets

    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return logits and ground truths to be passed in loss function

        Args:
            model (nn.Module): Model to use for prediction
            batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, inputs)
        """
        inputs, targets, source_mask, target_mask = self.parse_batch(batch)
        y_pred = model(
            inputs, targets, source_mask=source_mask, target_mask=target_mask
        )

        y_pred = y_pred.view(-1, y_pred.size(-1))
        targets = targets.view(-1)

        return y_pred, targets


class _BertSequenceClassification(_Predictor):
    """Bert Classification task"""

    def parse_batch(self, batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        """Parse incoming batch

        Input batch just contains inputs, targets and lengths.
        Comes from slp.data.collators.SequentialCollator.
        Create pad masks to be passed to BERT attention

        Args:
            batch (Tuple[torch.Tensor]): (inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (inputs, inputs)
        """
        inputs = batch[0]
        targets = batch[1]
        lengths = batch[2]

        attention_mask = pad_mask(lengths, inputs.size(1))

        return inputs, targets, attention_mask

    def get_predictions_and_targets(
        self, model: nn.Module, batch: Tuple[torch.Tensor, ...]
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return logits and ground truths to be passed in loss function

        Args:
            model (nn.Module): Model to use for prediction
            batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: (logits, inputs)
        """
        inputs, targets, attention_mask = self.parse_batch(batch)
        out = model(
            input_ids=inputs,
            attention_mask=attention_mask,
            labels=None,
            return_dict=False,
        )
        y_pred = out[0].view(-1, out[0].size(-1))
        targets = targets.view(-1)

        return y_pred.squeeze(), targets.squeeze()


class SimplePLModule(pl.LightningModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        predictor_cls=_Classification,
        calculate_perplexity: bool = False,  # for LM. Dirty but much more efficient
    ):
        """Wraps a (model, optimizer, criterion, lr_scheduler) tuple in a LightningModule

        Handles the boilerplate for metrics calculation and logging and defines the train_step / val_step / test_step
        with use of the predictor helper classes (e.g. _Classification, _RnnClassification)

        Args:
            model (nn.Module): Module to use for prediction
            optimizer (Union[Optimizer, List[Optimizer]]): Optimizers to use for training
            criterion (LossType): Task loss
            lr_scheduler (Union[_LRScheduler, List[_LRScheduler]], optional): Learning rate scheduler. Defaults to None.
            hparams (Configuration, optional): Hyperparameter values. This ensures they are logged with trainer.loggers. Defaults to None.
            metrics (Optional[Dict[str, pl.metrics.Metric]], optional): Metrics to track. Defaults to None.
            predictor_cls ([type], optional): Class that defines a parse_batch and a
                    get_predictions_and_targets method. Defaults to _Classification.
            calculate_perplexity (bool, optional): Whether to calculate perplexity.
                    Would be cleaner as a metric, but this is more efficient. Defaults to False.
        """
        super(SimplePLModule, self).__init__()
        self.calculate_perplexity = calculate_perplexity
        self.model = model
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler
        self.criterion = criterion

        if metrics is not None:
            self.train_metrics = nn.ModuleDict(metrics)
            self.val_metrics = nn.ModuleDict({k: v.clone() for k, v in metrics.items()})
            self.test_metrics = nn.ModuleDict(
                {k: v.clone() for k, v in metrics.items()}
            )
        else:
            self.train_metrics = nn.ModuleDict(modules=None)
            self.val_metrics = nn.ModuleDict(modules=None)
            self.test_metrics = nn.ModuleDict(modules=None)
        self.predictor = predictor_cls()

        if hparams is not None:
            if isinstance(hparams, Namespace):
                dict_params = vars(hparams)
            elif isinstance(hparams, DictConfig):
                dict_params = cast(Dict[str, Any], OmegaConf.to_container(hparams))
            else:
                dict_params = hparams
            # self.hparams = dict_params
            self.save_hyperparameters(dict_params)

    def configure_optimizers(self):
        """Return optimizers and learning rate schedulers

        Returns:
            Tuple[List[Optimizer], List[_LRScheduler]]: (optimizers, lr_schedulers)
        """

        if self.lr_scheduler is not None:
            scheduler = {
                "scheduler": self.lr_scheduler,
                "interval": "epoch",
                "monitor": "val_loss",
            }

            return [self.optimizer], [scheduler]

        return self.optimizer

    def forward(self, *args, **kwargs):
        """Call wrapped module forward"""

        return self.model(*args, **kwargs)

    def _compute_metrics(self, metrics, loss, y_hat, targets, mode="train"):
        """Compute all metrics and aggregate in a dict

        Args:
            metrics (Dict[str, pl.metrics.Metric]): metrics to compute
            loss (torch.Tensor): Computed loss
            y_hat (torch.Tensor): Logits
            targets (torch.Tensor): Ground Truths
            mode (str, optional): "train", "val" or "test". Defaults to "train".
        """

        def fmt(name):
            """Format metric name"""

            return f"{mode}_{name}"

        metrics = {f"{mode}_{k}": v(y_hat, targets) for k, v in metrics.items()}

        if mode == "train":
            metrics["loss"] = loss
        else:
            metrics[fmt("loss")] = loss

        if self.calculate_perplexity:
            metrics[fmt("ppl")] = torch.exp(loss)

        return metrics

    def log_to_console(self, metrics, mode="Training"):
        """Log metrics to console

        Args:
            metrics (Dict[str, torch.Tensor]): Computed metrics
            mode (str, optional): "Training", "Validation" or "Testing". Defaults to "Training".
        """
        logger.info("Epoch {} {} results".format(self.current_epoch + 1, mode))
        print_separator(symbol="-", n=50, print_fn=logger.info)

        for name, value in metrics.items():
            if name == "epoch":
                continue
            logger.info("{:<15} {:<15}".format(name, value))

        print_separator(symbol="%", n=50, print_fn=logger.info)

    def aggregate_epoch_metrics(self, outputs, mode="Training"):
        """Aggregate metrics over a whole epoch

        Args:
            outputs (List[Dict[str, torch.Tensor]]): Aggregated outputs from train_step, validation_step or test_step
            mode (str, optional): "Training", "Validation" or "Testing". Defaults to "Training".
        """

        def fmt(name):
            """Format metric name"""

            return f"{name}" if name != "loss" else "train_loss"

        keys = list(outputs[0].keys())
        aggregated = {fmt(k): torch.stack([x[k] for x in outputs]).mean() for k in keys}
        aggregated["epoch"] = self.current_epoch + 1

        self.log_dict(aggregated, logger=True, prog_bar=False, on_epoch=True)

        return aggregated

    def training_step(self, batch, batch_idx):
        """Compute loss for a single training step and log metrics to loggers

        Args:
            batch (Tuple[torch.Tensor, ...]): Input batch
            batch_idx (int): Index of batch

        Returns:
            Dict[str, torch.Tensor]: computed metrics
        """
        y_hat, targets = self.predictor.get_predictions_and_targets(self.model, batch)
        loss = self.criterion(y_hat, targets)
        metrics = self._compute_metrics(
            self.train_metrics, loss, y_hat, targets, mode="train"
        )

        self.log_dict(
            metrics,
            on_step=True,
            on_epoch=False,
            logger=True,
            prog_bar=False,
        )

        metrics["loss"] = loss

        return metrics

    def training_epoch_end(self, outputs):
        """Aggregate metrics of a training epoch

        Args:
            outputs (List[Dict[str, torch.Tensor]]): Aggregated outputs from train_step
        """
        outputs = self.aggregate_epoch_metrics(outputs, mode="Training")
        self.log_to_console(outputs, mode="Training")

    def validation_step(self, batch, batch_idx):
        """Compute loss for a single validation step and log metrics to loggers

        Args:
            batch (Tuple[torch.Tensor, ...]): Input batch
            batch_idx (int): Index of batch

        Returns:
            Dict[str, torch.Tensor]: computed metrics
        """
        y_hat, targets = self.predictor.get_predictions_and_targets(self, batch)
        loss = self.criterion(y_hat, targets)
        metrics = self._compute_metrics(
            self.val_metrics, loss, y_hat, targets, mode="val"
        )

        metrics[
            "best_score"
        ] = self.trainer.early_stopping_callback.best_score.detach().cpu()

        return metrics

    def validation_epoch_end(self, outputs):
        """Aggregate metrics of a validation epoch

        Args:
            outputs (List[Dict[str, torch.Tensor]]): Aggregated outputs from validation_step
        """
        outputs = self.aggregate_epoch_metrics(outputs, mode="Validation")

        if torch.isnan(outputs["val_loss"]) or torch.isinf(outputs["val_loss"]):
            outputs["val_loss"] = 1000000

        outputs["best_score"] = min(
            outputs[self.trainer.early_stopping_callback.monitor].detach().cpu(),
            self.trainer.early_stopping_callback.best_score.detach().cpu(),
        )
        self.log_to_console(outputs, mode="Validation")

    def test_step(self, batch, batch_idx):
        """Compute loss for a single test step and log metrics to loggers

        Args:
            batch (Tuple[torch.Tensor, ...]): Input batch
            batch_idx (int): Index of batch

        Returns:
            Dict[str, torch.Tensor]: computed metrics
        """
        y_hat, targets = self.predictor.get_predictions_and_targets(self, batch)
        loss = self.criterion(y_hat, targets)
        metrics = self._compute_metrics(
            self.test_metrics, loss, y_hat, targets, mode="test"
        )

        return metrics

    def test_epoch_end(self, outputs):
        """Aggregate metrics of a test epoch

        Args:
            outputs (List[Dict[str, torch.Tensor]]): Aggregated outputs from test_step
        """
        outputs = self.aggregate_epoch_metrics(outputs, mode="Test")
        self.log_to_console(outputs, mode="Test")


class PLModule(SimplePLModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        calculate_perplexity=False,
    ):
        """Pass arguments through to base class"""
        super(PLModule, self).__init__(
            model,
            optimizer,
            criterion,
            predictor_cls=_Classification,
            lr_scheduler=lr_scheduler,
            hparams=hparams,
            metrics=metrics,
            calculate_perplexity=calculate_perplexity,
        )


class AutoEncoderPLModule(SimplePLModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        calculate_perplexity=False,
    ):
        """Pass arguments through to base class"""
        super(AutoEncoderPLModule, self).__init__(
            model,
            optimizer,
            criterion,
            predictor_cls=_AutoEncoder,
            lr_scheduler=lr_scheduler,
            hparams=hparams,
            metrics=metrics,
            calculate_perplexity=calculate_perplexity,
        )


class RnnPLModule(SimplePLModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        calculate_perplexity=False,
    ):
        """Pass arguments through to base class"""
        super(RnnPLModule, self).__init__(
            model,
            optimizer,
            criterion,
            predictor_cls=_RnnClassification,
            lr_scheduler=lr_scheduler,
            hparams=hparams,
            metrics=metrics,
            calculate_perplexity=calculate_perplexity,
        )


class TransformerClassificationPLModule(SimplePLModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        calculate_perplexity=False,
    ):
        """Pass arguments through to base class"""
        super(TransformerClassificationPLModule, self).__init__(
            model,
            optimizer,
            criterion,
            predictor_cls=_TransformerClassification,
            lr_scheduler=lr_scheduler,
            hparams=hparams,
            metrics=metrics,
            calculate_perplexity=calculate_perplexity,
        )


class TransformerPLModule(SimplePLModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        calculate_perplexity=False,
    ):
        """Pass arguments through to base class"""
        super(TransformerPLModule, self).__init__(
            model,
            optimizer,
            criterion,
            predictor_cls=_Transformer,
            lr_scheduler=lr_scheduler,
            hparams=hparams,
            metrics=metrics,
            calculate_perplexity=calculate_perplexity,
        )


class BertPLModule(SimplePLModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        calculate_perplexity=False,
    ):
        """Pass arguments through to base class"""
        super(BertPLModule, self).__init__(
            model,
            optimizer,
            criterion,
            predictor_cls=_BertSequenceClassification,
            lr_scheduler=lr_scheduler,
            hparams=hparams,
            metrics=metrics,
            calculate_perplexity=calculate_perplexity,
        )


class MultimodalTransformerClassificationPLModule(SimplePLModule):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Union[Optimizer, List[Optimizer]],
        criterion: LossType,
        lr_scheduler: Union[_LRScheduler, List[_LRScheduler]] = None,
        hparams: Configuration = None,
        metrics: Optional[Dict[str, pl.metrics.Metric]] = None,
        calculate_perplexity=False,
    ):
        """Pass arguments through to base class"""
        super(MultimodalTransformerClassificationPLModule, self).__init__(
            model,
            optimizer,
            criterion,
            predictor_cls=_MultimodalTransformerClassification,
            lr_scheduler=lr_scheduler,
            hparams=hparams,
            metrics=metrics,
            calculate_perplexity=calculate_perplexity,
        )
Base predictor class      Define an interface that can be used to extend the lightning module to new tasks and models      * parse_batch: Parse input batch and extract necessery masks etc.     * get_predictions_and_targets: Perform a forward pass through the model to get the logits and return logits and targets Abstract parse_batch method to be implemented by child class          Args:             batch (Tuple[torch.Tensor, ...]): A tuple of tensors that contains inputs to the model and targets          Returns:             Tuple[torch.Tensor, ...]: The processed inputs, ready to provide to the model Abstract get_predictions and targets method to be implemented by child class          This method gets exposed to the PLModule classes          Args:             model (nn.Module): model to use for forward pass             batch (Tuple[torch.Tensor, torch.Tensor]): A tuple of tensors that contains inputs to the model and targets          **Note**: Maybe it should be useful to move loss calculation here. Then multitask learning and auxiliary losses should be easier          Returns:             Tuple[torch.Tensor, torch.Tensor]: (logits, ground_truths), ready to be passed to the loss function Classification task Parse incoming batch          Input batch just contains inputs and targets          Args:             batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, labels)          Returns:             Tuple[torch.Tensor, torch.Tensor]: (inputs, labels) Return logits and ground truths to be passed in loss function          Args:             model (nn.Module): Model to use for prediction             batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, labels)          Returns:             Tuple[torch.Tensor, torch.Tensor]: (logits, labels) Autoencoder task Parse incoming batch          Input batch just contains inputs. Targets are the same as inputs, because we are doing reconstruction.          Args:             batch (Tuple[torch.Tensor]): (inputs)          Returns:             Tuple[torch.Tensor, torch.Tensor]: (inputs, inputs) Return logits and ground truths to be passed in loss function          Args:             model (nn.Module): Model to use for prediction             batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)          Returns:             Tuple[torch.Tensor, torch.Tensor]: (logits, inputs) RNN classification task Parse incoming batch          Input batch just contains inputs, targets and lengths.         Comes from slp.data.collators.SequentialCollator.          Args:             batch (Tuple[torch.Tensor]): (inputs)          Returns:             Tuple[torch.Tensor, torch.Tensor]: (inputs, inputs) Return logits and ground truths to be passed in loss function          Args:             model (nn.Module): Model to use for prediction             batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)          Returns:             Tuple[torch.Tensor, torch.Tensor]: (logits, inputs) Transformer classification task Parse incoming batch          Input batch just contains inputs, targets and lengths.         Comes from slp.data.collators.SequentialCollator.         Create pad masks to be passed to transformer attention          Args:             batch (Tuple[torch.Tensor]): (inputs)          Returns:             Tuple[torch.Tensor, torch.Tensor]: (inputs, inputs) Return logits and ground truths to be passed in loss function          Args:             model (nn.Module): Model to use for prediction             batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)          Returns:             Tuple[torch.Tensor, torch.Tensor]: (logits, inputs) Transformer classification task Parse incoming batch          Input batch just contains inputs, targets and lengths.         Comes from slp.data.collators.SequentialCollator.         Create pad masks to be passed to transformer attention          Args:             batch (Tuple[Dict[str, torch.Tensor], torch.Tensor, Dict[str, torch.Tensor]]): (inputs, targets, lengths)          Returns:             Tuple[Dict[str, torch.Tensor], torch.Tensor, Dict[str, torch.Tensor]]: (inputs, targets, attention_masks) Return logits and ground truths to be passed in loss function          Args:             model (nn.Module): Model to use for prediction             batch (Tuple[Dict[str, torch.Tensor], torch.Tensor, Dict[str, torch.Tensor]]): (inputs, targets, lengths)          Returns:             Tuple[torch.Tensor, torch.Tensor]: (logits, targets) Generic transformer seq2seq task Parse incoming batch          Input batch just contains inputs, targets and lengths.         Comes from slp.data.collators.SequentialCollator.         Create pad masks and subsequent_masks to be passed to transformer attention          Args:             batch (Tuple[torch.Tensor]): (inputs)          Returns:             Tuple[torch.Tensor, torch.Tensor]: (inputs, targets) # type: ignore Return logits and ground truths to be passed in loss function          Args:             model (nn.Module): Model to use for prediction             batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)          Returns:             Tuple[torch.Tensor, torch.Tensor]: (logits, inputs) Bert Classification task Parse incoming batch          Input batch just contains inputs, targets and lengths.         Comes from slp.data.collators.SequentialCollator.         Create pad masks to be passed to BERT attention          Args:             batch (Tuple[torch.Tensor]): (inputs)          Returns:             Tuple[torch.Tensor, torch.Tensor]: (inputs, inputs) Return logits and ground truths to be passed in loss function          Args:             model (nn.Module): Model to use for prediction             batch (Tuple[torch.Tensor, torch.Tensor]): (inputs, inputs)          Returns:             Tuple[torch.Tensor, torch.Tensor]: (logits, inputs) # for LM. Dirty but much more efficient Wraps a (model, optimizer, criterion, lr_scheduler) tuple in a LightningModule          Handles the boilerplate for metrics calculation and logging and defines the train_step / val_step / test_step         with use of the predictor helper classes (e.g. _Classification, _RnnClassification)          Args:             model (nn.Module): Module to use for prediction             optimizer (Union[Optimizer, List[Optimizer]]): Optimizers to use for training             criterion (LossType): Task loss             lr_scheduler (Union[_LRScheduler, List[_LRScheduler]], optional): Learning rate scheduler. Defaults to None.             hparams (Configuration, optional): Hyperparameter values. This ensures they are logged with trainer.loggers. Defaults to None.             metrics (Optional[Dict[str, pl.metrics.Metric]], optional): Metrics to track. Defaults to None.             predictor_cls ([type], optional): Class that defines a parse_batch and a                     get_predictions_and_targets method. Defaults to _Classification.             calculate_perplexity (bool, optional): Whether to calculate perplexity.                     Would be cleaner as a metric, but this is more efficient. Defaults to False. # self.hparams = dict_params Return optimizers and learning rate schedulers          Returns:             Tuple[List[Optimizer], List[_LRScheduler]]: (optimizers, lr_schedulers) Call wrapped module forward Compute all metrics and aggregate in a dict          Args:             metrics (Dict[str, pl.metrics.Metric]): metrics to compute             loss (torch.Tensor): Computed loss             y_hat (torch.Tensor): Logits             targets (torch.Tensor): Ground Truths             mode (str, optional): "train", "val" or "test". Defaults to "train". Format metric name Log metrics to console          Args:             metrics (Dict[str, torch.Tensor]): Computed metrics             mode (str, optional): "Training", "Validation" or "Testing". Defaults to "Training". Aggregate metrics over a whole epoch          Args:             outputs (List[Dict[str, torch.Tensor]]): Aggregated outputs from train_step, validation_step or test_step             mode (str, optional): "Training", "Validation" or "Testing". Defaults to "Training". Format metric name Compute loss for a single training step and log metrics to loggers          Args:             batch (Tuple[torch.Tensor, ...]): Input batch             batch_idx (int): Index of batch          Returns:             Dict[str, torch.Tensor]: computed metrics Aggregate metrics of a training epoch          Args:             outputs (List[Dict[str, torch.Tensor]]): Aggregated outputs from train_step Compute loss for a single validation step and log metrics to loggers          Args:             batch (Tuple[torch.Tensor, ...]): Input batch             batch_idx (int): Index of batch          Returns:             Dict[str, torch.Tensor]: computed metrics Aggregate metrics of a validation epoch          Args:             outputs (List[Dict[str, torch.Tensor]]): Aggregated outputs from validation_step Compute loss for a single test step and log metrics to loggers          Args:             batch (Tuple[torch.Tensor, ...]): Input batch             batch_idx (int): Index of batch          Returns:             Dict[str, torch.Tensor]: computed metrics Aggregate metrics of a test epoch          Args:             outputs (List[Dict[str, torch.Tensor]]): Aggregated outputs from test_step Pass arguments through to base class Pass arguments through to base class Pass arguments through to base class Pass arguments through to base class Pass arguments through to base class Pass arguments through to base class Pass arguments through to base class