Repository: Lightning-Universe/lightning-flash

Issue:
Text classification tokenization currently failing
## 🐛 Bug

Text classification tokenization currently failing.

### To Reproduce

Run the `flash_notebooks/text_classification.ipynb` in colab.


Discussion:

Resolving patch:
diff --git a/flash/core/data/data_source.py b/flash/core/data/data_source.py
--- a/flash/core/data/data_source.py
+++ b/flash/core/data/data_source.py
@@ -158,11 +158,12 @@ class DefaultDataSources(LightningEnum):
     TENSORS = "tensors"
     CSV = "csv"
     JSON = "json"
+    PARQUET = "parquet"
     DATASETS = "datasets"
+    HUGGINGFACE_DATASET = "hf_dataset"
     FIFTYONE = "fiftyone"
     DATAFRAME = "data_frame"
     LISTS = "lists"
-    SENTENCES = "sentences"
     LABELSTUDIO = "labelstudio"
 
     # TODO: Create a FlashEnum class???
diff --git a/flash/text/classification/data.py b/flash/text/classification/data.py
--- a/flash/text/classification/data.py
+++ b/flash/text/classification/data.py
@@ -14,8 +14,8 @@
 from functools import partial
 from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Type, Union
 
-import pandas as pd
 import torch
+from pandas.core.frame import DataFrame
 from torch import Tensor
 from torch.utils.data.sampler import Sampler
 
@@ -23,13 +23,13 @@
 from flash.core.data.auto_dataset import AutoDataset
 from flash.core.data.callback import BaseDataFetcher
 from flash.core.data.data_module import DataModule
-from flash.core.data.data_source import DataSource, DefaultDataSources, LabelsState
+from flash.core.data.data_source import DataSource, DefaultDataKeys, DefaultDataSources, LabelsState
 from flash.core.data.process import Deserializer, Postprocess, Preprocess
 from flash.core.integrations.labelstudio.data_source import LabelStudioTextClassificationDataSource
 from flash.core.utilities.imports import _TEXT_AVAILABLE, requires
 
 if _TEXT_AVAILABLE:
-    from datasets import Dataset, DatasetDict, load_dataset
+    from datasets import Dataset, load_dataset
     from transformers import AutoTokenizer, default_data_collator
     from transformers.modeling_outputs import SequenceClassifierOutput
 
@@ -74,126 +74,75 @@ def _tokenize_fn(
         input: Optional[str] = None,
     ) -> Callable:
         """This function is used to tokenize sentences using the provided tokenizer."""
-        if isinstance(ex, dict):
-            ex = ex[input]
-        return self.tokenizer(ex, max_length=self.max_length, truncation=True, padding="max_length")
+        return self.tokenizer(ex[input], max_length=self.max_length, truncation=True, padding="max_length")
 
     @staticmethod
     def _transform_label(label_to_class_mapping: Dict[str, int], target: str, ex: Dict[str, Union[int, str]]):
         ex[target] = label_to_class_mapping[ex[target]]
         return ex
 
-    def __getstate__(self):  # TODO: Find out why this is being pickled
-        state = self.__dict__.copy()
-        state.pop("tokenizer")
-        return state
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-        self.tokenizer = AutoTokenizer.from_pretrained(self.backbone, use_fast=True)
-
-
-class TextFileDataSource(TextDataSource):
-    def __init__(self, filetype: str, backbone: str, max_length: int = 128):
-        super().__init__(backbone, max_length=max_length)
-
-        self.filetype = filetype
-
     @staticmethod
-    def _multilabel_target(targets, element):
+    def _multilabel_target(targets: List[str], element: Dict[str, Any]) -> Dict[str, Any]:
         targets = [element.pop(target) for target in targets]
-        element["labels"] = targets
+        element[DefaultDataKeys.TARGET] = targets
         return element
 
+    def _to_hf_dataset(self, data) -> Sequence[Mapping[str, Any]]:
+        """account for flash CI testing context."""
+        hf_dataset, *other = self.to_hf_dataset(data)
+
+        if flash._IS_TESTING and not torch.cuda.is_available():
+            # NOTE: must subset in this way to return a Dataset
+            hf_dataset = hf_dataset.select(range(20))
+
+        return (hf_dataset, *other)
+
     def load_data(
         self,
         data: Tuple[str, Union[str, List[str]], Union[str, List[str]]],
         dataset: Optional[Any] = None,
-        columns: Union[List[str], Tuple[str]] = ("input_ids", "attention_mask", "labels"),
-    ) -> Union[Sequence[Mapping[str, Any]]]:
-        if self.filetype == "json":
-            file, input, target, field = data
-        else:
-            file, input, target = data
-
-        data_files = {}
+    ) -> Sequence[Mapping[str, Any]]:
+        """Loads data into HuggingFace datasets.Dataset."""
 
-        stage = self.running_stage.value
-        data_files[stage] = str(file)
-
-        # FLASH_TESTING is set in the CI to run faster.
-        if flash._IS_TESTING and not torch.cuda.is_available():
-            try:
-                if self.filetype == "json" and field is not None:
-                    dataset_dict = DatasetDict(
-                        {
-                            stage: load_dataset(
-                                self.filetype, data_files=data_files, split=[f"{stage}[:20]"], field=field
-                            )[0]
-                        }
-                    )
-                else:
-                    dataset_dict = DatasetDict(
-                        {stage: load_dataset(self.filetype, data_files=data_files, split=[f"{stage}[:20]"])[0]}
-                    )
-            except Exception:
-                if self.filetype == "json" and field is not None:
-                    dataset_dict = load_dataset(self.filetype, data_files=data_files, field=field)
-                else:
-                    dataset_dict = load_dataset(self.filetype, data_files=data_files)
-        else:
-            if self.filetype == "json" and field is not None:
-                dataset_dict = load_dataset(self.filetype, data_files=data_files, field=field)
-            else:
-                dataset_dict = load_dataset(self.filetype, data_files=data_files)
+        hf_dataset, input, *other = self._to_hf_dataset(data)
 
         if not self.predicting:
