def METHOD_NAME(self) -> Dict[str, Union[datasets.Dataset, datasets.DatasetDict]]:
    if self.data_config.pad_to_max_length:
        padding = "max_length"
    else:
        # We will pad later, dynamically at batch creation to the max_seq_length in each batch.
        padding = False
    # We cannot use self.tokenizer as a non-local variable in the preprocess_function if we
    # want map to be able to cache the output of the tokenizer.  Hence, the preprocess_function
    # takes a tokenizer explicitly as an input and we create a closure using functools.partial.
    def preprocess_function(tokenizer, padding, max_length, examples):
        # Tokenize the texts
        return tokenizer(
            examples["premise"],
            examples["hypothesis"],
            padding=padding,
            max_length=max_length,
            truncation=True,
        )
    train_dataset = self.raw_datasets["train"].map(
        functools.partial(
            preprocess_function, self.tokenizer, padding, self.data_config.max_seq_length
        ),
        batched=True,
        load_from_cache_file=not self.data_config.overwrite_cache,
    )
    eval_dataset = self.raw_datasets["validation"].map(
        functools.partial(
            preprocess_function, self.tokenizer, padding, self.data_config.max_seq_length
        ),
        batched=True,
        load_from_cache_file=not self.data_config.overwrite_cache,
    )
    if self.data_config.pad_to_max_length:
        self.collator = transformers.default_data_collator
    else:
        collator = transformers.DataCollatorWithPadding(
            self.tokenizer, pad_to_multiple_of=8 if self.hparams.use_apex_amp else None
        )
        self.collator = lambda x: collator(x).data
    return {"train": train_dataset, "validation": eval_dataset}