def load_data(
        self,
        hf_dataset: Dataset,
        input_key: str,
        target_keys: Optional[Union[str, List[str]]] = None,
        target_formatter: Optional[TargetFormatter] = None,
    ) -> Dataset:
        """Loads data into HuggingFace datasets.Dataset."""
        if not self.predicting:
            hf_dataset = hf_dataset.map(partial(self._resolve_target, target_keys))
            targets = hf_dataset.to_dict()[DataKeys.TARGET]
            self.load_target_metadata(targets, target_formatter=target_formatter)

            # If we had binary multi-class targets then we also know the labels (column names)
            if (
                hasattr(self, "target_formatter")
                and isinstance(self.target_formatter, MultiBinaryTargetFormatter)
                and isinstance(target_keys, List)
            ):
                self.labels = target_keys

        # remove extra columns
        extra_columns = set(hf_dataset.column_names) - {input_key, DataKeys.TARGET}
        hf_dataset = hf_dataset.remove_columns(extra_columns)

        if input_key != DataKeys.INPUT:
            hf_dataset = hf_dataset.rename_column(input_key, DataKeys.INPUT)

        return hf_dataset
Loads data into HuggingFace datasets.Dataset.
src/flash/text/classification/input.py