+            target: Union[str, List[str]] = other.pop()
             if isinstance(target, List):
                 # multi-target
                 dataset.multi_label = True
-                dataset_dict = dataset_dict.map(partial(self._multilabel_target, target))
+                hf_dataset = hf_dataset.map(partial(self._multilabel_target, target))  # NOTE: renames target column
                 dataset.num_classes = len(target)
                 self.set_state(LabelsState(target))
             else:
                 dataset.multi_label = False
                 if self.training:
-                    labels = list(sorted(list(set(dataset_dict[stage][target]))))
+                    labels = list(sorted(list(set(hf_dataset[target]))))
                     dataset.num_classes = len(labels)
                     self.set_state(LabelsState(labels))
 
                 labels = self.get_state(LabelsState)
 
-                # convert labels to ids
+                # convert labels to ids (note: the target column get overwritten)
                 if labels is not None:
                     labels = labels.labels
                     label_to_class_mapping = {v: k for k, v in enumerate(labels)}
-                    dataset_dict = dataset_dict.map(partial(self._transform_label, label_to_class_mapping, target))
-
-                # Hugging Face models expect target to be named ``labels``.
-                if target != "labels":
-                    dataset_dict.rename_column_(target, "labels")
-
-        dataset_dict = dataset_dict.map(partial(self._tokenize_fn, input=input), batched=True)
-        dataset_dict.set_format("torch", columns=columns)
-
-        return dataset_dict[stage]
+                    hf_dataset = hf_dataset.map(partial(self._transform_label, label_to_class_mapping, target))
 
-    def predict_load_data(self, data: Any, dataset: AutoDataset):
-        return self.load_data(data, dataset, columns=["input_ids", "attention_mask"])
+                # rename label column
+                hf_dataset = hf_dataset.rename_column(target, DefaultDataKeys.TARGET)
 
-    def __getstate__(self):  # TODO: Find out why this is being pickled
-        state = self.__dict__.copy()
-        state.pop("tokenizer")
-        return state
+        # tokenize
+        hf_dataset = hf_dataset.map(partial(self._tokenize_fn, input=input), batched=True)
 
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-        self.tokenizer = AutoTokenizer.from_pretrained(self.backbone, use_fast=True)
+        # set format
+        hf_dataset = hf_dataset.remove_columns([input])  # just leave the numerical columns
+        hf_dataset.set_format("torch")
 
+        return hf_dataset
 
-class TextCSVDataSource(TextFileDataSource):
-    def __init__(self, backbone: str, max_length: int = 128):
-        super().__init__("csv", backbone, max_length=max_length)
+    def predict_load_data(self, data: Any, dataset: AutoDataset):
+        return self.load_data(data, dataset)
 
     def __getstate__(self):  # TODO: Find out why this is being pickled
         state = self.__dict__.copy()
@@ -205,156 +154,98 @@ def __setstate__(self, state):
         self.tokenizer = AutoTokenizer.from_pretrained(self.backbone, use_fast=True)
 
 
-class TextJSONDataSource(TextFileDataSource):
-    def __init__(self, backbone: str, max_length: int = 128):
-        super().__init__("json", backbone, max_length=max_length)
+class TextCSVDataSource(TextDataSource):
+    def to_hf_dataset(self, data: Tuple[str, str, str]) -> Tuple[Sequence[Mapping[str, Any]], str, str]:
+        file, *other = data
+        dataset_dict = load_dataset("csv", data_files={"train": str(file)})
+        return (dataset_dict["train"], *other)
 
-    def __getstate__(self):  # TODO: Find out why this is being pickled
-        state = self.__dict__.copy()
-        state.pop("tokenizer")
-        return state
 
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-        self.tokenizer = AutoTokenizer.from_pretrained(self.backbone, use_fast=True)
+class TextJSONDataSource(TextDataSource):
+    def to_hf_dataset(self, data: Tuple[str, str, str, str]) -> Tuple[Sequence[Mapping[str, Any]], str, str]:
+        file, *other, field = data
+        dataset_dict = load_dataset("json", data_files={"train": str(file)}, field=field)
+        return (dataset_dict["train"], *other)
 
 
 class TextDataFrameDataSource(TextDataSource):
-    @staticmethod
-    def _multilabel_target(targets, element):
-        targets = [element.pop(target) for target in targets]
-        element["labels"] = targets
-        return element
-
-    def load_data(
-        self,
-        data: Tuple[pd.DataFrame, Union[str, List[str]], Union[str, List[str]]],
-        dataset: Optional[Any] = None,
-        columns: Union[List[str], Tuple[str]] = ("input_ids", "attention_mask", "labels"),
-    ) -> Union[Sequence[Mapping[str, Any]]]:
-        df, input, target = data
+    def to_hf_dataset(self, data: Tuple[DataFrame, str, str]) -> Tuple[Sequence[Mapping[str, Any]], str, str]:
+        df, *other = data
         hf_dataset = Dataset.from_pandas(df)
+        return (hf_dataset, *other)
 
-        if not self.predicting:
-            if isinstance(target, List):
-                # multi-target
-                dataset.multi_label = True
-                hf_dataset = hf_dataset.map(partial(self._multilabel_target, target))
-                dataset.num_classes = len(target)
-                self.set_state(LabelsState(target))
-            else:
-                dataset.multi_label = False
-                if self.training:
-                    labels = list(sorted(list(set(hf_dataset[target]))))
-                    dataset.num_classes = len(labels)
-                    self.set_state(LabelsState(labels))
 
-                labels = self.get_state(LabelsState)
+class TextParquetDataSource(TextDataSource):
+    def to_hf_dataset(self, data: Tuple[str, str, str]) -> Tuple[Sequence[Mapping[str, Any]], str, str]:
+        file, *other = data
+        hf_dataset = Dataset.from_parquet(str(file))
+        return (hf_dataset, *other)
 
