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, sentence class
                ids, masks for the input ids, masks for  sentence class ids and source
                text. If train_model is True, it also contains the labels 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, masks for the sentence class ids and labels.
            Labels are only returned when train_mode is True.
        """

        if model_name.split("-")[0] in ["bert", "distilbert"]:
            if train_mode:
                batch = batch.to(device)
                # labels must be the last
                return {
                    "x": batch.src,
                    "segs": batch.segs,
                    "clss": batch.clss,
                    "mask": batch.mask,
                    "mask_cls": batch.mask_cls,
                    "labels": batch.labels,
                }
            else:
                batch = batch.to(device)
                return {
                    "x": batch.src,
                    "segs": batch.segs,
                    "clss": batch.clss,
                    "mask": batch.mask,
                    "mask_cls": batch.mask_cls,
                    # "labels": batch.labels,
                }
                """
                return {
                    "x": batch.src.to(device),
                    "segs": batch.segs.to(device),
                    "clss": batch.clss.to(device),
                    "mask": batch.mask.to(device),
                    "mask_cls": batch.mask_cls.to(device),
                }
                """
        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, sentence class
                ids, masks for the input ids, masks for  sentence class ids and source
                text. If train_model is True, it also contains the labels 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, masks for the sentence class ids and labels.
            Labels are only returned when train_mode is True.
        
utils_nlp/models/transformers/extractive_summarization.py