def get_inputs(batch, device, model_name, train_mode=True):
        """
        Creates an input dictionary given a model name.

        Args:
            batch (object): A Batch containing input ids, segment ids,
                masks for the input ids and source text. If train_mode is True, it
                also contains the target ids and the number of tokens
                in the target and target text.
            device (torch.device): A PyTorch device.
            model_name (bool): Model name used to format the inputs.
            train_mode (bool, optional): Training mode flag.
                Defaults to True.

        Returns:
            dict: Dictionary containing input ids, segment ids, sentence class ids,
            masks for the input ids. Target ids and number of tokens in the target are
            only returned when train_mode is True.
        """

        if model_name.split("-")[0] in ["bert"]:
            if train_mode:
                # labels must be the last

                return {
                    "src": batch.src,
                    "segs": batch.segs,
                    "mask_src": batch.mask_src,
                    "tgt": batch.tgt,
                    "tgt_num_tokens": batch.tgt_num_tokens,
                }
            else:
                return {
                    "src": batch.src,
                    "segs": batch.segs,
                    "mask_src": batch.mask_src,
                }
        else:
            raise ValueError("Model not supported: {}".format(model_name))

        Creates an input dictionary given a model name.

        Args:
            batch (object): A Batch containing input ids, segment ids,
                masks for the input ids and source text. If train_mode is True, it
                also contains the target ids and the number of tokens
                in the target and target text.
            device (torch.device): A PyTorch device.
            model_name (bool): Model name used to format the inputs.
            train_mode (bool, optional): Training mode flag.
                Defaults to True.

        Returns:
            dict: Dictionary containing input ids, segment ids, sentence class ids,
            masks for the input ids. Target ids and number of tokens in the target are
            only returned when train_mode is True.
        
utils_nlp/models/transformers/abstractive_summarization_bertsum.py