-                # convert labels to ids
-                if labels is not None:
-                    labels = labels.labels
-                    label_to_class_mapping = {v: k for k, v in enumerate(labels)}
-                    hf_dataset = hf_dataset.map(partial(self._transform_label, label_to_class_mapping, target))
 
-                # Hugging Face models expect target to be named ``labels``.
-                if target != "labels":
-                    hf_dataset.rename_column_(target, "labels")
+class TextHuggingFaceDatasetDataSource(TextDataSource):
+    def to_hf_dataset(self, data: Tuple[str, str, str]) -> Tuple[Sequence[Mapping[str, Any]], str, str]:
+        hf_dataset, *other = data
+        return (hf_dataset, *other)
 
-        hf_dataset = hf_dataset.map(partial(self._tokenize_fn, input=input), batched=True)
-        hf_dataset.set_format("torch", columns=columns)
 
-        return hf_dataset
-
-    def predict_load_data(self, data: Any, dataset: AutoDataset):
-        return self.load_data(data, dataset, columns=["input_ids", "attention_mask"])
+class TextListDataSource(TextDataSource):
+    def to_hf_dataset(
+        self, data: Union[Tuple[List[str], List[str]], List[str]]
+    ) -> Tuple[Sequence[Mapping[str, Any]], Optional[List[str]]]:
 
-    def __getstate__(self):  # TODO: Find out why this is being pickled
-        state = self.__dict__.copy()
-        state.pop("tokenizer")
-        return state
+        if isinstance(data, tuple):
+            input_list, target_list = data
+            # NOTE: here we already deal with multilabels
+            # NOTE: here we already rename to correct column names
+            hf_dataset = Dataset.from_dict({DefaultDataKeys.INPUT: input_list, DefaultDataKeys.TARGET: target_list})
+            return hf_dataset, target_list
 
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-        self.tokenizer = AutoTokenizer.from_pretrained(self.backbone, use_fast=True)
+        # predicting
+        hf_dataset = Dataset.from_dict({DefaultDataKeys.INPUT: data})
 
+        return (hf_dataset,)
 
-class TextListDataSource(TextDataSource):
     def load_data(
         self,
         data: Tuple[List[str], Union[List[Any], List[List[Any]]]],
         dataset: Optional[Any] = None,
-        columns: Union[List[str], Tuple[str]] = ("input_ids", "attention_mask", "labels"),
-    ) -> Union[Sequence[Mapping[str, Any]]]:
-        input, target = data
-        hf_dataset = Dataset.from_dict({"input": input, "labels": target})
+    ) -> Sequence[Mapping[str, Any]]:
+
+        hf_dataset, *other = self._to_hf_dataset(data)
 
         if not self.predicting:
-            if isinstance(target[0], List):
-                # multi-target
+            target_list = other.pop()
+            if isinstance(target_list[0], List):
+                # multi-target_list
                 dataset.multi_label = True
-                dataset.num_classes = len(target[0])
-                self.set_state(LabelsState(target))
+                dataset.num_classes = len(target_list[0])
+                self.set_state(LabelsState(target_list))
             else:
                 dataset.multi_label = False
                 if self.training:
-                    labels = list(sorted(list(set(hf_dataset["labels"]))))
+                    labels = list(sorted(list(set(hf_dataset[DefaultDataKeys.TARGET]))))
                     dataset.num_classes = len(labels)
                     self.set_state(LabelsState(labels))
 
-                    labels = self.get_state(LabelsState)
+                labels = self.get_state(LabelsState)
 
                 # convert labels to ids
                 if labels is not None:
                     labels = labels.labels
                     label_to_class_mapping = {v: k for k, v in enumerate(labels)}
-                    hf_dataset = hf_dataset.map(partial(self._transform_label, label_to_class_mapping, "labels"))
-
-        hf_dataset = hf_dataset.map(partial(self._tokenize_fn, input="input"), batched=True)
-        hf_dataset.set_format("torch", columns=columns)
-
-        return hf_dataset
-
-    def predict_load_data(self, data: Any, dataset: AutoDataset):
-        return self.load_data(data, dataset, columns=["input_ids", "attention_mask"])
-
-    def __getstate__(self):  # TODO: Find out why this is being pickled
-        state = self.__dict__.copy()
-        state.pop("tokenizer")
-        return state
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-        self.tokenizer = AutoTokenizer.from_pretrained(self.backbone, use_fast=True)
-
-
-class TextSentencesDataSource(TextDataSource):
-    def __init__(self, backbone: str, max_length: int = 128):
-        super().__init__(backbone, max_length=max_length)
-
-    def load_data(
-        self,
-        data: Union[str, List[str]],
-        dataset: Optional[Any] = None,
-    ) -> Union[Sequence[Mapping[str, Any]]]:
+                    # happens in-place and keeps the target column name
+                    hf_dataset = hf_dataset.map(
+                        partial(self._transform_label, label_to_class_mapping, DefaultDataKeys.TARGET)
+                    )
 
-        if isinstance(data, str):
-            data = [data]
-        return [
-            self._tokenize_fn(
-                s,
-            )
-            for s in data
-        ]
+        # tokenize
+        hf_dataset = hf_dataset.map(partial(self._tokenize_fn, input=DefaultDataKeys.INPUT), batched=True)
 
-    def __getstate__(self):  # TODO: Find out why this is being pickled
-        state = self.__dict__.copy()
-        state.pop("tokenizer")
-        return state
+        # set format
+        hf_dataset = hf_dataset.remove_columns([DefaultDataKeys.INPUT])  # just leave the numerical columns
+        hf_dataset.set_format("torch")
 
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-        self.tokenizer = AutoTokenizer.from_pretrained(self.backbone, use_fast=True)
+        return hf_dataset
 
 
 class TextClassificationPreprocess(Preprocess):
@@ -379,14 +270,17 @@ def __init__(
             data_sources={
                 DefaultDataSources.CSV: TextCSVDataSource(self.backbone, max_length=max_length),
                 DefaultDataSources.JSON: TextJSONDataSource(self.backbone, max_length=max_length),
+                DefaultDataSources.PARQUET: TextParquetDataSource(self.backbone, max_length=max_length),
+                DefaultDataSources.HUGGINGFACE_DATASET: TextHuggingFaceDatasetDataSource(
+                    self.backbone, max_length=max_length
+                ),
                 DefaultDataSources.DATAFRAME: TextDataFrameDataSource(self.backbone, max_length=max_length),
                 DefaultDataSources.LISTS: TextListDataSource(self.backbone, max_length=max_length),
-                DefaultDataSources.SENTENCES: TextSentencesDataSource(self.backbone, max_length=max_length),
                 DefaultDataSources.LABELSTUDIO: LabelStudioTextClassificationDataSource(
                     backbone=self.backbone, max_length=max_length
                 ),
             },
-            default_data_source=DefaultDataSources.SENTENCES,
+            default_data_source=DefaultDataSources.LISTS,
             deserializer=TextDeserializer(backbone, max_length),
         )
 
@@ -437,10 +331,10 @@ def from_data_frame(
         cls,
         input_field: str,
         target_fields: Union[str, Sequence[str]],
-        train_data_frame: Optional[pd.DataFrame] = None,
-        val_data_frame: Optional[pd.DataFrame] = None,
-        test_data_frame: Optional[pd.DataFrame] = None,
-        predict_data_frame: Optional[pd.DataFrame] = None,
+        train_data_frame: Optional[DataFrame] = None,
+        val_data_frame: Optional[DataFrame] = None,
+        test_data_frame: Optional[DataFrame] = None,
+        predict_data_frame: Optional[DataFrame] = None,
         train_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
         val_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
         test_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
@@ -582,3 +476,162 @@ def from_lists(
             sampler=sampler,
             **preprocess_kwargs,
         )
+
+    @classmethod
+    def from_parquet(
+        cls,
+        input_field: str,
+        target_fields: Optional[Union[str, Sequence[str]]] = None,
+        train_file: Optional[str] = None,
+        val_file: Optional[str] = None,
+        test_file: Optional[str] = None,
+        predict_file: Optional[str] = None,
+        train_transform: Optional[Dict[str, Callable]] = None,
+        val_transform: Optional[Dict[str, Callable]] = None,
+        test_transform: Optional[Dict[str, Callable]] = None,
+        predict_transform: Optional[Dict[str, Callable]] = None,
+        data_fetcher: Optional[BaseDataFetcher] = None,
+        preprocess: Optional[Preprocess] = None,
+        val_split: Optional[float] = None,
+        batch_size: int = 4,
+        num_workers: int = 0,
+        sampler: Optional[Type[Sampler]] = None,
+        **preprocess_kwargs: Any,
+    ) -> "DataModule":
+        """Creates a :class:`~flash.core.data.data_module.DataModule` object from the given PARQUET files using the
+        :class:`~flash.core.data.data_source.DataSource`
+        of name :attr:`~flash.core.data.data_source.DefaultDataSources.PARQUET`
+        from the passed or constructed :class:`~flash.core.data.process.Preprocess`.
+
+        Args:
+            input_fields: The field or fields (columns) in the PARQUET file to use for the input.
+            target_fields: The field or fields (columns) in the PARQUET file to use for the target.
+            train_file: The PARQUET file containing the training data.
+            val_file: The PARQUET file containing the validation data.
+            test_file: The PARQUET file containing the testing data.
+            predict_file: The PARQUET file containing the data to use when predicting.
+            train_transform: The dictionary of transforms to use during training which maps
+                :class:`~flash.core.data.process.Preprocess` hook names to callable transforms.
+            val_transform: The dictionary of transforms to use during validation which maps
+                :class:`~flash.core.data.process.Preprocess` hook names to callable transforms.
+            test_transform: The dictionary of transforms to use during testing which maps
+                :class:`~flash.core.data.process.Preprocess` hook names to callable transforms.
+            predict_transform: The dictionary of transforms to use during predicting which maps
+                :class:`~flash.core.data.process.Preprocess` hook names to callable transforms.
+            data_fetcher: The :class:`~flash.core.data.callback.BaseDataFetcher` to pass to the
+                :class:`~flash.core.data.data_module.DataModule`.
+            preprocess: The :class:`~flash.core.data.data.Preprocess` to pass to the
+                :class:`~flash.core.data.data_module.DataModule`. If ``None``, ``cls.preprocess_cls``
+                will be constructed and used.
+            val_split: The ``val_split`` argument to pass to the :class:`~flash.core.data.data_module.DataModule`.
+            batch_size: The ``batch_size`` argument to pass to the :class:`~flash.core.data.data_module.DataModule`.
+            num_workers: The ``num_workers`` argument to pass to the :class:`~flash.core.data.data_module.DataModule`.
+            sampler: The ``sampler`` to use for the ``train_dataloader``.
+            preprocess_kwargs: Additional keyword arguments to use when constructing the preprocess. Will only be used
+                if ``preprocess = None``.
+
+        Returns:
+            The constructed data module.
+
+        Examples::
+
+            data_module = DataModule.from_parquet(
+                "input",
+                "target",
+                train_file="train_data.parquet",
+                train_transform={
+                    "to_tensor_transform": torch.as_tensor,
+                },
+            )
+        """
+        return cls.from_data_source(
+            DefaultDataSources.PARQUET,
+            (train_file, input_field, target_fields),
+            (val_file, input_field, target_fields),
+            (test_file, input_field, target_fields),
+            (predict_file, input_field, target_fields),
+            train_transform=train_transform,
+            val_transform=val_transform,
+            test_transform=test_transform,
+            predict_transform=predict_transform,
+            data_fetcher=data_fetcher,
+            preprocess=preprocess,
+            val_split=val_split,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            sampler=sampler,
+            **preprocess_kwargs,
+        )
+
+    @classmethod
+    def from_hf_datasets(
+        cls,
+        input_field: str,
+        target_fields: Union[str, Sequence[str]],
+        train_hf_dataset: Optional[Sequence[Mapping[str, Any]]] = None,
+        val_hf_dataset: Optional[Sequence[Mapping[str, Any]]] = None,
+        test_hf_dataset: Optional[Sequence[Mapping[str, Any]]] = None,
+        predict_hf_dataset: Optional[Sequence[Mapping[str, Any]]] = None,
+        train_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        val_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        test_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        predict_transform: Optional[Dict[str, Callable]] = None,
+        data_fetcher: Optional[BaseDataFetcher] = None,
+        preprocess: Optional[Preprocess] = None,
+        val_split: Optional[float] = None,
+        batch_size: int = 4,
+        num_workers: int = 0,
+        sampler: Optional[Type[Sampler]] = None,
+        **preprocess_kwargs: Any,
+    ) -> "DataModule":
+        """Creates a :class:`~flash.text.classification.data.TextClassificationData` object from the given Hugging
+        Face datasets ``Dataset`` objects.
+
+        Args:
+            input_field: The field (column) in the pandas ``Dataset`` to use for the input.
+            target_fields: The field or fields (columns) in the pandas ``Dataset`` to use for the target.
+            train_hf_dataset: The pandas ``Dataset`` containing the training data.
+            val_hf_dataset: The pandas ``Dataset`` containing the validation data.
+            test_hf_dataset: The pandas ``Dataset`` containing the testing data.
+            predict_hf_dataset: The pandas ``Dataset`` containing the data to use when predicting.
+            train_transform: The dictionary of transforms to use during training which maps
+                :class:`~flash.core.data.process.Preprocess` hook names to callable transforms.
+            val_transform: The dictionary of transforms to use during validation which maps
+                :class:`~flash.core.data.process.Preprocess` hook names to callable transforms.
+            test_transform: The dictionary of transforms to use during testing which maps
+                :class:`~flash.core.data.process.Preprocess` hook names to callable transforms.
+            predict_transform: The dictionary of transforms to use during predicting which maps
+                :class:`~flash.core.data.process.Preprocess` hook names to callable transforms.
+            data_fetcher: The :class:`~flash.core.data.callback.BaseDataFetcher` to pass to the
+                :class:`~flash.core.data.data_module.DataModule`.
+            preprocess: The :class:`~flash.core.data.data.Preprocess` to pass to the
+                :class:`~flash.core.data.data_module.DataModule`. If ``None``, ``cls.preprocess_cls``
+                will be constructed and used.
+            val_split: The ``val_split`` argument to pass to the :class:`~flash.core.data.data_module.DataModule`.
+            batch_size: The ``batch_size`` argument to pass to the :class:`~flash.core.data.data_module.DataModule`.
+            num_workers: The ``num_workers`` argument to pass to the :class:`~flash.core.data.data_module.DataModule`.
+            sampler: The ``sampler`` to use for the ``train_dataloader``.
+            preprocess_kwargs: Additional keyword arguments to use when constructing the preprocess. Will only be used
+                if ``preprocess = None``.
+
+        Returns:
+            The constructed data module.
+        """
+        return cls.from_data_source(
+            DefaultDataSources.HUGGINGFACE_DATASET,
+            (train_hf_dataset, input_field, target_fields),
+            (val_hf_dataset, input_field, target_fields),
+            (test_hf_dataset, input_field, target_fields),
+            (predict_hf_dataset, input_field, target_fields),
+            train_transform=train_transform,
+            val_transform=val_transform,
+            test_transform=test_transform,
+            predict_transform=predict_transform,
+            data_fetcher=data_fetcher,
+            preprocess=preprocess,
+            val_split=val_split,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            sampler=sampler,
+            **preprocess_kwargs,
+        )
diff --git a/flash/text/classification/model.py b/flash/text/classification/model.py
--- a/flash/text/classification/model.py
+++ b/flash/text/classification/model.py
@@ -21,6 +21,7 @@
 from torchmetrics import Metric
 
 from flash.core.classification import ClassificationTask, Labels
+from flash.core.data.data_source import DefaultDataKeys
 from flash.core.data.process import Serializer
 from flash.core.registry import FlashRegistry
 from flash.core.utilities.imports import _TRANSFORMERS_AVAILABLE
@@ -115,7 +116,7 @@ def to_metrics_format(self, x) -> torch.Tensor:
         return super().to_metrics_format(x)
 
     def step(self, batch, batch_idx, metrics) -> dict:
-        target = batch.pop("labels")
+        target = batch.pop(DefaultDataKeys.TARGET)
         batch = (batch, target)
         return super().step(batch, batch_idx, metrics)
 


Test patch:
diff --git a/tests/text/classification/test_data.py b/tests/text/classification/test_data.py
--- a/tests/text/classification/test_data.py
+++ b/tests/text/classification/test_data.py
@@ -17,20 +17,22 @@
 import pandas as pd
 import pytest
 
+from flash.core.data.data_source import DefaultDataKeys
 from flash.core.utilities.imports import _TEXT_AVAILABLE
 from flash.text import TextClassificationData
 from flash.text.classification.data import (
     TextCSVDataSource,
     TextDataFrameDataSource,
     TextDataSource,
-    TextFileDataSource,
+    TextHuggingFaceDatasetDataSource,
     TextJSONDataSource,
     TextListDataSource,
-    TextSentencesDataSource,
+    TextParquetDataSource,
 )
 from tests.helpers.utils import _TEXT_TESTING
 
 if _TEXT_AVAILABLE:
+    from datasets import Dataset
     from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
 TEST_BACKBONE = "prajjwal1/bert-tiny"  # super small model for testing
@@ -41,18 +43,35 @@
 this is a sentence three,0
 """
 
+TEST_CSV_DATA_MULTILABEL = """sentence,lab1,lab2
+this is a sentence one,0,1
+this is a sentence two,1,0
+this is a sentence three,0,1
+"""
+
 TEST_JSON_DATA = """
 {"sentence": "this is a sentence one","lab":0}
 {"sentence": "this is a sentence two","lab":1}
 {"sentence": "this is a sentence three","lab":0}
 """
 
+TEST_JSON_DATA_MULTILABEL = """
+{"sentence": "this is a sentence one","lab1":0, "lab2": 1}
+{"sentence": "this is a sentence two","lab1":1, "lab2": 0}
+{"sentence": "this is a sentence three","lab1":0, "lab2": 1}
+"""
+
 TEST_JSON_DATA_FIELD = """{"data": [
 {"sentence": "this is a sentence one","lab":0},
 {"sentence": "this is a sentence two","lab":1},
 {"sentence": "this is a sentence three","lab":0}]}
 """
 
+TEST_JSON_DATA_FIELD_MULTILABEL = """{"data": [
+{"sentence": "this is a sentence one","lab1":0, "lab2": 1},
+{"sentence": "this is a sentence two","lab1":1, "lab2": 0},
+{"sentence": "this is a sentence three","lab1":0, "lab2": 1}]}
+"""
 
 TEST_DATA_FRAME_DATA = pd.DataFrame(
     {
@@ -62,81 +81,291 @@
     },
 )
 
-
 TEST_LIST_DATA = ["this is a sentence one", "this is a sentence two", "this is a sentence three"]
 TEST_LIST_TARGETS = [0, 1, 0]
 TEST_LIST_TARGETS_MULTILABEL = [[0, 1], [1, 0], [0, 1]]
 
 
-def csv_data(tmpdir):
+def csv_data(tmpdir, multilabel: bool):
     path = Path(tmpdir) / "data.csv"
-    path.write_text(TEST_CSV_DATA)
+    if multilabel:
+        path.write_text(TEST_CSV_DATA_MULTILABEL)
+    else:
+        path.write_text(TEST_CSV_DATA)
     return path
 
 
-def json_data(tmpdir):
+def json_data(tmpdir, multilabel: bool):
     path = Path(tmpdir) / "data.json"
-    path.write_text(TEST_JSON_DATA)
+    if multilabel:
+        path.write_text(TEST_JSON_DATA_MULTILABEL)
+    else:
+        path.write_text(TEST_JSON_DATA)
     return path
 
 
-def json_data_with_field(tmpdir):
+def json_data_with_field(tmpdir, multilabel: bool):
     path = Path(tmpdir) / "data.json"
-    path.write_text(TEST_JSON_DATA_FIELD)
+    if multilabel:
+        path.write_text(TEST_JSON_DATA_FIELD_MULTILABEL)
+    else:
+        path.write_text(TEST_JSON_DATA_FIELD)
+    return path
+
+
+def parquet_data(tmpdir):
+    path = Path(tmpdir) / "data.parquet"
+    TEST_DATA_FRAME_DATA.to_parquet(path)
     return path
 
 
 @pytest.mark.skipif(os.name == "nt", reason="Huggingface timing out on Windows")
 @pytest.mark.skipif(not _TEXT_TESTING, reason="text libraries aren't installed.")
 def test_from_csv(tmpdir):
-    csv_path = csv_data(tmpdir)
-    dm = TextClassificationData.from_csv("sentence", "label", backbone=TEST_BACKBONE, train_file=csv_path, batch_size=1)
+    csv_path = csv_data(tmpdir, multilabel=False)
+    dm = TextClassificationData.from_csv(
+        "sentence",
+        "label",
+        backbone=TEST_BACKBONE,
+        train_file=csv_path,
+        val_file=csv_path,
+        test_file=csv_path,
+        predict_file=csv_path,
+        batch_size=1,
+    )
+
     batch = next(iter(dm.train_dataloader()))
-    assert batch["labels"].item() in [0, 1]
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.val_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.test_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.predict_dataloader()))
     assert "input_ids" in batch
 
 
 @pytest.mark.skipif(os.name == "nt", reason="Huggingface timing out on Windows")
 @pytest.mark.skipif(not _TEXT_TESTING, reason="text libraries aren't installed.")
-def test_test_valid(tmpdir):
-    csv_path = csv_data(tmpdir)
+def test_from_csv_multilabel(tmpdir):
+    csv_path = csv_data(tmpdir, multilabel=True)
     dm = TextClassificationData.from_csv(
         "sentence",
-        "label",
+        ["lab1", "lab2"],
         backbone=TEST_BACKBONE,
         train_file=csv_path,
         val_file=csv_path,
         test_file=csv_path,
+        predict_file=csv_path,
         batch_size=1,
     )
+
+    batch = next(iter(dm.train_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
     batch = next(iter(dm.val_dataloader()))
-    assert batch["labels"].item() in [0, 1]
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
     assert "input_ids" in batch
 
     batch = next(iter(dm.test_dataloader()))
-    assert batch["labels"].item() in [0, 1]
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.predict_dataloader()))
     assert "input_ids" in batch
 
 
 @pytest.mark.skipif(os.name == "nt", reason="Huggingface timing out on Windows")
 @pytest.mark.skipif(not _TEXT_TESTING, reason="text libraries aren't installed.")
 def test_from_json(tmpdir):
-    json_path = json_data(tmpdir)
-    dm = TextClassificationData.from_json("sentence", "lab", backbone=TEST_BACKBONE, train_file=json_path, batch_size=1)
+    json_path = json_data(tmpdir, multilabel=False)
+    dm = TextClassificationData.from_json(
+        "sentence",
+        "lab",
+        backbone=TEST_BACKBONE,
+        train_file=json_path,
+        val_file=json_path,
+        test_file=json_path,
+        predict_file=json_path,
+        batch_size=1,
+    )
+
     batch = next(iter(dm.train_dataloader()))
-    assert batch["labels"].item() in [0, 1]
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.val_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.test_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.predict_dataloader()))
+    assert "input_ids" in batch
+
+
+@pytest.mark.skipif(os.name == "nt", reason="Huggingface timing out on Windows")
+@pytest.mark.skipif(not _TEXT_TESTING, reason="text libraries aren't installed.")
+def test_from_json_multilabel(tmpdir):
+    json_path = json_data(tmpdir, multilabel=True)
+    dm = TextClassificationData.from_json(
+        "sentence",
+        ["lab1", "lab2"],
+        backbone=TEST_BACKBONE,
+        train_file=json_path,
+        val_file=json_path,
+        test_file=json_path,
+        predict_file=json_path,
+        batch_size=1,
+    )
+
+    batch = next(iter(dm.train_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.val_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.test_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.predict_dataloader()))
     assert "input_ids" in batch
 
 
 @pytest.mark.skipif(os.name == "nt", reason="Huggingface timing out on Windows")
 @pytest.mark.skipif(not _TEXT_TESTING, reason="text libraries aren't installed.")
 def test_from_json_with_field(tmpdir):
-    json_path = json_data_with_field(tmpdir)
+    json_path = json_data_with_field(tmpdir, multilabel=False)
     dm = TextClassificationData.from_json(
-        "sentence", "lab", backbone=TEST_BACKBONE, train_file=json_path, batch_size=1, field="data"
+        "sentence",
+        "lab",
+        backbone=TEST_BACKBONE,
+        train_file=json_path,
+        val_file=json_path,
+        test_file=json_path,
+        predict_file=json_path,
+        batch_size=1,
+        field="data",
+    )
+
+    batch = next(iter(dm.train_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.val_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.test_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.predict_dataloader()))
+    assert "input_ids" in batch
+
+
+@pytest.mark.skipif(os.name == "nt", reason="Huggingface timing out on Windows")
+@pytest.mark.skipif(not _TEXT_TESTING, reason="text libraries aren't installed.")
+def test_from_json_with_field_multilabel(tmpdir):
+    json_path = json_data_with_field(tmpdir, multilabel=True)
+    dm = TextClassificationData.from_json(
+        "sentence",
+        ["lab1", "lab2"],
+        backbone=TEST_BACKBONE,
+        train_file=json_path,
+        val_file=json_path,
+        test_file=json_path,
+        predict_file=json_path,
+        batch_size=1,
+        field="data",
+    )
+
+    batch = next(iter(dm.train_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.val_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.test_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.predict_dataloader()))
+    assert "input_ids" in batch
+
+
+@pytest.mark.skipif(os.name == "nt", reason="Huggingface timing out on Windows")
+@pytest.mark.skipif(not _TEXT_TESTING, reason="text libraries aren't installed.")
+def test_from_parquet(tmpdir):
+    parquet_path = parquet_data(tmpdir)
+    dm = TextClassificationData.from_parquet(
+        "sentence",
+        "lab1",
+        backbone=TEST_BACKBONE,
+        train_file=parquet_path,
+        val_file=parquet_path,
+        test_file=parquet_path,
+        predict_file=parquet_path,
+        batch_size=1,
+    )
+
+    batch = next(iter(dm.train_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.val_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.test_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.predict_dataloader()))
+    assert "input_ids" in batch
+
+
+@pytest.mark.skipif(os.name == "nt", reason="Huggingface timing out on Windows")
+@pytest.mark.skipif(not _TEXT_TESTING, reason="text libraries aren't installed.")
+def test_from_parquet_multilabel(tmpdir):
+    parquet_path = parquet_data(tmpdir)
+    dm = TextClassificationData.from_parquet(
+        "sentence",
+        ["lab1", "lab2"],
+        backbone=TEST_BACKBONE,
+        train_file=parquet_path,
+        val_file=parquet_path,
+        test_file=parquet_path,
+        predict_file=parquet_path,
+        batch_size=1,
     )
+
     batch = next(iter(dm.train_dataloader()))
-    assert batch["labels"].item() in [0, 1]
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.val_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.test_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.predict_dataloader()))
     assert "input_ids" in batch
 
 
@@ -144,10 +373,29 @@ def test_from_json_with_field(tmpdir):
 @pytest.mark.skipif(not _TEXT_TESTING, reason="text libraries aren't installed.")
 def test_from_data_frame():
     dm = TextClassificationData.from_data_frame(
-        "sentence", "lab1", backbone=TEST_BACKBONE, train_data_frame=TEST_DATA_FRAME_DATA, batch_size=1
+        "sentence",
+        "lab1",
+        backbone=TEST_BACKBONE,
+        train_data_frame=TEST_DATA_FRAME_DATA,
+        val_data_frame=TEST_DATA_FRAME_DATA,
+        test_data_frame=TEST_DATA_FRAME_DATA,
+        predict_data_frame=TEST_DATA_FRAME_DATA,
+        batch_size=1,
     )
+
     batch = next(iter(dm.train_dataloader()))
-    assert batch["labels"].item() in [0, 1]
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.val_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.test_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.predict_dataloader()))
     assert "input_ids" in batch
 
 
@@ -155,10 +403,91 @@ def test_from_data_frame():
 @pytest.mark.skipif(not _TEXT_TESTING, reason="text libraries aren't installed.")
 def test_from_data_frame_multilabel():
     dm = TextClassificationData.from_data_frame(
-        "sentence", ["lab1", "lab2"], backbone=TEST_BACKBONE, train_data_frame=TEST_DATA_FRAME_DATA, batch_size=1
+        "sentence",
+        ["lab1", "lab2"],
+        backbone=TEST_BACKBONE,
+        train_data_frame=TEST_DATA_FRAME_DATA,
+        val_data_frame=TEST_DATA_FRAME_DATA,
+        test_data_frame=TEST_DATA_FRAME_DATA,
+        predict_data_frame=TEST_DATA_FRAME_DATA,
+        batch_size=1,
     )
+
     batch = next(iter(dm.train_dataloader()))
-    assert all([label in [0, 1] for label in batch["labels"][0]])
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.val_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.test_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.predict_dataloader()))
+    assert "input_ids" in batch
+
+
+@pytest.mark.skipif(os.name == "nt", reason="Huggingface timing out on Windows")
+@pytest.mark.skipif(not _TEXT_TESTING, reason="text libraries aren't installed.")
+def test_from_hf_datasets():
+    TEST_HF_DATASET_DATA = Dataset.from_pandas(TEST_DATA_FRAME_DATA)
+    dm = TextClassificationData.from_hf_datasets(
+        "sentence",
+        "lab1",
+        backbone=TEST_BACKBONE,
+        train_hf_dataset=TEST_HF_DATASET_DATA,
+        val_hf_dataset=TEST_HF_DATASET_DATA,
+        test_hf_dataset=TEST_HF_DATASET_DATA,
+        predict_hf_dataset=TEST_HF_DATASET_DATA,
+        batch_size=1,
+    )
+
+    batch = next(iter(dm.train_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.val_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.test_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.predict_dataloader()))
+    assert "input_ids" in batch
+
+
+@pytest.mark.skipif(os.name == "nt", reason="Huggingface timing out on Windows")
+@pytest.mark.skipif(not _TEXT_TESTING, reason="text libraries aren't installed.")
+def test_from_hf_datasets_multilabel():
+    TEST_HF_DATASET_DATA = Dataset.from_pandas(TEST_DATA_FRAME_DATA)
+    dm = TextClassificationData.from_hf_datasets(
+        "sentence",
+        ["lab1", "lab2"],
+        backbone=TEST_BACKBONE,
+        train_hf_dataset=TEST_HF_DATASET_DATA,
+        val_hf_dataset=TEST_HF_DATASET_DATA,
+        test_hf_dataset=TEST_HF_DATASET_DATA,
+        predict_hf_dataset=TEST_HF_DATASET_DATA,
+        batch_size=1,
+    )
+
+    batch = next(iter(dm.train_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.val_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.test_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.predict_dataloader()))
     assert "input_ids" in batch
 
 
@@ -166,10 +495,30 @@ def test_from_data_frame_multilabel():
 @pytest.mark.skipif(not _TEXT_TESTING, reason="text libraries aren't installed.")
 def test_from_lists():
     dm = TextClassificationData.from_lists(
-        backbone=TEST_BACKBONE, train_data=TEST_LIST_DATA, train_targets=TEST_LIST_TARGETS, batch_size=1
+        backbone=TEST_BACKBONE,
+        train_data=TEST_LIST_DATA,
+        train_targets=TEST_LIST_TARGETS,
+        val_data=TEST_LIST_DATA,
+        val_targets=TEST_LIST_TARGETS,
+        test_data=TEST_LIST_DATA,
+        test_targets=TEST_LIST_TARGETS,
+        predict_data=TEST_LIST_DATA,
+        batch_size=1,
     )
+
     batch = next(iter(dm.train_dataloader()))
-    assert batch["labels"].item() in [0, 1]
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.val_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.test_dataloader()))
+    assert batch[DefaultDataKeys.TARGET].item() in [0, 1]
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.predict_dataloader()))
     assert "input_ids" in batch
 
 
@@ -177,10 +526,30 @@ def test_from_lists():
 @pytest.mark.skipif(not _TEXT_TESTING, reason="text libraries aren't installed.")
 def test_from_lists_multilabel():
     dm = TextClassificationData.from_lists(
-        backbone=TEST_BACKBONE, train_data=TEST_LIST_DATA, train_targets=TEST_LIST_TARGETS_MULTILABEL, batch_size=1
+        backbone=TEST_BACKBONE,
+        train_data=TEST_LIST_DATA,
+        train_targets=TEST_LIST_TARGETS_MULTILABEL,
+        val_data=TEST_LIST_DATA,
+        val_targets=TEST_LIST_TARGETS_MULTILABEL,
+        test_data=TEST_LIST_DATA,
+        test_targets=TEST_LIST_TARGETS_MULTILABEL,
+        predict_data=TEST_LIST_DATA,
+        batch_size=1,
     )
+
     batch = next(iter(dm.train_dataloader()))
-    assert all([label in [0, 1] for label in batch["labels"][0]])
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.val_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.test_dataloader()))
+    assert all([label in [0, 1] for label in batch[DefaultDataKeys.TARGET][0]])
+    assert "input_ids" in batch
+
+    batch = next(iter(dm.predict_dataloader()))
     assert "input_ids" in batch
 
 
@@ -196,17 +565,17 @@ def test_text_module_not_found_error():
     "cls, kwargs",
     [
         (TextDataSource, {}),
-        (TextFileDataSource, {"filetype": "csv"}),
         (TextCSVDataSource, {}),
         (TextJSONDataSource, {}),
         (TextDataFrameDataSource, {}),
+        (TextParquetDataSource, {}),
+        (TextHuggingFaceDatasetDataSource, {}),
         (TextListDataSource, {}),
-        (TextSentencesDataSource, {}),
     ],
 )
 def test_tokenizer_state(cls, kwargs):
     """Tests that the tokenizer is not in __getstate__"""
-    instance = cls(backbone="sshleifer/tiny-mbart", **kwargs)
+    instance = cls(backbone=TEST_BACKBONE, **kwargs)
     state = instance.__getstate__()
     tokenizers = []
     for name, attribute in instance.__dict__.items():
diff --git a/tests/text/classification/test_model.py b/tests/text/classification/test_model.py
--- a/tests/text/classification/test_model.py
+++ b/tests/text/classification/test_model.py
@@ -20,6 +20,7 @@
 
 from flash import Trainer
 from flash.__main__ import main
+from flash.core.data.data_source import DefaultDataKeys
 from flash.core.utilities.imports import _TEXT_AVAILABLE
 from flash.text import TextClassifier
 from flash.text.classification.data import TextClassificationPostprocess, TextClassificationPreprocess
@@ -32,7 +33,7 @@ class DummyDataset(torch.utils.data.Dataset):
     def __getitem__(self, index):
         return {
             "input_ids": torch.randint(1000, size=(100,)),
-            "labels": torch.randint(2, size=(1,)).item(),
+            DefaultDataKeys.TARGET: torch.randint(2, size=(1,)).item(),
         }
 
     def __len__(self) -> int: