Repository: Lightning-Universe/lightning-transformers (language: python)

Issue:
Move away from configs, to just passing args to the class
## 🚀 Feature

Currently, the process of creating a module looks like such:

```python
...
dm = TextClassificationDataModule(
    cfg=TextClassificationDataConfig(
        batch_size=1,
        dataset_name="emotion",
        max_length=512,
    ),
    tokenizer=tokenizer,
)
```

This is fine for the "general" cases, however for the specific modules for datasets, it looks cumbersome:

```python
...
dm = WMT16TranslationDataModule( # I have to choose the WMT16 class
    cfg=TranslationDataConfig(
        dataset_name="wmt16", # I have to pass the name of the dataset as well? why is this not the default?
        ...
    ),
    tokenizer=tokenizer,
)
```

A solution would be to introduce a config class like below, however this adds even more lines.

```python
...
@dataclass
class WMT16TranslationDataConfig:
    dataset_name: str = "wmt16"
```

I suggest we opt to remove configs entirely, and just rely on the modules:

```python
dm = TextClassificationDataModule(
    batch_size=1,
    dataset_name="emotion",
    max_length=512,
    tokenizer=tokenizer,
)
dm = WMT16TranslationDataModule(tokenizer=tokenizer) # dataset_name="wmt16" is default!
```

From my understanding, this is supported by hydra, but all `cfg` objects will be converted into parameters.


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 49e2563a..60dbb062 100644
--- a/README.md
+++ b/README.md
@@ -49,18 +49,15 @@ from transformers import AutoTokenizer
 from lightning_transformers.task.nlp.text_classification import (
     TextClassificationDataModule,
     TextClassificationTransformer,
-    TextClassificationDataConfig,
 )
 
 tokenizer = AutoTokenizer.from_pretrained(
     pretrained_model_name_or_path="bert-base-cased"
 )
 dm = TextClassificationDataModule(
-    cfg=TextClassificationDataConfig(
-        batch_size=1,
-        dataset_name="emotion",
-        max_length=512,
-    ),
+    batch_size=1,
+    dataset_name="emotion",
+    max_length=512,
     tokenizer=tokenizer,
 )
 model = TextClassificationTransformer(
@@ -81,8 +78,6 @@ from transformers import AutoTokenizer
 from lightning_transformers.task.nlp.translation import (
     TranslationTransformer,
     WMT16TranslationDataModule,
-    TranslationConfig,
-    TranslationDataConfig,
 )
 
 tokenizer = AutoTokenizer.from_pretrained(
@@ -90,24 +85,19 @@ tokenizer = AutoTokenizer.from_pretrained(
 )
 model = TranslationTransformer(
     pretrained_model_name_or_path="google/mt5-base",
-    cfg=TranslationConfig(
-        n_gram=4,
-        smooth=False,
-        val_target_max_length=142,
-        num_beams=None,
-        compute_generate_metrics=True,
-    ),
+    n_gram=4,
+    smooth=False,
+    val_target_max_length=142,
+    num_beams=None,
+    compute_generate_metrics=True,
 )
 dm = WMT16TranslationDataModule(
-    cfg=TranslationDataConfig(
-        dataset_name="wmt16",
-        # WMT translation datasets: ['cs-en', 'de-en', 'fi-en', 'ro-en', 'ru-en', 'tr-en']
-        dataset_config_name="ro-en",
-        source_language="en",
-        target_language="ro",
-        max_source_length=128,
-        max_target_length=128,
-    ),
+    # WMT translation datasets: ['cs-en', 'de-en', 'fi-en', 'ro-en', 'ru-en', 'tr-en']
+    dataset_config_name="ro-en",
+    source_language="en",
+    target_language="ro",
+    max_source_length=128,
+    max_target_length=128,
     tokenizer=tokenizer,
 )
 trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
diff --git a/docs/source/advanced/nlp/language_modeling_data.rst b/docs/source/advanced/nlp/language_modeling_data.rst
index f1945b65..1cc705db 100644
--- a/docs/source/advanced/nlp/language_modeling_data.rst
+++ b/docs/source/advanced/nlp/language_modeling_data.rst
@@ -12,9 +12,6 @@ The base data module can be used to modify this code, and follows a simple patte
 
     class LanguageModelingDataModule(HFDataModule):
 
-        def __init__(self, cfg: LanguageModelingDataConfig = LanguageModelingDataConfig()):
-            super().__init__(cfg=cfg)
-
         def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset:
             # `process_data` converting the dataset into features.
             # The dataset is pre-loaded using `load_dataset`.
@@ -63,14 +60,13 @@ Below we have the pseudo code version to show where most of the changes happened
     from datasets import Dataset, Optional
     from transformers import PreTrainedTokenizerBase
 
-    from lightning_transformers.core.nlp.huggingface import HFTransformerDataConfig
     from lightning_transformers.task.nlp.language_modeling import LanguageModelingDataModule
 
 
     class MyLanguageModelingDataModule(LanguageModelingDataModule):
 
-        def __init__(self, cfg: HFTransformerDataConfig, tokenizer: PreTrainedTokenizerBase):
-            super().__init__(cfg, tokenizer)
+        def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs):
+            super().__init__(tokenizer, *args, **kwargs)
             self.tokenized_condition_term = tokenizer("This is a story: ")
 
         def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset:
diff --git a/docs/source/advanced/nlp/translation_data.rst b/docs/source/advanced/nlp/translation_data.rst
index f2dd0615..9c655da9 100644
--- a/docs/source/advanced/nlp/translation_data.rst
+++ b/docs/source/advanced/nlp/translation_data.rst
@@ -12,9 +12,6 @@ The base data module can be used to modify this code, and follows a simple patte
 
     class TranslationDataModule(Seq2SeqDataModule):
 
-        def __init__(self, cfg: TranslationDataConfig = TranslationDataConfig()):
-            super().__init__(cfg=cfg)
-
         @property
         def source_target_column_names(self) -> Tuple[str, str]:
             return self.cfg.source_language, self.cfg.target_language
@@ -23,9 +20,6 @@ The base data module can be used to modify this code, and follows a simple patte
 
     class Seq2SeqDataModule(HFDataModule):
 
-        def __init__(self, cfg: Seq2SeqDataConfig = Seq2SeqDataConfig()):
-            super().__init__(cfg=cfg)
-
         def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset:
             # `process_data` converting the dataset into features.
             # The dataset is pre-loaded using `load_dataset`.
diff --git a/docs/source/datasets/nlp/custom_subset_names.rst b/docs/source/datasets/nlp/custom_subset_names.rst
index 96853ab4..6bd7fcdc 100644
--- a/docs/source/datasets/nlp/custom_subset_names.rst
+++ b/docs/source/datasets/nlp/custom_subset_names.rst
@@ -10,18 +10,15 @@ An example for how to train and validate on MNLI would the the following:
 .. code-block:: python
 
     from lightning_transformers.task.nlp.text_classification import (
-        TextClassificationDataConfig,
         TextClassificationDataModule,
         TextClassificationTransformer,
     )
 
     dm = TextClassificationDataModule(
-        cfg=TextClassificationDataConfig(
-            batch_size=1,
-            dataset_name="glue",
-            dataset_config_name="mnli",
-            max_length=512,
-            validation_subset_name="validation_matched"
-        ),
+        batch_size=1,
+        dataset_name="glue",
+        dataset_config_name="mnli",
+        max_length=512,
+        validation_subset_name="validation_matched"
         tokenizer=tokenizer,
     )
diff --git a/docs/source/datasets/nlp/language_modeling_data.rst b/docs/source/datasets/nlp/language_modeling_data.rst
index 751a2156..82287346 100644
--- a/docs/source/datasets/nlp/language_modeling_data.rst
+++ b/docs/source/datasets/nlp/language_modeling_data.rst
@@ -17,15 +17,12 @@ Below we have defined a csv file to use as our input data.
 .. code-block:: python
 
     from lightning_transformers.task.nlp.language_modeling import (
-        LanguageModelingDataConfig,
         LanguageModelingDataModule,
     )
 
     dm = LanguageModelingDataModule(
-        cfg=LanguageModelingDataConfig(
-            batch_size=1,
-            train_file="path/train.csv",
-            validation_file="/path/valid.csv"
-        ),
+        batch_size=1,
+        train_file="path/train.csv",
+        validation_file="/path/valid.csv"
         tokenizer=tokenizer,
     )
diff --git a/docs/source/datasets/nlp/multiple_choice_data.rst b/docs/source/datasets/nlp/multiple_choice_data.rst
index aaa6023f..bbdad135 100644
--- a/docs/source/datasets/nlp/multiple_choice_data.rst
+++ b/docs/source/datasets/nlp/multiple_choice_data.rst
@@ -22,18 +22,14 @@ We override the dataset files, allowing us to still use the data transforms defi
 .. code-block:: python
 
     from lightning_transformers.task.nlp.multiple_choice import (
-        MultipleChoiceDataConfig,
         RaceMultipleChoiceDataModule,
     )
 
     dm = RaceMultipleChoiceDataModule(
-        cfg=MultipleChoiceDataConfig(
-            batch_size=1,
-            dataset_name="race",
-            dataset_config_name="all",
-            padding=False,
-            train_file="path/train.json",
-            validation_file="/path/valid.json"
-        ),
+        batch_size=1,
+        dataset_config_name="all",
+        padding=False,
+        train_file="path/train.json",
+        validation_file="/path/valid.json"
         tokenizer=tokenizer,
     )
diff --git a/docs/source/datasets/nlp/question_answering_data.rst b/docs/source/datasets/nlp/question_answering_data.rst
index fc270ff0..80794503 100644
--- a/docs/source/datasets/nlp/question_answering_data.rst
+++ b/docs/source/datasets/nlp/question_answering_data.rst
@@ -24,23 +24,19 @@ We override the dataset files, allowing us to still use the data transforms defi
 
 
     from lightning_transformers.task.nlp.question_answering import (
-        QuestionAnsweringDataConfig,
         SquadDataModule,
     )
 
     dm = SquadDataModule(
-        cfg=QuestionAnsweringDataConfig(
-            batch_size=1,
-            dataset_name="squad",
-            dataset_config_name="plain_text",
-            max_length=384,
-            version_2_with_negative=False,
-            null_score_diff_threshold=0.0,
-            doc_stride=128,
-            n_best_size=20,
-            max_answer_length=30,
-            train_file="path/train.csv",
-            validation_file="/path/valid.csv"
-        ),
+        batch_size=1,
+        dataset_config_name="plain_text",
+        max_length=384,
+        version_2_with_negative=False,
+        null_score_diff_threshold=0.0,
+        doc_stride=128,
+        n_best_size=20,
+        max_answer_length=30,
+        train_file="path/train.csv",
+        validation_file="/path/valid.csv"
         tokenizer=tokenizer,
     )
diff --git a/docs/source/datasets/nlp/summarization_data.rst b/docs/source/datasets/nlp/summarization_data.rst
index 388ce27c..b88c5a3a 100644
--- a/docs/source/datasets/nlp/summarization_data.rst
+++ b/docs/source/datasets/nlp/summarization_data.rst
@@ -15,19 +15,14 @@ We override the dataset files, allowing us to still use the data transforms defi
 .. code-block:: python
 
     from lightning_transformers.task.nlp.summarization import (
-        SummarizationConfig,
-        SummarizationDataConfig,
         XsumSummarizationDataModule,
     )
 
     dm = XsumSummarizationDataModule(
-        cfg=SummarizationDataConfig(
-            batch_size=1,
-            dataset_name="xsum",
-            max_source_length=128,
-            max_target_length=128,
-            train_file="path/train.csv",
-            validation_file="/path/valid.csv"
-        ),
+        batch_size=1,
+        max_source_length=128,
+        max_target_length=128,
+        train_file="path/train.csv",
+        validation_file="/path/valid.csv"
         tokenizer=tokenizer,
     )
diff --git a/docs/source/datasets/nlp/text_classification_data.rst b/docs/source/datasets/nlp/text_classification_data.rst
index 529ff6a8..5bf6b8de 100644
--- a/docs/source/datasets/nlp/text_classification_data.rst
+++ b/docs/source/datasets/nlp/text_classification_data.rst
@@ -15,17 +15,14 @@ The label mapping is automatically generated from the training dataset labels if
 .. code-block:: python
 
     from lightning_transformers.task.nlp.text_classification import (
-        TextClassificationDataConfig,
         TextClassificationDataModule,
         TextClassificationTransformer,
     )
 
     dm = TextClassificationDataModule(
-        cfg=TextClassificationDataConfig(
-            batch_size=1,
-            max_length=512,
-            train_file="path/train.json",
-            validation_file="/path/valid.json"
-        ),
+        batch_size=1,
+        max_length=512,
+        train_file="path/train.json",
+        validation_file="/path/valid.json"
         tokenizer=tokenizer,
     )
diff --git a/docs/source/datasets/nlp/token_classification_data.rst b/docs/source/datasets/nlp/token_classification_data.rst
index 5835334a..15514a83 100644
--- a/docs/source/datasets/nlp/token_classification_data.rst
+++ b/docs/source/datasets/nlp/token_classification_data.rst
@@ -12,21 +12,16 @@ To use custom text files, the files should contain new line delimited json objec
 
 .. code-block:: python
 
-    from lightning_transformers.task.nlp.token_classification import (
-        TokenClassificationDataConfig,
-        TokenClassificationDataModule,
-    )
+    from lightning_transformers.task.nlp.token_classification import TokenClassificationDataModule
 
     dm = TokenClassificationDataModule(
-        cfg=TokenClassificationDataConfig(
-            batch_size=1,
-            task_name="ner",
-            dataset_name="conll2003",
-            preprocessing_num_workers=1,
-            label_all_tokens=False,
-            revision="master",
-            train_file="path/train.json",
-            validation_file="/path/valid.json"
-        ),
+        batch_size=1,
+        task_name="ner",
+        dataset_name="conll2003",
+        preprocessing_num_workers=1,
+        label_all_tokens=False,
+        revision="master",
+        train_file="path/train.json",
+        validation_file="/path/valid.json"
         tokenizer=tokenizer,
     )
diff --git a/docs/source/datasets/nlp/translation_data.rst b/docs/source/datasets/nlp/translation_data.rst
index a8e51a9a..09672bb8 100644
--- a/docs/source/datasets/nlp/translation_data.rst
+++ b/docs/source/datasets/nlp/translation_data.rst
@@ -14,22 +14,16 @@ We override the dataset files, allowing us to still use the data transforms defi
 
 .. code-block:: python
 
-    from lightning_transformers.task.nlp.translation import (
-        TranslationDataConfig,
-        WMT16TranslationDataModule,
-    )
+    from lightning_transformers.task.nlp.translation import WMT16TranslationDataModule
 
     dm = WMT16TranslationDataModule(
-        cfg=TranslationDataConfig(
-            dataset_name="wmt16",
-            # WMT translation datasets: ['cs-en', 'de-en', 'fi-en', 'ro-en', 'ru-en', 'tr-en']
-            dataset_config_name="ro-en",
-            source_language="en",
-            target_language="ro",
-            max_source_length=128,
-            max_target_length=128,
-            train_file="path/train.json",
-            validation_file="/path/valid.json"
-        ),
+        # WMT translation datasets: ['cs-en', 'de-en', 'fi-en', 'ro-en', 'ru-en', 'tr-en']
+        dataset_config_name="ro-en",
+        source_language="en",
+        target_language="ro",
+        max_source_length=128,
+        max_target_length=128,
+        train_file="path/train.json",
+        validation_file="/path/valid.json"
         tokenizer=tokenizer,
     )
diff --git a/docs/source/optimizations/hf_save.rst b/docs/source/features/hf_save.rst
similarity index 87%
rename from docs/source/optimizations/hf_save.rst
rename to docs/source/features/hf_save.rst
index 60b45985..f7ccd4a1 100644
--- a/docs/source/optimizations/hf_save.rst
+++ b/docs/source/features/hf_save.rst
@@ -27,19 +27,16 @@ To save an additional HF Checkpoint everytime the checkpoint callback saves, pas
 
     from lightning_transformers.plugins.checkpoint import HFSaveCheckpoint
     from lightning_transformers.task.nlp.text_classification import (
-        TextClassificationDataConfig,
         TextClassificationDataModule,
         TextClassificationTransformer,
     )
 
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     dm = TextClassificationDataModule(
-        cfg=TextClassificationDataConfig(
-            batch_size=1,
-            dataset_name="glue",
-            dataset_config_name="sst2",
-            max_length=512,
-        ),
+        batch_size=1,
+        dataset_name="glue",
+        dataset_config_name="sst2",
+        max_length=512,
         tokenizer=tokenizer,
     )
     model = TextClassificationTransformer(pretrained_model_name_or_path="prajjwal1/bert-tiny")
diff --git a/docs/source/optimizations/large_model.rst b/docs/source/features/large_model.rst
similarity index 96%
rename from docs/source/optimizations/large_model.rst
rename to docs/source/features/large_model.rst
index fb94fd6b..ba622edb 100644
--- a/docs/source/optimizations/large_model.rst
+++ b/docs/source/features/large_model.rst
@@ -1,7 +1,7 @@
 .. _large_model:
 
-Inference for Big Transformers
-==============================
+Big Transformer Models Inference
+================================
 
 Lightning Transformers provides out of the box support for running inference with very large billion parameter models. Under-the-hood we use HF Accelerates' Transformer support to auto-select devices for optimal throughput and memory usage.
 
diff --git a/docs/source/features/sparseml.rst b/docs/source/features/sparseml.rst
new file mode 100644
index 00000000..6a403e76
--- /dev/null
+++ b/docs/source/features/sparseml.rst
@@ -0,0 +1,27 @@
+.. _sparseml:
+
+SparseML
+========
+
+`SparseML <https://github.com/neuralmagic/sparseml>`__ provides GPU-class performance on CPUs through sparsification, pruning, and quantization.
+For more details, see `SparseML docs <https://docs.neuralmagic.com/sparseml/>`__.
+
+With multiple machines, the command has to be run on all machines either manually, or using an orchestration system such as SLURM or TorchElastic. More information can be seen in the Pytorch Lightning `Computing Cluster <https://pytorch-lightning.readthedocs.io/en/latest/advanced/cluster.html#computing-cluster>`_.
+
+We provide out of the box configs to use SparseML. Just pass the SparseML Callback when training.
+
+.. code-block:: python
+
+    import pytorch_lightning as pl
+    from lightning_transformers.callbacks import TransformerSparseMLCallback
+
+    pl.Trainer(
+        callbacks=TransformerSparseMLCallback(
+            output_dir="/content/MODELS",
+            recipe_path="/content/recipe.yaml"
+        )
+    )
+
+These commands are only useful when a recipe has already been created. Example recipes can be found `here <https://github.com/neuralmagic/sparseml/tree/main/integrations/huggingface-transformers/recipes>`__.
+
+After training, this will leave two ONNX models in the trainer.callbacks.output_dir folder: small_model.onnx and model.onnx. small_model.onnx is excellent for demos. For reliable inference, it is recommended to optimize model.onnx with your compression algorithm.
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 5674edb7..29986a75 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -29,12 +29,12 @@ Lightning Transformers
 
 .. toctree::
    :maxdepth: 1
-   :name: optimization
-   :caption: Optimizations
+   :name: features
+   :caption: Features
 
-   optimizations/sparseml
-   optimizations/hf_save
-   optimizations/large_model
+   features/sparseml
+   features/hf_save
+   features/large_model
 
 .. toctree::
    :maxdepth: 1
diff --git a/docs/source/optimizations/sparseml.rst b/docs/source/optimizations/sparseml.rst
deleted file mode 100644
index d076d0cc..00000000
--- a/docs/source/optimizations/sparseml.rst
+++ /dev/null
@@ -1,47 +0,0 @@
-.. _sparseml:
-
-SparseML
-========
-
-`SparseML <https://github.com/neuralmagic/sparseml>`__ provides GPU-class performance on CPUs through sparsification, pruning, and quantization.
-For more details, see `SparseML docs <https://docs.neuralmagic.com/sparseml/>`__.
-
-With multiple machines, the command has to be run on all machines either manually, or using an orchestration system such as SLURM or TorchElastic. More information can be seen in the Pytorch Lightning `Computing Cluster <https://pytorch-lightning.readthedocs.io/en/latest/advanced/cluster.html#computing-cluster>`_.
-
-GPU Config
-----------
-
-We provide out of the box configs to use SparseML. Below is an example of how you can swap to the default trainer config for SparseML when using the text classification task.
-
-.. code-block:: bash
-
-    python train.py task=nlp/text_classification                 \
-        dataset=nlp/text_classification/glue                     \
-        dataset.cfg.dataset_config_name=sst2                     \
-        backbone.pretrained_model_name_or_path=bert-base-uncased \
-        trainer=sparseml                                         \
-        ++trainer.callbacks.output_dir=/content/MODELS           \
-        ++trainer.callbacks.recipe_path=/content/recipe.yaml     \
-        log=True                                                 \
-
-These commands are only useful when a recipe has already been created. Example recipes can be found `here <https://github.com/neuralmagic/sparseml/tree/main/integrations/huggingface-transformers/recipes>`__.
-
-After training, this will leave two ONNX models in the trainer.callbacks.output_dir folder: small_model.onnx and model.onnx. small_model.onnx is excellent for demos. For reliable inference, it is recommended to optimize model.onnx with your compression algorithm.
-
-CPU Config
-----------
-
-If you want to use SparseML only as a callback, or if you would like to train/test on CPU, use a variation of the following config:
-
-.. code-block:: bash
-
-    python train.py task=nlp/text_classification                 \
-        dataset=nlp/text_classification/glue                     \
-        dataset.cfg.dataset_config_name=sst2                     \
-        backbone.pretrained_model_name_or_path=bert-base-uncased \
-        trainer.gpus=0                                           \
-        +trainer/callbacks=sparseml                              \
-        +trainer/logger=sparsewandb                              \
-        ++trainer.callbacks.output_dir=/content/MODELS           \
-        ++trainer.callbacks.recipe_path=/content/recipe.yaml     \
-        log=True                                                 \
diff --git a/docs/source/quickstart.rst b/docs/source/quickstart.rst
index d87269fa..23d6467a 100644
--- a/docs/source/quickstart.rst
+++ b/docs/source/quickstart.rst
@@ -27,18 +27,15 @@ Here is an example of training `bert-base-cased <https://huggingface.co/bert-bas
     from lightning_transformers.task.nlp.text_classification import (
         TextClassificationDataModule,
         TextClassificationTransformer,
-        TextClassificationDataConfig,
     )
 
     tokenizer = AutoTokenizer.from_pretrained(
         pretrained_model_name_or_path="bert-base-cased"
     )
     dm = TextClassificationDataModule(
-        cfg=TextClassificationDataConfig(
-            batch_size=1,
-            dataset_name="emotion",
-            max_length=512,
-        ),
+        batch_size=1,
+        dataset_name="emotion",
+        max_length=512,
         tokenizer=tokenizer,
     )
     model = TextClassificationTransformer(pretrained_model_name_or_path="bert-base-cased")
@@ -63,7 +60,6 @@ Swapping to the RMSProp optimizer:
     from lightning_transformers.task.nlp.text_classification import (
         TextClassificationDataModule,
         TextClassificationTransformer,
-        TextClassificationDataConfig,
     )
 
 
@@ -85,11 +81,9 @@ Swapping to the RMSProp optimizer:
         pretrained_model_name_or_path="bert-base-cased"
     )
     dm = TextClassificationDataModule(
-        cfg=TextClassificationDataConfig(
-            batch_size=1,
-            dataset_name="emotion",
-            max_length=512,
-        ),
+        batch_size=1,
+        dataset_name="emotion",
+        max_length=512,
         tokenizer=tokenizer,
     )
     model = RMSPropTransformer(pretrained_model_name_or_path="bert-base-cased")
diff --git a/docs/source/tasks/nlp/language_modeling.rst b/docs/source/tasks/nlp/language_modeling.rst
index c2076341..8abcbda7 100644
--- a/docs/source/tasks/nlp/language_modeling.rst
+++ b/docs/source/tasks/nlp/language_modeling.rst
@@ -28,7 +28,6 @@ Language Models pre-trained or fine-tuned to the Causal Language Modeling task c
     from transformers import AutoTokenizer
 
     from lightning_transformers.task.nlp.language_modeling import (
-        LanguageModelingDataConfig,
         LanguageModelingDataModule,
         LanguageModelingTransformer,
     )
@@ -36,11 +35,9 @@ Language Models pre-trained or fine-tuned to the Causal Language Modeling task c
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="gpt2")
     model = LanguageModelingTransformer(pretrained_model_name_or_path="gpt2")
     dm = LanguageModelingDataModule(
-        cfg=LanguageModelingDataConfig(
-            batch_size=1,
-            dataset_name="wikitext",
-            dataset_config_name="wikitext-2-raw-v1",
-        ),
+        batch_size=1,
+        dataset_name="wikitext",
+        dataset_config_name="wikitext-2-raw-v1",
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
diff --git a/docs/source/tasks/nlp/multiple_choice.rst b/docs/source/tasks/nlp/multiple_choice.rst
index a1c34d86..b48afa1e 100644
--- a/docs/source/tasks/nlp/multiple_choice.rst
+++ b/docs/source/tasks/nlp/multiple_choice.rst
@@ -34,7 +34,6 @@ Training
     from transformers import AutoTokenizer
 
     from lightning_transformers.task.nlp.multiple_choice import (
-        MultipleChoiceDataConfig,
         MultipleChoiceTransformer,
         SwagMultipleChoiceDataModule,
     )
@@ -42,12 +41,9 @@ Training
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="bert-base-uncased")
     model = MultipleChoiceTransformer(pretrained_model_name_or_path="bert-base-uncased")
     dm = SwagMultipleChoiceDataModule(
-        cfg=MultipleChoiceDataConfig(
-            batch_size=1,
-            dataset_name="swag",
-            dataset_config_name="regular",
-            padding=False,
-        ),
+        batch_size=1,
+        dataset_config_name="regular",
+        padding=False,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
diff --git a/docs/source/tasks/nlp/question_answering.rst b/docs/source/tasks/nlp/question_answering.rst
index f47413f4..d2c4efda 100644
--- a/docs/source/tasks/nlp/question_answering.rst
+++ b/docs/source/tasks/nlp/question_answering.rst
@@ -31,7 +31,6 @@ Training
     from transformers import AutoTokenizer
 
     from lightning_transformers.task.nlp.question_answering import (
-        QuestionAnsweringDataConfig,
         QuestionAnsweringTransformer,
         SquadDataModule,
     )
@@ -39,17 +38,14 @@ Training
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="bert-base-uncased")
     model = QuestionAnsweringTransformer(pretrained_model_name_or_path="bert-base-uncased")
     dm = SquadDataModule(
-        cfg=QuestionAnsweringDataConfig(
-            batch_size=1,
-            dataset_name="squad",
-            dataset_config_name="plain_text",
-            max_length=384,
-            version_2_with_negative=False,
-            null_score_diff_threshold=0.0,
-            doc_stride=128,
-            n_best_size=20,
-            max_answer_length=30,
-        ),
+        batch_size=1,
+        dataset_config_name="plain_text",
+        max_length=384,
+        version_2_with_negative=False,
+        null_score_diff_threshold=0.0,
+        doc_stride=128,
+        n_best_size=20,
+        max_answer_length=30,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
diff --git a/docs/source/tasks/nlp/summarization.rst b/docs/source/tasks/nlp/summarization.rst
index e39b2027..bf692881 100644
--- a/docs/source/tasks/nlp/summarization.rst
+++ b/docs/source/tasks/nlp/summarization.rst
@@ -30,8 +30,6 @@ To use this task, we must select a Seq2Seq Encoder/Decoder based model, such as
     from transformers import AutoTokenizer
 
     from lightning_transformers.task.nlp.summarization import (
-        SummarizationConfig,
-        SummarizationDataConfig,
         SummarizationTransformer,
         XsumSummarizationDataModule,
     )
@@ -39,20 +37,15 @@ To use this task, we must select a Seq2Seq Encoder/Decoder based model, such as
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="t5-base")
     model = SummarizationTransformer(
         pretrained_model_name_or_path="t5-base",
-        cfg=SummarizationConfig(
-            use_stemmer=True,
-            val_target_max_length=142,
-            num_beams=None,
-            compute_generate_metrics=True,
-        ),
+        use_stemmer=True,
+        val_target_max_length=142,
+        num_beams=None,
+        compute_generate_metrics=True,
     )
     dm = XsumSummarizationDataModule(
-        cfg=SummarizationDataConfig(
-            batch_size=1,
-            dataset_name="xsum",
-            max_source_length=128,
-            max_target_length=128,
-        ),
+        batch_size=1,
+        max_source_length=128,
+        max_target_length=128,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(accelerator="auto", devices=1, max_epochs=1)
diff --git a/docs/source/tasks/nlp/text_classification.rst b/docs/source/tasks/nlp/text_classification.rst
index 68c13174..2bf24365 100644
--- a/docs/source/tasks/nlp/text_classification.rst
+++ b/docs/source/tasks/nlp/text_classification.rst
@@ -28,19 +28,16 @@ For this task, you can rely on most Transformer models as your backbone.
     from transformers import AutoTokenizer
 
     from lightning_transformers.task.nlp.text_classification import (
-        TextClassificationDataConfig,
         TextClassificationDataModule,
         TextClassificationTransformer,
     )
 
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="bert-base-uncased")
     dm = TextClassificationDataModule(
-        cfg=TextClassificationDataConfig(
-            batch_size=1,
-            dataset_name="glue",
-            dataset_config_name="sst2",
-            max_length=512,
-        ),
+        batch_size=1,
+        dataset_name="glue",
+        dataset_config_name="sst2",
+        max_length=512,
         tokenizer=tokenizer,
     )
     model = TextClassificationTransformer(pretrained_model_name_or_path="bert-base-uncased", num_labels=dm.num_classes)
diff --git a/docs/source/tasks/nlp/token_classification.rst b/docs/source/tasks/nlp/token_classification.rst
index 9086aa3e..0d60ffc7 100644
--- a/docs/source/tasks/nlp/token_classification.rst
+++ b/docs/source/tasks/nlp/token_classification.rst
@@ -21,21 +21,18 @@ Training
     from transformers import AutoTokenizer
 
     from lightning_transformers.task.nlp.token_classification import (
-        TokenClassificationDataConfig,
         TokenClassificationDataModule,
         TokenClassificationTransformer,
     )
 
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="bert-base-uncased")
     dm = TokenClassificationDataModule(
-        cfg=TokenClassificationDataConfig(
-            batch_size=1,
-            task_name="ner",
-            dataset_name="conll2003",
-            preprocessing_num_workers=1,
-            label_all_tokens=False,
-            revision="master",
-        ),
+        batch_size=1,
+        task_name="ner",
+        dataset_name="conll2003",
+        preprocessing_num_workers=1,
+        label_all_tokens=False,
+        revision="master",
         tokenizer=tokenizer,
     )
     model = TokenClassificationTransformer(pretrained_model_name_or_path="bert-base-uncased", labels=dm.labels)
diff --git a/docs/source/tasks/nlp/translation.rst b/docs/source/tasks/nlp/translation.rst
index eaa2f75f..5269ecbd 100644
--- a/docs/source/tasks/nlp/translation.rst
+++ b/docs/source/tasks/nlp/translation.rst
@@ -29,8 +29,6 @@ In addition, you also need a tokenizer that has been created on multi-lingual te
     from transformers import AutoTokenizer
 
     from lightning_transformers.task.nlp.translation import (
-        TranslationConfig,
-        TranslationDataConfig,
         TranslationTransformer,
         WMT16TranslationDataModule,
     )
@@ -38,24 +36,19 @@ In addition, you also need a tokenizer that has been created on multi-lingual te
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="t5-base")
     model = TranslationTransformer(
         pretrained_model_name_or_path="t5-base",
-        cfg=TranslationConfig(
-            n_gram=4,
-            smooth=False,
-            val_target_max_length=142,
-            num_beams=None,
-            compute_generate_metrics=True,
-        ),
+        n_gram=4,
+        smooth=False,
+        val_target_max_length=142,
+        num_beams=None,
+        compute_generate_metrics=True,
     )
     dm = WMT16TranslationDataModule(
-        cfg=TranslationDataConfig(
-            dataset_name="wmt16",
-            # WMT translation datasets: ['cs-en', 'de-en', 'fi-en', 'ro-en', 'ru-en', 'tr-en']
-            dataset_config_name="ro-en",
-            source_language="en",
-            target_language="ro",
-            max_source_length=128,
-            max_target_length=128,
-        ),
+        # WMT translation datasets: ['cs-en', 'de-en', 'fi-en', 'ro-en', 'ru-en', 'tr-en']
+        dataset_config_name="ro-en",
+        source_language="en",
+        target_language="ro",
+        max_source_length=128,
+        max_target_length=128,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
diff --git a/docs/source/tasks/vision/image_classification.rst b/docs/source/tasks/vision/image_classification.rst
index 68a58907..d1150cb1 100644
--- a/docs/source/tasks/vision/image_classification.rst
+++ b/docs/source/tasks/vision/image_classification.rst
@@ -23,14 +23,15 @@ Language Models pre-trained or fine-tuned to the Causal Language Modeling task c
     from transformers import AutoFeatureExtractor
 
     from lightning_transformers.task.vision.image_classification import (
-        ImageClassificationDataConfig,
         ImageClassificationDataModule,
         ImageClassificationTransformer,
     )
 
     feature_extractor = AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path="nateraw/vit-base-beans")
     dm = ImageClassificationDataModule(
-        cfg=ImageClassificationDataConfig(batch_size=8, dataset_name="beans", num_workers=8),
+        batch_size=8,
+        dataset_name="beans",
+        num_workers=8,
         feature_extractor=feature_extractor,
     )
     model = ImageClassificationTransformer(pretrained_model_name_or_path="nateraw/vit-base-beans", num_labels=dm.num_classes)
diff --git a/examples/image_classification.py b/examples/image_classification.py
index 9015b837..f4428572 100644
--- a/examples/image_classification.py
+++ b/examples/image_classification.py
@@ -2,14 +2,15 @@
 from transformers import AutoFeatureExtractor
 
 from lightning_transformers.task.vision.image_classification import (
-    ImageClassificationDataConfig,
     ImageClassificationDataModule,
     ImageClassificationTransformer,
 )
 
 feature_extractor = AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path="nateraw/vit-base-beans")
 dm = ImageClassificationDataModule(
-    cfg=ImageClassificationDataConfig(batch_size=8, dataset_name="beans", num_workers=8),
+    batch_size=8,
+    dataset_name="beans",
+    num_workers=8,
     feature_extractor=feature_extractor,
 )
 model = ImageClassificationTransformer(
diff --git a/examples/language_modeling.py b/examples/language_modeling.py
index 21b71a4f..960f15ec 100644
--- a/examples/language_modeling.py
+++ b/examples/language_modeling.py
@@ -1,17 +1,15 @@
 import pytorch_lightning as pl
 from transformers import AutoTokenizer
 
-from lightning_transformers.task.nlp.language_modeling import (
-    LanguageModelingDataConfig,
-    LanguageModelingDataModule,
-    LanguageModelingTransformer,
-)
+from lightning_transformers.task.nlp.language_modeling import LanguageModelingDataModule, LanguageModelingTransformer
 
 if __name__ == "__main__":
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="gpt2")
     model = LanguageModelingTransformer(pretrained_model_name_or_path="gpt2")
     dm = LanguageModelingDataModule(
-        cfg=LanguageModelingDataConfig(batch_size=1, dataset_name="wikitext", dataset_config_name="wikitext-2-raw-v1"),
+        batch_size=1,
+        dataset_name="wikitext",
+        dataset_config_name="wikitext-2-raw-v1",
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
diff --git a/examples/masked_language_modeling.py b/examples/masked_language_modeling.py
index 1b57ab5b..95476091 100644
--- a/examples/masked_language_modeling.py
+++ b/examples/masked_language_modeling.py
@@ -2,7 +2,6 @@
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.masked_language_modeling import (
-    MaskedLanguageModelingDataConfig,
     MaskedLanguageModelingDataModule,
     MaskedLanguageModelingTransformer,
 )
@@ -11,12 +10,10 @@
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="bert-base-uncased")
     model = MaskedLanguageModelingTransformer(pretrained_model_name_or_path="bert-base-uncased")
     dm = MaskedLanguageModelingDataModule(
-        cfg=MaskedLanguageModelingDataConfig(
-            batch_size=1,
-            dataset_name="wikitext",
-            dataset_config_name="wikitext-2-raw-v1",
-            max_length=512,
-        ),
+        batch_size=1,
+        dataset_name="wikitext",
+        dataset_config_name="wikitext-2-raw-v1",
+        max_length=512,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
diff --git a/examples/multiple_choice.py b/examples/multiple_choice.py
index 271915f9..85380c3e 100644
--- a/examples/multiple_choice.py
+++ b/examples/multiple_choice.py
@@ -1,22 +1,16 @@
 import pytorch_lightning as pl
 from transformers import AutoTokenizer
 
-from lightning_transformers.task.nlp.multiple_choice import (
-    MultipleChoiceDataConfig,
-    MultipleChoiceTransformer,
-    SwagMultipleChoiceDataModule,
-)
+from lightning_transformers.task.nlp.multiple_choice import MultipleChoiceTransformer, SwagMultipleChoiceDataModule
 
 if __name__ == "__main__":
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="bert-base-uncased")
     model = MultipleChoiceTransformer(pretrained_model_name_or_path="bert-base-uncased")
     dm = SwagMultipleChoiceDataModule(
-        cfg=MultipleChoiceDataConfig(
-            batch_size=1,
-            dataset_name="swag",
-            dataset_config_name="regular",
-            padding=False,
-        ),
+        batch_size=1,
+        dataset_name="swag",
+        dataset_config_name="regular",
+        padding=False,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
diff --git a/examples/question_answering_squad.py b/examples/question_answering_squad.py
index 06a5a40f..d0367547 100644
--- a/examples/question_answering_squad.py
+++ b/examples/question_answering_squad.py
@@ -1,27 +1,20 @@
 import pytorch_lightning as pl
 from transformers import AutoTokenizer
 
-from lightning_transformers.task.nlp.question_answering import (
-    QuestionAnsweringDataConfig,
-    QuestionAnsweringTransformer,
-    SquadDataModule,
-)
+from lightning_transformers.task.nlp.question_answering import QuestionAnsweringTransformer, SquadDataModule
 
 if __name__ == "__main__":
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="bert-base-uncased")
     model = QuestionAnsweringTransformer(pretrained_model_name_or_path="bert-base-uncased")
     dm = SquadDataModule(
-        cfg=QuestionAnsweringDataConfig(
-            batch_size=1,
-            dataset_name="squad",
-            dataset_config_name="plain_text",
-            max_length=384,
-            version_2_with_negative=False,
-            null_score_diff_threshold=0.0,
-            doc_stride=128,
-            n_best_size=20,
-            max_answer_length=30,
-        ),
+        batch_size=1,
+        dataset_config_name="plain_text",
+        max_length=384,
+        version_2_with_negative=False,
+        null_score_diff_threshold=0.0,
+        doc_stride=128,
+        n_best_size=20,
+        max_answer_length=30,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
diff --git a/examples/summarization.py b/examples/summarization.py
index b0ac3321..d6ad66da 100644
--- a/examples/summarization.py
+++ b/examples/summarization.py
@@ -1,31 +1,21 @@
 import pytorch_lightning as pl
 from transformers import AutoTokenizer
 
-from lightning_transformers.task.nlp.summarization import (
-    SummarizationConfig,
-    SummarizationDataConfig,
-    SummarizationTransformer,
-    XsumSummarizationDataModule,
-)
+from lightning_transformers.task.nlp.summarization import SummarizationTransformer, XsumSummarizationDataModule
 
 if __name__ == "__main__":
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="t5-base")
     model = SummarizationTransformer(
         pretrained_model_name_or_path="t5-base",
-        cfg=SummarizationConfig(
-            use_stemmer=True,
-            val_target_max_length=142,
-            num_beams=None,
-            compute_generate_metrics=True,
-        ),
+        use_stemmer=True,
+        val_target_max_length=142,
+        num_beams=None,
+        compute_generate_metrics=True,
     )
     dm = XsumSummarizationDataModule(
-        cfg=SummarizationDataConfig(
-            batch_size=1,
-            dataset_name="xsum",
-            max_source_length=128,
-            max_target_length=128,
-        ),
+        batch_size=1,
+        max_source_length=128,
+        max_target_length=128,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(accelerator="auto", devices=1, max_epochs=1)
diff --git a/examples/text_classification.py b/examples/text_classification.py
index 05716d4c..0a401ac8 100644
--- a/examples/text_classification.py
+++ b/examples/text_classification.py
@@ -2,7 +2,6 @@
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.text_classification import (
-    TextClassificationDataConfig,
     TextClassificationDataModule,
     TextClassificationTransformer,
 )
@@ -10,12 +9,10 @@
 if __name__ == "__main__":
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="bert-base-uncased")
     dm = TextClassificationDataModule(
-        cfg=TextClassificationDataConfig(
-            batch_size=1,
-            dataset_name="glue",
-            dataset_config_name="sst2",
-            max_length=512,
-        ),
+        batch_size=1,
+        dataset_name="glue",
+        dataset_config_name="sst2",
+        max_length=512,
         tokenizer=tokenizer,
     )
     model = TextClassificationTransformer(pretrained_model_name_or_path="bert-base-uncased", num_labels=dm.num_classes)
diff --git a/examples/token_classification.py b/examples/token_classification.py
index 69a62dbb..5846f760 100644
--- a/examples/token_classification.py
+++ b/examples/token_classification.py
@@ -2,7 +2,6 @@
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.token_classification import (
-    TokenClassificationDataConfig,
     TokenClassificationDataModule,
     TokenClassificationTransformer,
 )
@@ -10,14 +9,12 @@
 if __name__ == "__main__":
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="bert-base-uncased")
     dm = TokenClassificationDataModule(
-        cfg=TokenClassificationDataConfig(
-            batch_size=1,
-            task_name="ner",
-            dataset_name="conll2003",
-            preprocessing_num_workers=1,
-            label_all_tokens=False,
-            revision="master",
-        ),
+        batch_size=1,
+        task_name="ner",
+        dataset_name="conll2003",
+        preprocessing_num_workers=1,
+        label_all_tokens=False,
+        revision="master",
         tokenizer=tokenizer,
     )
     model = TokenClassificationTransformer(pretrained_model_name_or_path="bert-base-uncased", labels=dm.labels)
diff --git a/examples/translation_wmt.py b/examples/translation_wmt.py
index 6c4d91d3..62688852 100644
--- a/examples/translation_wmt.py
+++ b/examples/translation_wmt.py
@@ -1,35 +1,25 @@
 import pytorch_lightning as pl
 from transformers import AutoTokenizer
 
-from lightning_transformers.task.nlp.translation import (
-    TranslationConfig,
-    TranslationDataConfig,
-    TranslationTransformer,
-    WMT16TranslationDataModule,
-)
+from lightning_transformers.task.nlp.translation import TranslationTransformer, WMT16TranslationDataModule
 
 if __name__ == "__main__":
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="t5-base")
     model = TranslationTransformer(
         pretrained_model_name_or_path="t5-base",
-        cfg=TranslationConfig(
-            n_gram=4,
-            smooth=False,
-            val_target_max_length=142,
-            num_beams=None,
-            compute_generate_metrics=True,
-        ),
+        n_gram=4,
+        smooth=False,
+        val_target_max_length=142,
+        num_beams=None,
+        compute_generate_metrics=True,
     )
     dm = WMT16TranslationDataModule(
-        cfg=TranslationDataConfig(
-            dataset_name="wmt16",
-            # WMT translation datasets: ['cs-en', 'de-en', 'fi-en', 'ro-en', 'ru-en', 'tr-en']
-            dataset_config_name="ro-en",
-            source_language="en",
-            target_language="ro",
-            max_source_length=128,
-            max_target_length=128,
-        ),
+        # WMT translation datasets: ['cs-en', 'de-en', 'fi-en', 'ro-en', 'ru-en', 'tr-en']
+        dataset_config_name="ro-en",
+        source_language="en",
+        target_language="ro",
+        max_source_length=128,
+        max_target_length=128,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
diff --git a/lightning_transformers/core/config.py b/lightning_transformers/core/config.py
deleted file mode 100644
index 5a2739db..00000000
--- a/lightning_transformers/core/config.py
+++ /dev/null
@@ -1,47 +0,0 @@
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Optional, Union
-
-from datasets import Version
-
-
-@dataclass
-class TransformerDataConfig:
-    batch_size: int = 32
-    num_workers: int = 0
-    dataset_name: Optional[str] = None
-    dataset_config_name: Optional[str] = None
-    revision: Optional[Union[str, Version]] = None
-    train_val_split: Optional[int] = None
-    train_file: Optional[str] = None
-    test_file: Optional[str] = None
-    predict_file: Optional[str] = None
-    validation_file: Optional[str] = None
-    padding: Union[str, bool] = "max_length"
-    truncation: str = "only_first"
-    max_length: int = 128
-    preprocessing_num_workers: int = 1
-    load_from_cache_file: bool = True
-    cache_dir: Optional[Union[Path, str]] = None
-    limit_train_samples: Optional[int] = None
-    limit_val_samples: Optional[int] = None
-    limit_test_samples: Optional[int] = None
-    train_subset_name: Optional[str] = None
-    validation_subset_name: Optional[str] = None
-    test_subset_name: Optional[str] = None
-    predict_subset_name: Optional[str] = None
-    streaming: bool = False
diff --git a/lightning_transformers/core/data.py b/lightning_transformers/core/data.py
index 0766538f..c38398af 100644
--- a/lightning_transformers/core/data.py
+++ b/lightning_transformers/core/data.py
@@ -1,13 +1,27 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import os
+from pathlib import Path
 from typing import Any, Callable, Dict, Optional, Union
 
 import pytorch_lightning as pl
-from datasets import Dataset, DatasetDict, load_dataset
+from datasets import Dataset, DatasetDict, Version, load_dataset
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from torch.utils.data import DataLoader
 from transformers import PreTrainedTokenizerBase
 
-from lightning_transformers.core.config import TransformerDataConfig
 from lightning_transformers.core.iterable import IterableDataLoader
 
 
@@ -17,18 +31,62 @@ class TransformerDataModule(pl.LightningDataModule):
 
     Args:
         tokenizer: ``PreTrainedTokenizerBase`` for tokenizing data.
-        cfg: Contains data specific parameters when processing/loading the dataset (Default ``HFTransformerDataConfig``)
     """
 
-    cfg: TransformerDataConfig
-    tokenizer: PreTrainedTokenizerBase
-
     def __init__(
-        self, tokenizer: PreTrainedTokenizerBase, cfg: TransformerDataConfig = TransformerDataConfig()
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        batch_size: int = 32,
+        num_workers: int = 0,
+        dataset_name: Optional[str] = None,
+        dataset_config_name: Optional[str] = None,
+        revision: Optional[Union[str, Version]] = None,
+        train_val_split: Optional[int] = None,
+        train_file: Optional[str] = None,
+        test_file: Optional[str] = None,
+        predict_file: Optional[str] = None,
+        validation_file: Optional[str] = None,
+        padding: Union[str, bool] = "max_length",
+        truncation: str = "only_first",
+        max_length: int = 128,
+        preprocessing_num_workers: int = 1,
+        load_from_cache_file: bool = True,
+        cache_dir: Optional[Union[Path, str]] = None,
+        limit_train_samples: Optional[int] = None,
+        limit_val_samples: Optional[int] = None,
+        limit_test_samples: Optional[int] = None,
+        train_subset_name: Optional[str] = None,
+        validation_subset_name: Optional[str] = None,
+        test_subset_name: Optional[str] = None,
+        predict_subset_name: Optional[str] = None,
+        streaming: bool = False,
     ) -> None:
         super().__init__()
         self.tokenizer = tokenizer
-        self.cfg = cfg
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.dataset_name = dataset_name
+        self.dataset_config_name = dataset_config_name
+        self.revision = revision
+        self.train_val_split = train_val_split
+        self.train_file = train_file
+        self.test_file = test_file
+        self.predict_file = predict_file
+        self.validation_file = validation_file
+        self.padding = padding
+        self.truncation = truncation
+        self.max_length = max_length
+        self.preprocessing_num_workers = preprocessing_num_workers
+        self.load_from_cache_file = load_from_cache_file
+        self.cache_dir = cache_dir
+        self.limit_train_samples = limit_train_samples
+        self.limit_val_samples = limit_val_samples
+        self.limit_test_samples = limit_test_samples
+        self.train_subset_name = train_subset_name
+        self.validation_subset_name = validation_subset_name
+        self.test_subset_name = test_subset_name
+        self.predict_subset_name = predict_subset_name
+        self.streaming = streaming
         os.environ["TOKENIZERS_PARALLELISM"] = "TRUE"  # todo: smarter handling of this env variable
 
     def setup(self, stage: Optional[str] = None):
@@ -45,28 +103,28 @@ def process_data(
     def load_dataset(self) -> Dataset:
         # Allow custom data files when loading the dataset
         data_files = {}
-        if self.cfg.train_file is not None:
-            data_files["train"] = self.cfg.train_file
-        if self.cfg.validation_file is not None:
-            data_files["validation"] = self.cfg.validation_file
-        if self.cfg.test_file is not None:
-            data_files["test"] = self.cfg.test_file
+        if self.train_file is not None:
+            data_files["train"] = self.train_file
+        if self.validation_file is not None:
+            data_files["validation"] = self.validation_file
+        if self.test_file is not None:
+            data_files["test"] = self.test_file
 
         data_files = data_files if data_files else None
-        if self.cfg.dataset_name is not None:
+        if self.dataset_name is not None:
             # Download and load the Huggingface dataset.
             dataset = load_dataset(
-                path=self.cfg.dataset_name,
-                name=self.cfg.dataset_config_name,
-                cache_dir=self.cfg.cache_dir,
+                path=self.dataset_name,
+                name=self.dataset_config_name,
+                cache_dir=self.cache_dir,
                 data_files=data_files,
-                revision=self.cfg.revision,
-                streaming=self.cfg.streaming,
+                revision=self.revision,
+                streaming=self.streaming,
             )
 
         # Load straight from data files
-        elif self.cfg.datafiles:
-            extension = self.cfg.train_file.split(".")[-1]
+        elif data_files:
+            extension = self.train_file.split(".")[-1]
             dataset = load_dataset(extension, data_files=data_files)
 
         else:
@@ -77,8 +135,8 @@ def load_dataset(self) -> Dataset:
         # Use special subset names if provided, and rename them back to standard ones
         for subset in ("train", "validation", "test", "predict"):
             config_attr = f"{subset}_subset_name"
-            if getattr(self.cfg, config_attr) is not None:
-                special_subset_name = getattr(self.cfg, config_attr)
+            if getattr(self, config_attr) is not None:
+                special_subset_name = getattr(self, config_attr)
                 if special_subset_name not in dataset:
                     raise KeyError(
                         f"Special {subset} subset name {special_subset_name} provided but not found in the dataset"
@@ -88,8 +146,8 @@ def load_dataset(self) -> Dataset:
         return dataset
 
     def split_dataset(self, dataset: Union[Dataset, DatasetDict]) -> Union[Dataset, DatasetDict]:
-        if self.cfg.train_val_split is not None:
-            split = dataset["train"].train_test_split(self.cfg.train_val_split)
+        if self.train_val_split is not None:
+            split = dataset["train"].train_test_split(self.train_val_split)
             dataset["train"] = split["train"]
             dataset["validation"] = split["test"]
         dataset = self._select_samples(dataset)
@@ -97,9 +155,9 @@ def split_dataset(self, dataset: Union[Dataset, DatasetDict]) -> Union[Dataset,
 
     def _select_samples(self, dataset: Union[Dataset, DatasetDict]) -> Union[Dataset, DatasetDict]:
         samples = (
-            ("train", self.cfg.limit_train_samples),
-            ("validation", self.cfg.limit_val_samples),
-            ("test", self.cfg.limit_test_samples),
+            ("train", self.limit_train_samples),
+            ("validation", self.limit_val_samples),
+            ("test", self.limit_test_samples),
         )
         for column_name, n_samples in samples:
             if n_samples is not None and column_name in dataset:
@@ -114,47 +172,43 @@ def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
         self.tokenizer = checkpoint["tokenizer"]
 
     def train_dataloader(self) -> DataLoader:
-        cls = DataLoader if not self.cfg.streaming else IterableDataLoader
+        cls = DataLoader if not self.streaming else IterableDataLoader
         return cls(
             self.ds["train"],
             batch_size=self.batch_size,
-            num_workers=self.cfg.num_workers,
+            num_workers=self.num_workers,
             collate_fn=self.collate_fn,
         )
 
     def val_dataloader(self) -> DataLoader:
-        cls = DataLoader if not self.cfg.streaming else IterableDataLoader
+        cls = DataLoader if not self.streaming else IterableDataLoader
         return cls(
             self.ds["validation"],
             batch_size=self.batch_size,
-            num_workers=self.cfg.num_workers,
+            num_workers=self.num_workers,
             collate_fn=self.collate_fn,
         )
 
     def test_dataloader(self) -> Optional[DataLoader]:
         if "test" in self.ds:
-            cls = DataLoader if not self.cfg.streaming else IterableDataLoader
+            cls = DataLoader if not self.streaming else IterableDataLoader
             return cls(
                 self.ds["test"],
                 batch_size=self.batch_size,
-                num_workers=self.cfg.num_workers,
+                num_workers=self.num_workers,
                 collate_fn=self.collate_fn,
             )
 
     def predict_dataloader(self) -> Optional[DataLoader]:
         if "predict" in self.ds:
-            cls = DataLoader if not self.cfg.streaming else IterableDataLoader
+            cls = DataLoader if not self.streaming else IterableDataLoader
             return cls(
                 self.ds["predict"],
                 batch_size=self.batch_size,
-                num_workers=self.cfg.num_workers,
+                num_workers=self.num_workers,
                 collate_fn=self.collate_fn,
             )
 
-    @property
-    def batch_size(self) -> int:
-        return self.cfg.batch_size
-
     @property
     def collate_fn(self) -> Optional[Callable]:
         return None
diff --git a/lightning_transformers/core/seq2seq/__init__.py b/lightning_transformers/core/seq2seq/__init__.py
index e113d97a..96dc25b7 100644
--- a/lightning_transformers/core/seq2seq/__init__.py
+++ b/lightning_transformers/core/seq2seq/__init__.py
@@ -1,4 +1,3 @@
-from lightning_transformers.core.seq2seq.config import Seq2SeqConfig, Seq2SeqDataConfig  # noqa: F401
 from lightning_transformers.core.seq2seq.data import Seq2SeqDataModule  # noqa: F401
 from lightning_transformers.core.seq2seq.finetuning import FreezeEmbeddings  # noqa: F401
 from lightning_transformers.core.seq2seq.model import Seq2SeqTransformer  # noqa: F401
diff --git a/lightning_transformers/core/seq2seq/config.py b/lightning_transformers/core/seq2seq/config.py
deleted file mode 100644
index 7207ee85..00000000
--- a/lightning_transformers/core/seq2seq/config.py
+++ /dev/null
@@ -1,18 +0,0 @@
-from dataclasses import dataclass
-from typing import Optional
-
-from lightning_transformers.core.config import TransformerDataConfig
-
-
-@dataclass
-class Seq2SeqConfig:
-    val_target_max_length: Optional[int] = 128
-    num_beams: Optional[int] = 1
-    compute_generate_metrics: bool = True
-
-
-@dataclass
-class Seq2SeqDataConfig(TransformerDataConfig):
-    max_target_length: int = 128
-    max_source_length: int = 1024
-    padding: str = "longest"
diff --git a/lightning_transformers/core/seq2seq/data.py b/lightning_transformers/core/seq2seq/data.py
index 5d65646e..2c007ff6 100644
--- a/lightning_transformers/core/seq2seq/data.py
+++ b/lightning_transformers/core/seq2seq/data.py
@@ -5,7 +5,6 @@
 from transformers import PreTrainedTokenizerBase, default_data_collator
 
 from lightning_transformers.core import TransformerDataModule
-from lightning_transformers.core.seq2seq.config import Seq2SeqDataConfig
 
 
 class Seq2SeqDataModule(TransformerDataModule):
@@ -13,15 +12,16 @@ class Seq2SeqDataModule(TransformerDataModule):
 
     Args:
         *args: ``HFDataModule`` specific arguments.
-        cfg: Contains data specific parameters when processing/loading the dataset
-            (Default ``Seq2SeqDataConfig``)
         **kwargs: ``HFDataModule`` specific arguments.
     """
 
-    cfg: Seq2SeqDataConfig
-
-    def __init__(self, *args, cfg: Seq2SeqDataConfig = Seq2SeqDataConfig(), **kwargs) -> None:
-        super().__init__(*args, cfg=cfg, **kwargs)
+    def __init__(
+        self, *args, max_target_length: int = 128, max_source_length: int = 1024, padding: str = "longest", **kwargs
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.max_target_length = max_target_length
+        self.max_source_length = max_source_length
+        self.padding = padding
 
     def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset:
         src_text_column_name, tgt_text_column_name = self.source_target_column_names
@@ -29,17 +29,17 @@ def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset
         convert_to_features = partial(
             self.convert_to_features,
             tokenizer=self.tokenizer,
-            padding=self.cfg.padding,
-            max_source_length=self.cfg.max_source_length,
-            max_target_length=self.cfg.max_target_length,
+            padding=self.padding,
+            max_source_length=self.max_source_length,
+            max_target_length=self.max_target_length,
             src_text_column_name=src_text_column_name,
             tgt_text_column_name=tgt_text_column_name,
         )
         dataset = dataset.map(
             convert_to_features,
             batched=True,
-            num_proc=self.cfg.preprocessing_num_workers,
-            load_from_cache_file=self.cfg.load_from_cache_file,
+            num_proc=self.preprocessing_num_workers,
+            load_from_cache_file=self.load_from_cache_file,
         )
 
         cols_to_keep = [x for x in ["input_ids", "attention_mask", "labels"] if x in dataset["train"].features]
diff --git a/lightning_transformers/core/seq2seq/model.py b/lightning_transformers/core/seq2seq/model.py
index 02559e43..848134d7 100644
--- a/lightning_transformers/core/seq2seq/model.py
+++ b/lightning_transformers/core/seq2seq/model.py
@@ -1,16 +1,24 @@
-from typing import Any, List
+from typing import Any, List, Optional
 
 import torch
 
 from lightning_transformers.core import TaskTransformer
-from lightning_transformers.core.seq2seq.config import Seq2SeqConfig
 from lightning_transformers.core.seq2seq.utils import _pad_tensors_to_max_len
 
 
 class Seq2SeqTransformer(TaskTransformer):
-    def __init__(self, *args, cfg: Seq2SeqConfig = Seq2SeqConfig(), **kwargs) -> None:
+    def __init__(
+        self,
+        *args,
+        val_target_max_length: Optional[int] = 128,
+        num_beams: Optional[int] = 1,
+        compute_generate_metrics: bool = True,
+        **kwargs
+    ) -> None:
         super().__init__(*args, **kwargs)
-        self.cfg = cfg
+        self.val_target_max_length = val_target_max_length
+        self.num_beams = num_beams
+        self.should_compute_generate_metrics = compute_generate_metrics
 
     def training_step(self, batch: Any, batch_idx: int) -> torch.Tensor:
         outputs = self.model(**batch)
@@ -21,7 +29,7 @@ def training_step(self, batch: Any, batch_idx: int) -> torch.Tensor:
     def common_step(self, prefix: str, batch: Any) -> torch.Tensor:
         outputs = self.model(**batch)
         loss, logits = outputs[:2]
-        if self.cfg.compute_generate_metrics:
+        if self.should_compute_generate_metrics:
             self.compute_generate_metrics(batch, prefix)
         return loss
 
@@ -35,8 +43,8 @@ def compute_generate_metrics(self, batch, prefix):
         raise NotImplementedError
 
     def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> List[str]:
-        max_length = self.cfg.val_target_max_length if self.cfg.val_target_max_length else self.model.config.max_length
-        num_beams = self.cfg.num_beams if self.cfg.num_beams else self.model.config.num_beams
+        max_length = self.val_target_max_length if self.val_target_max_length else self.model.config.max_length
+        num_beams = self.num_beams if self.num_beams else self.model.config.num_beams
         generated_tokens = self.model.generate(
             input_ids=input_ids, attention_mask=attention_mask, max_length=max_length, num_beams=num_beams
         )
diff --git a/lightning_transformers/task/nlp/language_modeling/__init__.py b/lightning_transformers/task/nlp/language_modeling/__init__.py
index e6fb8350..da9ca8c5 100644
--- a/lightning_transformers/task/nlp/language_modeling/__init__.py
+++ b/lightning_transformers/task/nlp/language_modeling/__init__.py
@@ -1,3 +1,2 @@
-from lightning_transformers.task.nlp.language_modeling.config import LanguageModelingDataConfig  # noqa: F401
 from lightning_transformers.task.nlp.language_modeling.data import LanguageModelingDataModule  # noqa: F401
 from lightning_transformers.task.nlp.language_modeling.model import LanguageModelingTransformer  # noqa: F401
diff --git a/lightning_transformers/task/nlp/language_modeling/config.py b/lightning_transformers/task/nlp/language_modeling/config.py
deleted file mode 100644
index e3242bbd..00000000
--- a/lightning_transformers/task/nlp/language_modeling/config.py
+++ /dev/null
@@ -1,21 +0,0 @@
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from dataclasses import dataclass
-
-from lightning_transformers.core.config import TransformerDataConfig
-
-
-@dataclass
-class LanguageModelingDataConfig(TransformerDataConfig):
-    block_size: int = 128
diff --git a/lightning_transformers/task/nlp/language_modeling/data.py b/lightning_transformers/task/nlp/language_modeling/data.py
index 5ee36418..581bfda2 100644
--- a/lightning_transformers/task/nlp/language_modeling/data.py
+++ b/lightning_transformers/task/nlp/language_modeling/data.py
@@ -19,7 +19,6 @@
 from transformers import PreTrainedTokenizerBase, default_data_collator
 
 from lightning_transformers.core import TransformerDataModule
-from lightning_transformers.task.nlp.language_modeling.config import LanguageModelingDataConfig
 
 
 class LanguageModelingDataModule(TransformerDataModule):
@@ -27,15 +26,12 @@ class LanguageModelingDataModule(TransformerDataModule):
 
     Args:
         *args: ``HFDataModule`` specific arguments.
-        cfg: Contains data specific parameters when processing/loading the dataset
-            (Default ``LanguageModelingDataConfig``)
         **kwargs: ``HFDataModule`` specific arguments.
     """
 
-    cfg: LanguageModelingDataConfig
-
-    def __init__(self, *args, cfg: LanguageModelingDataConfig = LanguageModelingDataConfig(), **kwargs) -> None:
-        super().__init__(*args, cfg=cfg, **kwargs)
+    def __init__(self, *args, block_size: int = 128, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.block_size = block_size
 
     def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset:
         dataset_split = dataset["train" if stage == "fit" else "validation"]
@@ -58,36 +54,36 @@ def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset
         dataset = dataset.map(
             tokenize_function,
             batched=True,
-            num_proc=self.cfg.preprocessing_num_workers,
+            num_proc=self.preprocessing_num_workers,
             remove_columns=column_names,
-            load_from_cache_file=self.cfg.load_from_cache_file,
+            load_from_cache_file=self.load_from_cache_file,
         )
 
         return dataset.map(
             convert_to_features,
             batched=True,
-            num_proc=self.cfg.preprocessing_num_workers,
-            load_from_cache_file=self.cfg.load_from_cache_file,
+            num_proc=self.preprocessing_num_workers,
+            load_from_cache_file=self.load_from_cache_file,
         )
 
     @property
     def effective_block_size(self) -> int:
-        if self.cfg.block_size is None:
+        if self.block_size is None:
             block_size = self.tokenizer.model_max_length
             if block_size > 1024:
                 log.warn(
                     f"The tokenizer picked seems to have a very large `model_max_length` "
                     f"({self.tokenizer.model_max_length}). "
-                    "Picking 1024 instead. You can change that default value by passing dataset.cfg.block_size=x."
+                    "Picking 1024 instead. You can change that default value by passing block_size=x to the DataModule."
                 )
             block_size = 1024
         else:
-            if self.cfg.block_size > self.tokenizer.model_max_length:
+            if self.block_size > self.tokenizer.model_max_length:
                 log.warn(
-                    f"The block_size passed ({self.cfg.block_size}) is larger than the maximum length for the model"
+                    f"The block_size passed ({self.block_size}) is larger than the maximum length for the model"
                     f"({self.tokenizer.model_max_length}). Using block_size={self.tokenizer.model_max_length}."
                 )
-            block_size = min(self.cfg.block_size, self.tokenizer.model_max_length)
+            block_size = min(self.block_size, self.tokenizer.model_max_length)
         return block_size
 
     @staticmethod
diff --git a/lightning_transformers/task/nlp/language_modeling/model.py b/lightning_transformers/task/nlp/language_modeling/model.py
index 17798dae..9bc67484 100644
--- a/lightning_transformers/task/nlp/language_modeling/model.py
+++ b/lightning_transformers/task/nlp/language_modeling/model.py
@@ -27,9 +27,9 @@ class LanguageModelingTransformer(TaskTransformer):
     """Defines ``LightningModule`` for the Language Modeling Task.
 
     Args:
-        *args: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
+        *args: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
         downstream_model_type: Downstream HuggingFace AutoModel to load. (default ``transformers.AutoModelForCausalLM``)
-        **kwargs: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
+        **kwargs: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
     """
 
     def __init__(
diff --git a/lightning_transformers/task/nlp/masked_language_modeling/__init__.py b/lightning_transformers/task/nlp/masked_language_modeling/__init__.py
index cba63b0d..78536200 100644
--- a/lightning_transformers/task/nlp/masked_language_modeling/__init__.py
+++ b/lightning_transformers/task/nlp/masked_language_modeling/__init__.py
@@ -1,6 +1,3 @@
-from lightning_transformers.task.nlp.masked_language_modeling.config import (  # noqa: F401
-    MaskedLanguageModelingDataConfig,
-)
 from lightning_transformers.task.nlp.masked_language_modeling.data import MaskedLanguageModelingDataModule  # noqa: F401
 from lightning_transformers.task.nlp.masked_language_modeling.model import (  # noqa: F401
     MaskedLanguageModelingTransformer,
diff --git a/lightning_transformers/task/nlp/masked_language_modeling/config.py b/lightning_transformers/task/nlp/masked_language_modeling/config.py
deleted file mode 100644
index bc9f245d..00000000
--- a/lightning_transformers/task/nlp/masked_language_modeling/config.py
+++ /dev/null
@@ -1,32 +0,0 @@
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from dataclasses import dataclass, field
-
-from lightning_transformers.core.config import TransformerDataConfig
-
-
-@dataclass
-class MaskedLanguageModelingDataConfig(TransformerDataConfig):
-    mlm_probability: float = field(
-        default=0.15,
-        metadata={"help": "Ratio of tokens to mask for masked language modeling loss"},
-    )
-    line_by_line: bool = field(
-        default=False,
-        metadata={"help": "Whether distinct lines of text in the dataset are to be handled as distinct sequences."},
-    )
-    wwm: bool = field(
-        default=False,
-        metadata={"help": "Whether or not use whole word masking"},
-    )
diff --git a/lightning_transformers/task/nlp/masked_language_modeling/data.py b/lightning_transformers/task/nlp/masked_language_modeling/data.py
index 667e4f1b..8c3e66f2 100644
--- a/lightning_transformers/task/nlp/masked_language_modeling/data.py
+++ b/lightning_transformers/task/nlp/masked_language_modeling/data.py
@@ -18,7 +18,6 @@
 from transformers import DataCollatorForLanguageModeling, DataCollatorForWholeWordMask, PreTrainedTokenizerBase
 
 from lightning_transformers.core import TransformerDataModule
-from lightning_transformers.task.nlp.masked_language_modeling.config import MaskedLanguageModelingDataConfig
 
 
 class MaskedLanguageModelingDataModule(TransformerDataModule):
@@ -26,17 +25,19 @@ class MaskedLanguageModelingDataModule(TransformerDataModule):
 
     Args:
         *args: ``HFDataModule`` specific arguments.
-        cfg: Contains data specific parameters when processing/loading the dataset
-            (Default ``MaskedLanguageModelingDataConfig``)
+        mlm_probability: Ratio of tokens to mask for masked language modeling loss
+        line_by_line: Whether distinct lines of text in the dataset are to be handled as distinct sequences.
+        wwm: Whether or not use whole word masking.
         **kwargs: ``HFDataModule`` specific arguments.
     """
 
-    cfg: MaskedLanguageModelingDataConfig
-
     def __init__(
-        self, *args, cfg: MaskedLanguageModelingDataConfig = MaskedLanguageModelingDataConfig(), **kwargs
+        self, *args, mlm_probability: float = 0.15, line_by_line: bool = False, wwm: bool = False, **kwargs
     ) -> None:
-        super().__init__(*args, cfg=cfg, **kwargs)
+        super().__init__(*args, **kwargs)
+        self.mlm_probability = mlm_probability
+        self.line_by_line = line_by_line
+        self.wwm = wwm
 
     def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset:
         column_names = dataset["train" if stage == "fit" else "validation"].column_names
@@ -45,30 +46,30 @@ def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset
             self.tokenize_function,
             tokenizer=self.tokenizer,
             text_column_name=text_column_name,
-            line_by_line=self.cfg.line_by_line,
-            padding=self.cfg.padding,
-            max_length=self.cfg.max_length,
+            line_by_line=self.line_by_line,
+            padding=self.padding,
+            max_length=self.max_length,
         )
 
         dataset = dataset.map(
             tokenize_function,
             batched=True,
-            num_proc=self.cfg.preprocessing_num_workers,
+            num_proc=self.preprocessing_num_workers,
             remove_columns=column_names,
-            load_from_cache_file=self.cfg.load_from_cache_file,
+            load_from_cache_file=self.load_from_cache_file,
         )
 
-        if not self.cfg.line_by_line:
+        if not self.line_by_line:
             convert_to_features = partial(
                 self.convert_to_features,
-                max_seq_length=self.cfg.max_length,
+                max_seq_length=self.max_length,
             )
 
             dataset = dataset.map(
                 convert_to_features,
                 batched=True,
-                num_proc=self.cfg.preprocessing_num_workers,
-                load_from_cache_file=self.cfg.load_from_cache_file,
+                num_proc=self.preprocessing_num_workers,
+                load_from_cache_file=self.load_from_cache_file,
             )
 
         return dataset
@@ -118,7 +119,7 @@ def convert_to_features(examples, max_seq_length: int, **kwargs):
 
     @property
     def collate_fn(self) -> Callable:
-        if self.cfg.wwm:
-            return DataCollatorForWholeWordMask(self.tokenizer, mlm_probability=self.cfg.mlm_probability)
+        if self.wwm:
+            return DataCollatorForWholeWordMask(self.tokenizer, mlm_probability=self.mlm_probability)
         else:
-            return DataCollatorForLanguageModeling(self.tokenizer, mlm_probability=self.cfg.mlm_probability)
+            return DataCollatorForLanguageModeling(self.tokenizer, mlm_probability=self.mlm_probability)
diff --git a/lightning_transformers/task/nlp/masked_language_modeling/model.py b/lightning_transformers/task/nlp/masked_language_modeling/model.py
index 06f0080e..3e71c241 100644
--- a/lightning_transformers/task/nlp/masked_language_modeling/model.py
+++ b/lightning_transformers/task/nlp/masked_language_modeling/model.py
@@ -25,9 +25,9 @@ class MaskedLanguageModelingTransformer(TaskTransformer):
     """Defines ``LightningModule`` for the Masked Language Modeling Task.
 
     Args:
-        *args: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
+        *args: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
         downstream_model_type: Downstream HuggingFace AutoModel to load. (default ``transformers.AutoModelForMaskedLM``)
-        **kwargs: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
+        **kwargs: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
     """
 
     def __init__(
diff --git a/lightning_transformers/task/nlp/multiple_choice/__init__.py b/lightning_transformers/task/nlp/multiple_choice/__init__.py
index 556d1d24..29c48baa 100644
--- a/lightning_transformers/task/nlp/multiple_choice/__init__.py
+++ b/lightning_transformers/task/nlp/multiple_choice/__init__.py
@@ -1,4 +1,3 @@
-from lightning_transformers.task.nlp.multiple_choice.config import MultipleChoiceDataConfig  # noqa: F401
 from lightning_transformers.task.nlp.multiple_choice.data import MultipleChoiceDataModule  # noqa: F401
 from lightning_transformers.task.nlp.multiple_choice.datasets import (  # noqa: F401
     RaceMultipleChoiceDataModule,
diff --git a/lightning_transformers/task/nlp/multiple_choice/config.py b/lightning_transformers/task/nlp/multiple_choice/config.py
deleted file mode 100644
index 97c177db..00000000
--- a/lightning_transformers/task/nlp/multiple_choice/config.py
+++ /dev/null
@@ -1,21 +0,0 @@
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from dataclasses import dataclass
-
-from lightning_transformers.core.config import TransformerDataConfig
-
-
-@dataclass
-class MultipleChoiceDataConfig(TransformerDataConfig):
-    ...
diff --git a/lightning_transformers/task/nlp/multiple_choice/data.py b/lightning_transformers/task/nlp/multiple_choice/data.py
index d102d970..a218d4f2 100644
--- a/lightning_transformers/task/nlp/multiple_choice/data.py
+++ b/lightning_transformers/task/nlp/multiple_choice/data.py
@@ -16,21 +16,15 @@
 from transformers import default_data_collator
 
 from lightning_transformers.core import TransformerDataModule
-from lightning_transformers.task.nlp.multiple_choice.config import MultipleChoiceDataConfig
 from lightning_transformers.task.nlp.multiple_choice.utils import DataCollatorForMultipleChoice
 
 
 class MultipleChoiceDataModule(TransformerDataModule):
     """Defines the ``LightningDataModule`` for Multiple Choice Datasets."""
 
-    cfg: MultipleChoiceDataConfig
-
-    def __init__(self, *args, cfg: MultipleChoiceDataConfig = MultipleChoiceDataConfig(), **kwargs) -> None:
-        super().__init__(*args, cfg=cfg, **kwargs)
-
     @property
     def pad_to_max_length(self):
-        return self.cfg.padding == "max_length"
+        return self.padding == "max_length"
 
     @property
     def collate_fn(self) -> callable:
diff --git a/lightning_transformers/task/nlp/multiple_choice/datasets/race/data.py b/lightning_transformers/task/nlp/multiple_choice/datasets/race/data.py
index 551ce187..c9c02d0b 100644
--- a/lightning_transformers/task/nlp/multiple_choice/datasets/race/data.py
+++ b/lightning_transformers/task/nlp/multiple_choice/datasets/race/data.py
@@ -21,6 +21,9 @@
 
 
 class RaceMultipleChoiceDataModule(MultipleChoiceDataModule):
+    def __init__(self, *args, dataset_name: str = "race", dataset_config_name="all", **kwargs):
+        super().__init__(*args, dataset_name=dataset_name, dataset_config_name=dataset_config_name, **kwargs)
+
     @property
     def choices(self) -> list:
         return ["A", "B", "C", "D"]
@@ -34,15 +37,15 @@ def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset
             question_header_name=self.question_header_name,
             answer_column_name=self.answer_column_name,
             options_column_name=self.options_column_name,
-            max_length=self.cfg.max_length,
-            padding=self.cfg.padding,
+            max_length=self.max_length,
+            padding=self.padding,
         )
 
         dataset = dataset.map(
             convert_to_features,
             batched=True,
-            num_proc=self.cfg.preprocessing_num_workers,
-            load_from_cache_file=self.cfg.load_from_cache_file,
+            num_proc=self.preprocessing_num_workers,
+            load_from_cache_file=self.load_from_cache_file,
         )
 
         cols_to_keep = [
diff --git a/lightning_transformers/task/nlp/multiple_choice/datasets/swag/data.py b/lightning_transformers/task/nlp/multiple_choice/datasets/swag/data.py
index 1d6aef76..dda160c3 100644
--- a/lightning_transformers/task/nlp/multiple_choice/datasets/swag/data.py
+++ b/lightning_transformers/task/nlp/multiple_choice/datasets/swag/data.py
@@ -24,6 +24,9 @@
 class SwagMultipleChoiceDataModule(MultipleChoiceDataModule):
     num_choices: int = 4  # there are four different endings to select in the SWAG dataset
 
+    def __init__(self, *args, dataset_name: str = "swag", dataset_config_name="regular", **kwargs):
+        super().__init__(*args, dataset_name=dataset_name, dataset_config_name=dataset_config_name, **kwargs)
+
     @property
     def num_classes(self) -> int:
         return len(self.ending_column_names)
@@ -36,15 +39,15 @@ def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset
             context_name=self.context_name,
             question_header_name=self.question_header_name,
             ending_names=self.ending_column_names,
-            max_length=self.cfg.max_length,
-            padding=self.cfg.padding,
+            max_length=self.max_length,
+            padding=self.padding,
         )
 
         dataset = dataset.map(
             convert_to_features,
             batched=True,
-            num_proc=self.cfg.preprocessing_num_workers,
-            load_from_cache_file=self.cfg.load_from_cache_file,
+            num_proc=self.preprocessing_num_workers,
+            load_from_cache_file=self.load_from_cache_file,
         )
 
         cols_to_keep = [
diff --git a/lightning_transformers/task/nlp/multiple_choice/model.py b/lightning_transformers/task/nlp/multiple_choice/model.py
index 175698af..6503654c 100644
--- a/lightning_transformers/task/nlp/multiple_choice/model.py
+++ b/lightning_transformers/task/nlp/multiple_choice/model.py
@@ -27,10 +27,10 @@ class MultipleChoiceTransformer(TaskTransformer):
     """Defines ``LightningModule`` for the Multiple Choice Task.
 
     Args:
-        *args: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
+        *args: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
         downstream_model_type: Downstream HuggingFace AutoModel to load.
             (default ``transformers.AutoModelForMultipleChoice``)
-        **kwargs: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
+        **kwargs: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
     """
 
     def __init__(
diff --git a/lightning_transformers/task/nlp/question_answering/__init__.py b/lightning_transformers/task/nlp/question_answering/__init__.py
index 9094c9d6..72a9231b 100644
--- a/lightning_transformers/task/nlp/question_answering/__init__.py
+++ b/lightning_transformers/task/nlp/question_answering/__init__.py
@@ -1,4 +1,3 @@
-from lightning_transformers.task.nlp.question_answering.config import QuestionAnsweringDataConfig  # noqa: F401
 from lightning_transformers.task.nlp.question_answering.data import QuestionAnsweringDataModule  # noqa: F401
 from lightning_transformers.task.nlp.question_answering.datasets import SquadDataModule  # noqa: F401
 from lightning_transformers.task.nlp.question_answering.model import QuestionAnsweringTransformer  # noqa: F401
diff --git a/lightning_transformers/task/nlp/question_answering/config.py b/lightning_transformers/task/nlp/question_answering/config.py
deleted file mode 100644
index b2b1fb2a..00000000
--- a/lightning_transformers/task/nlp/question_answering/config.py
+++ /dev/null
@@ -1,78 +0,0 @@
-# Copyright 2020 The HuggingFace Team All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Adapted from:
-# https://github.com/huggingface/transformers/blob/master/examples/question-answering/run_qa.py
-# yapf: off
-
-from dataclasses import dataclass, field
-from typing import Optional
-
-from lightning_transformers.core.config import TransformerDataConfig
-
-
-@dataclass
-class QuestionAnsweringDataConfig(TransformerDataConfig):
-    """Arguments pertaining to what data we are going to input our model for training and eval."""
-
-    max_length: int = field(
-        default=384,
-        metadata={
-            "help": "The maximum total input sequence length after tokenization. Sequences longer "
-            "than this will be truncated, sequences shorter will be padded."
-        },
-    )
-    version_2_with_negative: bool = field(
-        default=False, metadata={"help": "If true, some of the examples do not have an answer."}
-    )
-    null_score_diff_threshold: float = field(
-        default=0.0,
-        metadata={
-            "help": "The threshold used to select the null answer: if the best answer has a score that is less than "
-            "the score of the null answer minus this threshold, the null answer is selected for this example. "
-            "Only useful when `version_2_with_negative=True`."
-        },
-    )
-    doc_stride: int = field(
-        default=128,
-        metadata={"help": "When splitting up a long document into chunks, how much stride to take between chunks."},
-    )
-    n_best_size: int = field(
-        default=20,
-        metadata={"help": "The total number of n-best predictions to generate when looking for an answer."},
-    )
-    max_answer_length: int = field(
-        default=30,
-        metadata={
-            "help": "The maximum length of an answer that can be generated. This is needed because the start "
-            "and end predictions are not conditioned on one another."
-        },
-    )
-    output_dir: Optional[str] = field(
-        default=None,
-        metadata={
-            "help": "If provided, the dictionaries of predictions, "
-            "n_best predictions (with their scores and logits) and, "
-            "if `version_2_with_negative=True`, the dictionary of the scores differences "
-            "between best and null answers, are saved in `output_dir`."
-        },
-    )
-
-    def __post_init__(self):
-        if self.train_file is not None:
-            extension = self.train_file.split(".")[-1]
-            assert extension in ["csv", "json"], "`train_file` should be a csv or a json file."
-        if self.validation_file is not None:
-            extension = self.validation_file.split(".")[-1]
-            assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file."
diff --git a/lightning_transformers/task/nlp/question_answering/data.py b/lightning_transformers/task/nlp/question_answering/data.py
index 31532fb1..15677b91 100644
--- a/lightning_transformers/task/nlp/question_answering/data.py
+++ b/lightning_transformers/task/nlp/question_answering/data.py
@@ -20,7 +20,6 @@
 from transformers import DataCollatorWithPadding, PreTrainedTokenizerBase, default_data_collator
 
 from lightning_transformers.core import TransformerDataModule
-from lightning_transformers.task.nlp.question_answering.config import QuestionAnsweringDataConfig
 
 
 class QuestionAnsweringDataModule(TransformerDataModule):
@@ -28,16 +27,44 @@ class QuestionAnsweringDataModule(TransformerDataModule):
 
     Args:
         *args: ``HFDataModule`` specific arguments.
-        cfg: Contains data specific parameters when processing/loading the dataset
-            (Default ``QuestionAnsweringDataConfig``)
+        max_length: The maximum total input sequence length after tokenization. Sequences longer
+            than this will be truncated, sequences shorter will be padded.
+        version_2_with_negative: If true, some of the examples do not have an answer.
+        null_score_diff_threshold: The threshold used to select the null answer:
+            if the best answer has a score that is less than
+            the score of the null answer minus this threshold, the null answer is selected for this example.
+            Only useful when `version_2_with_negative=True`.
+        doc_stride: When splitting up a long document into chunks, how much stride to take between chunks.
+        n_best_size: The total number of n-best predictions to generate when looking for an answer.
+        max_answer_length: The maximum length of an answer that can be generated. This is needed because the start
+            and end predictions are not conditioned on one another.
+        output_dir: If provided, the dictionaries of predictions,
+            n_best predictions (with their scores and logits) and,
+            if `version_2_with_negative=True`, the dictionary of the scores differences
+            between best and null answers, are saved in `output_dir`.
         **kwargs: ``HFDataModule`` specific arguments.
     """
 
-    cfg: QuestionAnsweringDataConfig
-
-    def __init__(self, *args, cfg: QuestionAnsweringDataConfig = QuestionAnsweringDataConfig(), **kwargs) -> None:
-        super().__init__(*args, cfg=cfg, **kwargs)
+    def __init__(
+        self,
+        *args,
+        max_length: int = 384,
+        version_2_with_negative: bool = False,
+        null_score_diff_threshold: float = 0.0,
+        doc_stride: int = 128,
+        n_best_size: int = 20,
+        max_answer_length: int = 30,
+        output_dir: Optional[str] = None,
+        **kwargs
+    ) -> None:
+        super().__init__(*args, max_length=max_length, **kwargs)
         self.example_id_strings = OrderedDict()
+        self.version_2_with_negative = version_2_with_negative
+        self.null_score_diff_threshold = null_score_diff_threshold
+        self.doc_stride = doc_stride
+        self.n_best_size = n_best_size
+        self.max_answer_length = max_answer_length
+        self.output_dir = output_dir
 
     def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset:
         train = stage == "fit"
@@ -54,9 +81,9 @@ def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset
             "question_column_name": question_column_name,
             "context_column_name": context_column_name,
             "answer_column_name": answer_column_name,
-            "max_length": self.cfg.max_length,
-            "doc_stride": self.cfg.doc_stride,
-            "padding": self.cfg.padding,
+            "max_length": self.max_length,
+            "doc_stride": self.doc_stride,
+            "padding": self.padding,
         }
 
         prepare_train_features = partial(self.convert_to_train_features, **kwargs)
@@ -65,9 +92,9 @@ def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset
             dataset["train"] = dataset["train"].map(
                 prepare_train_features,
                 batched=True,
-                num_proc=self.cfg.preprocessing_num_workers,
+                num_proc=self.preprocessing_num_workers,
                 remove_columns=column_names,
-                load_from_cache_file=self.cfg.load_from_cache_file,
+                load_from_cache_file=self.load_from_cache_file,
             )
 
         if "test" not in dataset:
@@ -79,7 +106,7 @@ def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset
             dataset["validation"] = dataset["validation"].map(
                 prepare_validation_features,
                 batched=True,
-                num_proc=self.cfg.preprocessing_num_workers,
+                num_proc=self.preprocessing_num_workers,
                 remove_columns=column_names,
                 # Legacy code `prepare_validation_features` must run to populate `self.example_id_strings`
                 # Therefore we cannot load from cache here
@@ -89,7 +116,7 @@ def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset
 
     @property
     def pad_to_max_length(self) -> bool:
-        return self.cfg.padding == "max_length"
+        return self.padding == "max_length"
 
     @property
     def collate_fn(self) -> Callable:
diff --git a/lightning_transformers/task/nlp/question_answering/datasets/squad/data.py b/lightning_transformers/task/nlp/question_answering/datasets/squad/data.py
index cdcc6ab5..b7776871 100644
--- a/lightning_transformers/task/nlp/question_answering/datasets/squad/data.py
+++ b/lightning_transformers/task/nlp/question_answering/datasets/squad/data.py
@@ -25,8 +25,8 @@
 
 
 class SquadDataModule(QuestionAnsweringDataModule):
-    def __init__(self, *args, cfg, **kwargs):
-        super().__init__(*args, cfg=cfg, **kwargs)
+    def __init__(self, *args, dataset_name: str = "squad", **kwargs):
+        super().__init__(*args, dataset_name=dataset_name, **kwargs)
 
     @staticmethod
     def convert_to_train_features(*args, **kwargs):
@@ -49,9 +49,9 @@ def postprocess_func(
             answer_column_name=self.answer_column_name,
             features=validation_dataset,
             examples=original_validation_dataset,
-            version_2_with_negative=self.cfg.version_2_with_negative,
-            n_best_size=self.cfg.n_best_size,
-            max_answer_length=self.cfg.max_answer_length,
-            null_score_diff_threshold=self.cfg.null_score_diff_threshold,
-            output_dir=self.cfg.output_dir,
+            version_2_with_negative=self.version_2_with_negative,
+            n_best_size=self.n_best_size,
+            max_answer_length=self.max_answer_length,
+            null_score_diff_threshold=self.null_score_diff_threshold,
+            output_dir=self.output_dir,
         )
diff --git a/lightning_transformers/task/nlp/question_answering/model.py b/lightning_transformers/task/nlp/question_answering/model.py
index 4a635870..bc181863 100644
--- a/lightning_transformers/task/nlp/question_answering/model.py
+++ b/lightning_transformers/task/nlp/question_answering/model.py
@@ -29,10 +29,10 @@ class QuestionAnsweringTransformer(TaskTransformer):
     """Defines ``LightningModule`` for the Question Answering Task.
 
     Args:
-        *args: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
+        *args: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
         downstream_model_type: Downstream HuggingFace AutoModel to load.
             (default ``transformers.AutoModelForQuestionAnswering``)
-        **kwargs: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
+        **kwargs: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
     """
 
     def __init__(
diff --git a/lightning_transformers/task/nlp/summarization/__init__.py b/lightning_transformers/task/nlp/summarization/__init__.py
index 07b75074..c64b1226 100644
--- a/lightning_transformers/task/nlp/summarization/__init__.py
+++ b/lightning_transformers/task/nlp/summarization/__init__.py
@@ -1,7 +1,3 @@
-from lightning_transformers.task.nlp.summarization.config import (  # noqa: F401
-    SummarizationConfig,
-    SummarizationDataConfig,
-)
 from lightning_transformers.task.nlp.summarization.data import SummarizationDataModule  # noqa: F401
 from lightning_transformers.task.nlp.summarization.datasets import (  # noqa: F401
     CNNDailyMailSummarizationDataModule,
diff --git a/lightning_transformers/task/nlp/summarization/config.py b/lightning_transformers/task/nlp/summarization/config.py
deleted file mode 100644
index 51336fdc..00000000
--- a/lightning_transformers/task/nlp/summarization/config.py
+++ /dev/null
@@ -1,25 +0,0 @@
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from dataclasses import dataclass
-
-from lightning_transformers.core.seq2seq.config import Seq2SeqConfig, Seq2SeqDataConfig
-
-
-@dataclass
-class SummarizationConfig(Seq2SeqConfig):
-    use_stemmer: bool = True
-
-
-class SummarizationDataConfig(Seq2SeqDataConfig):
-    ...
diff --git a/lightning_transformers/task/nlp/summarization/data.py b/lightning_transformers/task/nlp/summarization/data.py
index 9a295ade..a75906bb 100644
--- a/lightning_transformers/task/nlp/summarization/data.py
+++ b/lightning_transformers/task/nlp/summarization/data.py
@@ -12,13 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from lightning_transformers.core.seq2seq.data import Seq2SeqDataModule
-from lightning_transformers.task.nlp.summarization.config import SummarizationDataConfig
 
 
 class SummarizationDataModule(Seq2SeqDataModule):
     """Defines the ``LightningDataModule`` for Summarization Datasets."""
 
-    cfg: SummarizationDataConfig
-
-    def __init__(self, *args, cfg: SummarizationDataConfig = SummarizationDataConfig(), **kwargs) -> None:
-        super().__init__(*args, cfg=cfg, **kwargs)
+    ...
diff --git a/lightning_transformers/task/nlp/summarization/datasets/cnn_dailymail.py b/lightning_transformers/task/nlp/summarization/datasets/cnn_dailymail.py
index 6cdce740..32eece09 100644
--- a/lightning_transformers/task/nlp/summarization/datasets/cnn_dailymail.py
+++ b/lightning_transformers/task/nlp/summarization/datasets/cnn_dailymail.py
@@ -17,6 +17,9 @@
 
 
 class CNNDailyMailSummarizationDataModule(SummarizationDataModule):
+    def __init__(self, *args, dataset_name: str = "cnn_dailymail", **kwargs):
+        super().__init__(*args, dataset_name=dataset_name, **kwargs)
+
     @property
     def source_target_column_names(self) -> Tuple[str, str]:
         return "article", "highlights"
diff --git a/lightning_transformers/task/nlp/summarization/datasets/xsum.py b/lightning_transformers/task/nlp/summarization/datasets/xsum.py
index 51e60f92..fda9797c 100644
--- a/lightning_transformers/task/nlp/summarization/datasets/xsum.py
+++ b/lightning_transformers/task/nlp/summarization/datasets/xsum.py
@@ -17,6 +17,9 @@
 
 
 class XsumSummarizationDataModule(SummarizationDataModule):
+    def __init__(self, *args, dataset_name: str = "xsum", **kwargs):
+        super().__init__(*args, dataset_name=dataset_name, **kwargs)
+
     @property
     def source_target_column_names(self) -> Tuple[str, str]:
         return "document", "summary"
diff --git a/lightning_transformers/task/nlp/summarization/model.py b/lightning_transformers/task/nlp/summarization/model.py
index 11041e35..a1dc486c 100644
--- a/lightning_transformers/task/nlp/summarization/model.py
+++ b/lightning_transformers/task/nlp/summarization/model.py
@@ -17,7 +17,6 @@
 from torchmetrics.text.rouge import ROUGEScore
 
 from lightning_transformers.core.seq2seq.model import Seq2SeqTransformer
-from lightning_transformers.task.nlp.summarization.config import SummarizationConfig
 
 if TYPE_CHECKING:
     from transformers import AutoModel
@@ -27,21 +26,23 @@ class SummarizationTransformer(Seq2SeqTransformer):
     """Defines ``LightningModule`` for the Summarization Task.
 
     Args:
-        *args: :class:`lightning_transformers.core.nlp.seq2seq.Seq2SeqTransformer` arguments.
+        *args: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
         downstream_model_type: Downstream HuggingFace AutoModel to load.
             (default ``transformers.AutoModelForSeq2SeqLM``)
-        **kwargs: :class:`lightning_transformers.core.nlp.seq2seq.Seq2SeqTransformer` arguments.
+        use_stemmer: Use Porter stemmer to strip word suffixes to improve matching.
+        **kwargs: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
     """
 
     def __init__(
         self,
         *args,
         downstream_model_type: Type["AutoModel"] = transformers.AutoModelForSeq2SeqLM,
-        cfg: SummarizationConfig = SummarizationConfig(),
+        use_stemmer: bool = True,
         **kwargs
     ) -> None:
-        super().__init__(downstream_model_type, *args, cfg=cfg, **kwargs)
+        super().__init__(downstream_model_type, *args, **kwargs)
         self.rouge = None
+        self.use_stemmer = use_stemmer
 
     def compute_generate_metrics(self, batch, prefix):
         tgt_lns = self.tokenize_labels(batch["labels"])
@@ -50,7 +51,7 @@ def compute_generate_metrics(self, batch, prefix):
         self.log_dict(result, on_step=False, on_epoch=True)
 
     def configure_metrics(self, stage: str):
-        self.rouge = ROUGEScore(use_stemmer=self.cfg.use_stemmer)
+        self.rouge = ROUGEScore(use_stemmer=self.use_stemmer)
 
     @property
     def hf_pipeline_task(self) -> str:
diff --git a/lightning_transformers/task/nlp/text_classification/__init__.py b/lightning_transformers/task/nlp/text_classification/__init__.py
index 908e094e..f4dcf5fe 100644
--- a/lightning_transformers/task/nlp/text_classification/__init__.py
+++ b/lightning_transformers/task/nlp/text_classification/__init__.py
@@ -1,3 +1,2 @@
-from lightning_transformers.task.nlp.text_classification.config import TextClassificationDataConfig  # noqa: F401
 from lightning_transformers.task.nlp.text_classification.data import TextClassificationDataModule  # noqa: F401
 from lightning_transformers.task.nlp.text_classification.model import TextClassificationTransformer  # noqa: F401
diff --git a/lightning_transformers/task/nlp/text_classification/config.py b/lightning_transformers/task/nlp/text_classification/config.py
deleted file mode 100644
index 84621274..00000000
--- a/lightning_transformers/task/nlp/text_classification/config.py
+++ /dev/null
@@ -1,8 +0,0 @@
-from dataclasses import dataclass
-
-from lightning_transformers.core.config import TransformerDataConfig
-
-
-@dataclass
-class TextClassificationDataConfig(TransformerDataConfig):
-    ...
diff --git a/lightning_transformers/task/nlp/text_classification/data.py b/lightning_transformers/task/nlp/text_classification/data.py
index 786a0bfd..46a16fa8 100644
--- a/lightning_transformers/task/nlp/text_classification/data.py
+++ b/lightning_transformers/task/nlp/text_classification/data.py
@@ -18,16 +18,13 @@
 from transformers import PreTrainedTokenizerBase
 
 from lightning_transformers.core import TransformerDataModule
-from lightning_transformers.task.nlp.text_classification.config import TextClassificationDataConfig
 
 
 class TextClassificationDataModule(TransformerDataModule):
     """Defines the ``LightningDataModule`` for Text Classification Datasets."""
 
-    cfg: TextClassificationDataConfig
-
-    def __init__(self, *args, cfg: TextClassificationDataConfig = TextClassificationDataConfig(), **kwargs) -> None:
-        super().__init__(*args, cfg=cfg, **kwargs)
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
         self.labels = None
 
     def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset:
@@ -36,9 +33,9 @@ def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset
             dataset,
             tokenizer=self.tokenizer,
             input_feature_fields=input_feature_fields,
-            padding=self.cfg.padding,
-            truncation=self.cfg.truncation,
-            max_length=self.cfg.max_length,
+            padding=self.padding,
+            truncation=self.truncation,
+            max_length=self.max_length,
         )
         cols_to_keep = [
             x for x in ["input_ids", "attention_mask", "token_type_ids", "labels"] if x in dataset["train"].features
diff --git a/lightning_transformers/task/nlp/text_classification/model.py b/lightning_transformers/task/nlp/text_classification/model.py
index 228eac8d..49391293 100644
--- a/lightning_transformers/task/nlp/text_classification/model.py
+++ b/lightning_transformers/task/nlp/text_classification/model.py
@@ -27,10 +27,10 @@ class TextClassificationTransformer(TaskTransformer):
     """Defines ``LightningModule`` for the Text Classification Task.
 
     Args:
-        *args: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
+        *args: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
         downstream_model_type: Downstream HuggingFace AutoModel to load.
             (default ``transformers.AutoModelForSequenceClassification``)
-        **kwargs: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
+        **kwargs: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
     """
 
     def __init__(
diff --git a/lightning_transformers/task/nlp/token_classification/__init__.py b/lightning_transformers/task/nlp/token_classification/__init__.py
index 5aca45a5..6239a00b 100644
--- a/lightning_transformers/task/nlp/token_classification/__init__.py
+++ b/lightning_transformers/task/nlp/token_classification/__init__.py
@@ -1,3 +1,2 @@
-from lightning_transformers.task.nlp.token_classification.config import TokenClassificationDataConfig  # noqa: F401
 from lightning_transformers.task.nlp.token_classification.data import TokenClassificationDataModule  # noqa: F401
 from lightning_transformers.task.nlp.token_classification.model import TokenClassificationTransformer  # noqa: F401
diff --git a/lightning_transformers/task/nlp/token_classification/config.py b/lightning_transformers/task/nlp/token_classification/config.py
deleted file mode 100644
index de49e702..00000000
--- a/lightning_transformers/task/nlp/token_classification/config.py
+++ /dev/null
@@ -1,22 +0,0 @@
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from dataclasses import dataclass
-
-from lightning_transformers.core.config import TransformerDataConfig
-
-
-@dataclass
-class TokenClassificationDataConfig(TransformerDataConfig):
-    task_name: str = "ner"
-    label_all_tokens: bool = False
diff --git a/lightning_transformers/task/nlp/token_classification/data.py b/lightning_transformers/task/nlp/token_classification/data.py
index e5361455..b918a7d0 100644
--- a/lightning_transformers/task/nlp/token_classification/data.py
+++ b/lightning_transformers/task/nlp/token_classification/data.py
@@ -19,7 +19,6 @@
 from transformers import DataCollatorForTokenClassification, PreTrainedTokenizerBase
 
 from lightning_transformers.core import TransformerDataModule
-from lightning_transformers.task.nlp.token_classification.config import TokenClassificationDataConfig
 
 
 class TokenClassificationDataModule(TransformerDataModule):
@@ -27,16 +26,14 @@ class TokenClassificationDataModule(TransformerDataModule):
 
     Args:
         *args: ``HFDataModule`` specific arguments.
-        cfg: Contains data specific parameters when processing/loading the dataset
-            (Default ``TokenClassificationDataConfig``)
         **kwargs: ``HFDataModule`` specific arguments.
     """
 
-    cfg: TokenClassificationDataConfig
-
-    def __init__(self, *args, cfg: TokenClassificationDataConfig = TokenClassificationDataConfig(), **kwargs) -> None:
+    def __init__(self, *args, task_name: str = "ner", label_all_tokens: bool = False, **kwargs) -> None:
         self.labels = None
-        super().__init__(*args, cfg=cfg, **kwargs)
+        super().__init__(*args, **kwargs)
+        self.task_name = task_name
+        self.label_all_tokens = label_all_tokens
 
     def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset:
         features, label_column_name, text_column_name = self._setup_input_fields(dataset, stage)
@@ -46,8 +43,8 @@ def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset
         convert_to_features = partial(
             TokenClassificationDataModule.convert_to_features,
             tokenizer=self.tokenizer,
-            padding=self.cfg.padding,
-            label_all_tokens=self.cfg.label_all_tokens,
+            padding=self.padding,
+            label_all_tokens=self.label_all_tokens,
             label_to_id=self.label_to_id,
             text_column_name=text_column_name,
             label_column_name=label_column_name,
@@ -55,8 +52,8 @@ def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset
         dataset = dataset.map(
             convert_to_features,
             batched=True,
-            num_proc=self.cfg.preprocessing_num_workers,
-            load_from_cache_file=self.cfg.load_from_cache_file,
+            num_proc=self.preprocessing_num_workers,
+            load_from_cache_file=self.load_from_cache_file,
         )
         cols_to_keep = [
             x
@@ -71,9 +68,7 @@ def _setup_input_fields(self, dataset, stage):
         column_names = dataset[split].column_names
         features = dataset[split].features
         text_column_name = "tokens" if "tokens" in column_names else column_names[0]
-        label_column_name = (
-            f"{self.cfg.task_name}_tags" if f"{self.cfg.task_name}_tags" in column_names else column_names[1]
-        )
+        label_column_name = f"{self.task_name}_tags" if f"{self.task_name}_tags" in column_names else column_names[1]
         return features, label_column_name, text_column_name
 
     def _prepare_labels(self, dataset, features, label_column_name):
diff --git a/lightning_transformers/task/nlp/token_classification/model.py b/lightning_transformers/task/nlp/token_classification/model.py
index 4a2b0286..a475deb8 100644
--- a/lightning_transformers/task/nlp/token_classification/model.py
+++ b/lightning_transformers/task/nlp/token_classification/model.py
@@ -27,10 +27,10 @@ class TokenClassificationTransformer(TaskTransformer):
     """Defines ``LightningModule`` for the Text Classification Task.
 
     Args:
-        *args: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
+        *args: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
         downstream_model_type: Downstream HuggingFace AutoModel to load.
             (default ``transformers.AutoModelForTokenClassification``)
-        **kwargs: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
+        **kwargs: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
     """
 
     def __init__(
diff --git a/lightning_transformers/task/nlp/translation/__init__.py b/lightning_transformers/task/nlp/translation/__init__.py
index 205f172e..2e7a8f1f 100644
--- a/lightning_transformers/task/nlp/translation/__init__.py
+++ b/lightning_transformers/task/nlp/translation/__init__.py
@@ -1,4 +1,3 @@
-from lightning_transformers.task.nlp.translation.config import TranslationConfig, TranslationDataConfig  # noqa: F401
 from lightning_transformers.task.nlp.translation.data import TranslationDataModule  # noqa: F401
 from lightning_transformers.task.nlp.translation.datasets import WMT16TranslationDataModule  # noqa: F401
 from lightning_transformers.task.nlp.translation.model import TranslationTransformer  # noqa: F401
diff --git a/lightning_transformers/task/nlp/translation/config.py b/lightning_transformers/task/nlp/translation/config.py
deleted file mode 100644
index 251a205f..00000000
--- a/lightning_transformers/task/nlp/translation/config.py
+++ /dev/null
@@ -1,28 +0,0 @@
-# Copyright The PyTorch Lightning team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from dataclasses import dataclass
-
-from lightning_transformers.core.seq2seq.config import Seq2SeqConfig, Seq2SeqDataConfig
-
-
-@dataclass
-class TranslationConfig(Seq2SeqConfig):
-    n_gram: int = 4
-    smooth: bool = False
-
-
-@dataclass
-class TranslationDataConfig(Seq2SeqDataConfig):
-    source_language: str = ""
-    target_language: str = ""
diff --git a/lightning_transformers/task/nlp/translation/data.py b/lightning_transformers/task/nlp/translation/data.py
index 608b49bb..9bc8a808 100644
--- a/lightning_transformers/task/nlp/translation/data.py
+++ b/lightning_transformers/task/nlp/translation/data.py
@@ -14,7 +14,6 @@
 from typing import Tuple
 
 from lightning_transformers.core.seq2seq.data import Seq2SeqDataModule
-from lightning_transformers.task.nlp.translation import TranslationDataConfig
 
 
 class TranslationDataModule(Seq2SeqDataModule):
@@ -22,16 +21,14 @@ class TranslationDataModule(Seq2SeqDataModule):
 
     Args:
         *args: ``Seq2SeqDataModule`` specific arguments.
-        cfg: Contains data specific parameters when processing/loading the dataset
-            (Default ``TranslationDataConfig``)
         **kwargs: ``Seq2SeqDataModule`` specific arguments.
     """
 
-    cfg: TranslationDataConfig
-
-    def __init__(self, *args, cfg: TranslationDataConfig = TranslationDataConfig(), **kwargs) -> None:
-        super().__init__(*args, cfg=cfg, **kwargs)
+    def __init__(self, *args, source_language: str = "", target_language: str = "", **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.source_language = source_language
+        self.target_language = target_language
 
     @property
     def source_target_column_names(self) -> Tuple[str, str]:
-        return self.cfg.source_language, self.cfg.target_language
+        return self.source_language, self.target_language
diff --git a/lightning_transformers/task/nlp/translation/datasets/wmt16.py b/lightning_transformers/task/nlp/translation/datasets/wmt16.py
index 0c4df0f1..1f7157e3 100644
--- a/lightning_transformers/task/nlp/translation/datasets/wmt16.py
+++ b/lightning_transformers/task/nlp/translation/datasets/wmt16.py
@@ -19,6 +19,9 @@
 
 
 class WMT16TranslationDataModule(TranslationDataModule):
+    def __init__(self, dataset_name: str = "wmt16", *args, **kwargs) -> None:
+        super().__init__(*args, dataset_name=dataset_name, **kwargs)
+
     @staticmethod
     def convert_to_features(
         examples: Any,
diff --git a/lightning_transformers/task/nlp/translation/model.py b/lightning_transformers/task/nlp/translation/model.py
index 0793adae..4a60af10 100644
--- a/lightning_transformers/task/nlp/translation/model.py
+++ b/lightning_transformers/task/nlp/translation/model.py
@@ -18,7 +18,7 @@
 from transformers import MBartTokenizer
 
 from lightning_transformers.core.seq2seq.model import Seq2SeqTransformer
-from lightning_transformers.task.nlp.translation.config import TranslationConfig, TranslationDataConfig
+from lightning_transformers.task.nlp.translation import TranslationDataModule
 
 if TYPE_CHECKING:
     from transformers import AutoModel
@@ -28,21 +28,26 @@ class TranslationTransformer(Seq2SeqTransformer):
     """Defines ``LightningModule`` for the Translation Task.
 
     Args:
-        *args: :class:`lightning_transformers.core.nlp.seq2seq.Seq2SeqTransformer` arguments.
+        *args: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
         downstream_model_type: Downstream HuggingFace AutoModel to load.
             (default ``transformers.AutoModelForSeq2SeqLM``)
-        **kwargs: :class:`lightning_transformers.core.nlp.seq2seq.Seq2SeqTransformer` arguments.
+        n_gram: Gram value ranged from 1 to 4.
+        smooth: Whether or not to apply smoothing.
+        **kwargs: :class:`lightning_transformers.core.model.TaskTransformer` arguments.
     """
 
     def __init__(
         self,
         *args,
         downstream_model_type: Type["AutoModel"] = transformers.AutoModelForSeq2SeqLM,
-        cfg: TranslationConfig = TranslationConfig(),
+        n_gram: int = 4,
+        smooth: bool = False,
         **kwargs,
     ) -> None:
-        super().__init__(downstream_model_type, *args, cfg=cfg, **kwargs)
+        super().__init__(downstream_model_type, *args, **kwargs)
         self.bleu = None
+        self.n_gram = n_gram
+        self.smooth = smooth
 
     def compute_generate_metrics(self, batch, prefix):
         tgt_lns = self.tokenize_labels(batch["labels"])
@@ -52,13 +57,13 @@ def compute_generate_metrics(self, batch, prefix):
         self.log(f"{prefix}_bleu_score", result, on_step=False, on_epoch=True, prog_bar=True)
 
     def configure_metrics(self, stage: str):
-        self.bleu = BLEUScore(self.cfg.n_gram, self.cfg.smooth)
+        self.bleu = BLEUScore(self.n_gram, self.smooth)
 
     def initialize_model_specific_parameters(self):
         super().initialize_model_specific_parameters()
         if isinstance(self.tokenizer, MBartTokenizer):
-            cfg: TranslationDataConfig = self.trainer.datamodule.cfg
-            tgt_lang = cfg.target_language
+            dm: TranslationDataModule = self.trainer.datamodule
+            tgt_lang = dm.target_language
             # set decoder_start_token_id for MBart
             if self.model.config.decoder_start_token_id is None:
                 assert tgt_lang is not None, "mBart requires --target_language"
diff --git a/lightning_transformers/task/vision/image_classification/__init__.py b/lightning_transformers/task/vision/image_classification/__init__.py
index c7412818..a35364f2 100644
--- a/lightning_transformers/task/vision/image_classification/__init__.py
+++ b/lightning_transformers/task/vision/image_classification/__init__.py
@@ -1,3 +1,2 @@
-from lightning_transformers.task.vision.image_classification.config import ImageClassificationDataConfig  # noqa: F401
 from lightning_transformers.task.vision.image_classification.data import ImageClassificationDataModule  # noqa: F401
 from lightning_transformers.task.vision.image_classification.model import ImageClassificationTransformer  # noqa: F401
diff --git a/lightning_transformers/task/vision/image_classification/config.py b/lightning_transformers/task/vision/image_classification/config.py
deleted file mode 100644
index c7eff257..00000000
--- a/lightning_transformers/task/vision/image_classification/config.py
+++ /dev/null
@@ -1,8 +0,0 @@
-from dataclasses import dataclass
-
-from lightning_transformers.core.config import TransformerDataConfig
-
-
-@dataclass
-class ImageClassificationDataConfig(TransformerDataConfig):
-    ...
diff --git a/lightning_transformers/task/vision/image_classification/data.py b/lightning_transformers/task/vision/image_classification/data.py
index a150bbd4..f6db3193 100644
--- a/lightning_transformers/task/vision/image_classification/data.py
+++ b/lightning_transformers/task/vision/image_classification/data.py
@@ -17,18 +17,13 @@
 from pytorch_lightning.utilities import rank_zero_warn
 
 from lightning_transformers.core import TransformerDataModule
-from lightning_transformers.task.vision.image_classification.config import ImageClassificationDataConfig
 
 
 class ImageClassificationDataModule(TransformerDataModule):
     """Defines the ``LightningDataModule`` for Image Classification Datasets."""
 
-    cfg: ImageClassificationDataConfig
-
-    def __init__(
-        self, feature_extractor, *args, cfg: ImageClassificationDataConfig = ImageClassificationDataConfig(), **kwargs
-    ) -> None:
-        super().__init__(tokenizer=feature_extractor, *args, cfg=cfg, **kwargs)
+    def __init__(self, feature_extractor, *args, **kwargs) -> None:
+        super().__init__(tokenizer=feature_extractor, *args, **kwargs)
         self.labels = None
 
     def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset:


Test patch:
diff --git a/tests/core/nlp/test_nlp_model.py b/tests/core/nlp/test_nlp_model.py
index 8c85bed8..3471c98e 100644
--- a/tests/core/nlp/test_nlp_model.py
+++ b/tests/core/nlp/test_nlp_model.py
@@ -23,7 +23,6 @@
 
 from lightning_transformers.core import TaskTransformer
 from lightning_transformers.task.nlp.text_classification import (
-    TextClassificationDataConfig,
     TextClassificationDataModule,
     TextClassificationTransformer,
 )
@@ -124,16 +123,14 @@ def hf_pipeline_task(self):
 def test_task_transformer_default_optimizer_scheduler(hf_cache_path):
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     dm = TextClassificationDataModule(
-        cfg=TextClassificationDataConfig(
-            batch_size=1,
-            dataset_name="glue",
-            dataset_config_name="sst2",
-            max_length=512,
-            limit_test_samples=64,
-            limit_val_samples=64,
-            limit_train_samples=64,
-            cache_dir=hf_cache_path,
-        ),
+        batch_size=1,
+        dataset_name="glue",
+        dataset_config_name="sst2",
+        max_length=512,
+        limit_test_samples=64,
+        limit_val_samples=64,
+        limit_train_samples=64,
+        cache_dir=hf_cache_path,
         tokenizer=tokenizer,
     )
     model = TextClassificationTransformer(pretrained_model_name_or_path="prajjwal1/bert-tiny")
diff --git a/tests/core/test_callbacks.py b/tests/core/test_callbacks.py
index 80e6198f..e8dabfc5 100644
--- a/tests/core/test_callbacks.py
+++ b/tests/core/test_callbacks.py
@@ -8,7 +8,6 @@
 from lightning_transformers.callbacks import TransformerSparseMLCallback
 from lightning_transformers.core.loggers import WABLogger
 from lightning_transformers.task.nlp.text_classification import (
-    TextClassificationDataConfig,
     TextClassificationDataModule,
     TextClassificationTransformer,
 )
@@ -75,16 +74,14 @@ def test_training_steps(max_epochs, num_processes, limit_train_batches, modifier
 
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     dm = TextClassificationDataModule(
-        cfg=TextClassificationDataConfig(
-            batch_size=1,
-            dataset_name="glue",
-            dataset_config_name="sst2",
-            max_length=512,
-            limit_test_samples=64,
-            limit_val_samples=64,
-            limit_train_samples=64,
-            cache_dir=hf_cache_path,
-        ),
+        batch_size=1,
+        dataset_name="glue",
+        dataset_config_name="sst2",
+        max_length=512,
+        limit_test_samples=64,
+        limit_val_samples=64,
+        limit_train_samples=64,
+        cache_dir=hf_cache_path,
         tokenizer=tokenizer,
     )
     model = TextClassificationTransformer(pretrained_model_name_or_path="prajjwal1/bert-tiny")
diff --git a/tests/core/test_model_step.py b/tests/core/test_model_step.py
index 821a80ac..7ed4c552 100644
--- a/tests/core/test_model_step.py
+++ b/tests/core/test_model_step.py
@@ -3,7 +3,6 @@
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.text_classification import (
-    TextClassificationDataConfig,
     TextClassificationDataModule,
     TextClassificationTransformer,
 )
@@ -42,16 +41,14 @@ def test_training_and_warmup_steps(
 ):
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     dm = TextClassificationDataModule(
-        cfg=TextClassificationDataConfig(
-            batch_size=1,
-            dataset_name="glue",
-            dataset_config_name="sst2",
-            max_length=512,
-            limit_test_samples=64,
-            limit_val_samples=64,
-            limit_train_samples=64,
-            cache_dir=hf_cache_path,
-        ),
+        batch_size=1,
+        dataset_name="glue",
+        dataset_config_name="sst2",
+        max_length=512,
+        limit_test_samples=64,
+        limit_val_samples=64,
+        limit_train_samples=64,
+        cache_dir=hf_cache_path,
         tokenizer=tokenizer,
     )
     dm.setup("fit")
diff --git a/tests/task/nlp/test_language_modeling.py b/tests/task/nlp/test_language_modeling.py
index 7004c973..515b054e 100644
--- a/tests/task/nlp/test_language_modeling.py
+++ b/tests/task/nlp/test_language_modeling.py
@@ -8,11 +8,7 @@
 import transformers
 from transformers import AutoTokenizer
 
-from lightning_transformers.task.nlp.language_modeling import (
-    LanguageModelingDataConfig,
-    LanguageModelingDataModule,
-    LanguageModelingTransformer,
-)
+from lightning_transformers.task.nlp.language_modeling import LanguageModelingDataModule, LanguageModelingTransformer
 from lightning_transformers.utilities.imports import _ACCELERATE_AVAILABLE
 
 if _ACCELERATE_AVAILABLE:
@@ -25,13 +21,11 @@ def test_smoke_train(hf_cache_path, streaming):
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     model = LanguageModelingTransformer(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     dm = LanguageModelingDataModule(
-        cfg=LanguageModelingDataConfig(
-            batch_size=1,
-            dataset_name="wikitext",
-            dataset_config_name="wikitext-2-raw-v1",
-            cache_dir=hf_cache_path,
-            streaming=streaming,
-        ),
+        batch_size=1,
+        dataset_name="wikitext",
+        dataset_config_name="wikitext-2-raw-v1",
+        cache_dir=hf_cache_path,
+        streaming=streaming,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(max_steps=10, max_epochs=1, limit_val_batches=1, limit_test_batches=1, limit_train_batches=1)
@@ -55,10 +49,9 @@ def test_model_has_correct_cfg():
     assert isinstance(model.model, transformers.BertLMHeadModel)
 
 
-def test_datamodule_has_correct_cfg():
+def test_datamodule_has_tokenizer():
     tokenizer = MagicMock()
     dm = LanguageModelingDataModule(tokenizer)
-    assert isinstance(dm.cfg, LanguageModelingDataConfig)
     assert dm.tokenizer is tokenizer
 
 
diff --git a/tests/task/nlp/test_masked_language_modeling.py b/tests/task/nlp/test_masked_language_modeling.py
index bb316af5..2857274b 100644
--- a/tests/task/nlp/test_masked_language_modeling.py
+++ b/tests/task/nlp/test_masked_language_modeling.py
@@ -7,7 +7,6 @@
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.masked_language_modeling import (
-    MaskedLanguageModelingDataConfig,
     MaskedLanguageModelingDataModule,
     MaskedLanguageModelingTransformer,
 )
@@ -18,12 +17,10 @@ def test_smoke_train(hf_cache_path):
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     model = MaskedLanguageModelingTransformer(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     dm = MaskedLanguageModelingDataModule(
-        cfg=MaskedLanguageModelingDataConfig(
-            batch_size=1,
-            dataset_name="wikitext",
-            dataset_config_name="wikitext-2-raw-v1",
-            cache_dir=hf_cache_path,
-        ),
+        batch_size=1,
+        dataset_name="wikitext",
+        dataset_config_name="wikitext-2-raw-v1",
+        cache_dir=hf_cache_path,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(fast_dev_run=True)
@@ -47,8 +44,7 @@ def test_model_has_correct_cfg():
     assert isinstance(model.model, transformers.BertForMaskedLM)
 
 
-def test_datamodule_has_correct_cfg():
+def test_datamodule_has_tokenizer():
     tokenizer = MagicMock()
     dm = MaskedLanguageModelingDataModule(tokenizer)
-    assert isinstance(dm.cfg, MaskedLanguageModelingDataConfig)
     assert dm.tokenizer is tokenizer
diff --git a/tests/task/nlp/test_multiple_choice.py b/tests/task/nlp/test_multiple_choice.py
index 60c3cb78..f1753637 100644
--- a/tests/task/nlp/test_multiple_choice.py
+++ b/tests/task/nlp/test_multiple_choice.py
@@ -7,7 +7,6 @@
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.multiple_choice import (
-    MultipleChoiceDataConfig,
     MultipleChoiceDataModule,
     MultipleChoiceTransformer,
     SwagMultipleChoiceDataModule,
@@ -19,13 +18,10 @@ def test_smoke_train(hf_cache_path):
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     model = MultipleChoiceTransformer(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     dm = SwagMultipleChoiceDataModule(
-        cfg=MultipleChoiceDataConfig(
-            batch_size=1,
-            dataset_name="swag",
-            dataset_config_name="regular",
-            padding=False,
-            cache_dir=hf_cache_path,
-        ),
+        batch_size=1,
+        dataset_config_name="regular",
+        padding=False,
+        cache_dir=hf_cache_path,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(fast_dev_run=True)
@@ -38,8 +34,7 @@ def test_model_has_correct_cfg():
     assert isinstance(model.model, transformers.BertForMultipleChoice)
 
 
-def test_datamodule_has_correct_cfg():
+def test_datamodule_has_tokenizer():
     tokenizer = MagicMock()
     dm = MultipleChoiceDataModule(tokenizer)
-    assert isinstance(dm.cfg, MultipleChoiceDataConfig)
     assert dm.tokenizer is tokenizer
diff --git a/tests/task/nlp/test_pipeline.py b/tests/task/nlp/test_pipeline.py
index 4cd51327..3ac51738 100644
--- a/tests/task/nlp/test_pipeline.py
+++ b/tests/task/nlp/test_pipeline.py
@@ -1,11 +1,7 @@
 import pytorch_lightning as pl
 from transformers import AutoTokenizer
 
-from lightning_transformers.task.nlp.language_modeling import (
-    LanguageModelingDataConfig,
-    LanguageModelingDataModule,
-    LanguageModelingTransformer,
-)
+from lightning_transformers.task.nlp.language_modeling import LanguageModelingDataModule, LanguageModelingTransformer
 
 
 def test_kwargs_load_from_checkpoint(hf_cache_path, tmpdir):
@@ -14,12 +10,10 @@ def test_kwargs_load_from_checkpoint(hf_cache_path, tmpdir):
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="sshleifer/tiny-gpt2")
     model = LanguageModelingTransformer(pretrained_model_name_or_path="sshleifer/tiny-gpt2")
     dm = LanguageModelingDataModule(
-        cfg=LanguageModelingDataConfig(
-            batch_size=1,
-            dataset_name="wikitext",
-            dataset_config_name="wikitext-2-raw-v1",
-            cache_dir=hf_cache_path,
-        ),
+        batch_size=1,
+        dataset_name="wikitext",
+        dataset_config_name="wikitext-2-raw-v1",
+        cache_dir=hf_cache_path,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(default_root_dir=tmpdir, fast_dev_run=True)
diff --git a/tests/task/nlp/test_question_answering.py b/tests/task/nlp/test_question_answering.py
index e6d6d4e4..c08ece2d 100644
--- a/tests/task/nlp/test_question_answering.py
+++ b/tests/task/nlp/test_question_answering.py
@@ -7,7 +7,6 @@
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.question_answering import (
-    QuestionAnsweringDataConfig,
     QuestionAnsweringDataModule,
     QuestionAnsweringTransformer,
     SquadDataModule,
@@ -20,21 +19,18 @@ def test_smoke_train(hf_cache_path):
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="sshleifer/tiny-gpt2")
     model = QuestionAnsweringTransformer(pretrained_model_name_or_path="sshleifer/tiny-gpt2")
     dm = SquadDataModule(
-        cfg=QuestionAnsweringDataConfig(
-            batch_size=1,
-            dataset_name="squad",
-            dataset_config_name="plain_text",
-            max_length=384,
-            version_2_with_negative=False,
-            null_score_diff_threshold=0.0,
-            doc_stride=128,
-            n_best_size=20,
-            max_answer_length=30,
-            limit_train_samples=64,
-            limit_test_samples=64,
-            limit_val_samples=64,
-            cache_dir=hf_cache_path,
-        ),
+        batch_size=1,
+        dataset_config_name="plain_text",
+        max_length=384,
+        version_2_with_negative=False,
+        null_score_diff_threshold=0.0,
+        doc_stride=128,
+        n_best_size=20,
+        max_answer_length=30,
+        limit_train_samples=64,
+        limit_test_samples=64,
+        limit_val_samples=64,
+        cache_dir=hf_cache_path,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(fast_dev_run=True)
@@ -59,8 +55,7 @@ def test_model_has_correct_cfg():
     assert isinstance(model.model, transformers.BertForQuestionAnswering)
 
 
-def test_datamodule_has_correct_cfg():
+def test_datamodule_has_tokenizer():
     tokenizer = MagicMock()
     dm = QuestionAnsweringDataModule(tokenizer)
-    assert isinstance(dm.cfg, QuestionAnsweringDataConfig)
     assert dm.tokenizer is tokenizer
diff --git a/tests/task/nlp/test_summarization.py b/tests/task/nlp/test_summarization.py
index 2f73e13c..39209f68 100644
--- a/tests/task/nlp/test_summarization.py
+++ b/tests/task/nlp/test_summarization.py
@@ -7,8 +7,6 @@
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.summarization import (
-    SummarizationConfig,
-    SummarizationDataConfig,
     SummarizationDataModule,
     SummarizationTransformer,
     XsumSummarizationDataModule,
@@ -20,24 +18,19 @@ def test_smoke_train(hf_cache_path):
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="patrickvonplaten/t5-tiny-random")
     model = SummarizationTransformer(
         pretrained_model_name_or_path="patrickvonplaten/t5-tiny-random",
-        cfg=SummarizationConfig(
-            use_stemmer=True,
-            val_target_max_length=142,
-            num_beams=None,
-            compute_generate_metrics=True,
-        ),
+        use_stemmer=True,
+        val_target_max_length=142,
+        num_beams=None,
+        compute_generate_metrics=True,
     )
     dm = XsumSummarizationDataModule(
-        cfg=SummarizationDataConfig(
-            limit_train_samples=64,
-            limit_val_samples=64,
-            limit_test_samples=64,
-            batch_size=1,
-            dataset_name="xsum",
-            max_source_length=128,
-            max_target_length=128,
-            cache_dir=hf_cache_path,
-        ),
+        limit_train_samples=64,
+        limit_val_samples=64,
+        limit_test_samples=64,
+        batch_size=1,
+        max_source_length=128,
+        max_target_length=128,
+        cache_dir=hf_cache_path,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(fast_dev_run=True)
@@ -62,14 +55,12 @@ def test_smoke_predict():
     assert 2 <= len(output.split()) <= 12
 
 
-def test_model_has_correct_cfg():
+def test_model_has_correct_class():
     model = SummarizationTransformer(pretrained_model_name_or_path="t5-base")
     assert isinstance(model.model, transformers.T5ForConditionalGeneration)
-    assert type(model.cfg) is SummarizationConfig
 
 
-def test_datamodule_has_correct_cfg():
+def test_datamodule_has_tokenizer():
     tokenizer = MagicMock()
     dm = SummarizationDataModule(tokenizer)
-    assert isinstance(dm.cfg, SummarizationDataConfig)
     assert dm.tokenizer is tokenizer
diff --git a/tests/task/nlp/test_text_classification.py b/tests/task/nlp/test_text_classification.py
index 928b2055..ed819764 100644
--- a/tests/task/nlp/test_text_classification.py
+++ b/tests/task/nlp/test_text_classification.py
@@ -10,7 +10,6 @@
 
 from lightning_transformers.plugins.checkpoint import HFSaveCheckpoint
 from lightning_transformers.task.nlp.text_classification import (
-    TextClassificationDataConfig,
     TextClassificationDataModule,
     TextClassificationTransformer,
 )
@@ -20,16 +19,14 @@
 def test_smoke_train(hf_cache_path):
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     dm = TextClassificationDataModule(
-        cfg=TextClassificationDataConfig(
-            batch_size=1,
-            dataset_name="glue",
-            dataset_config_name="sst2",
-            max_length=512,
-            limit_test_samples=64,
-            limit_val_samples=64,
-            limit_train_samples=64,
-            cache_dir=hf_cache_path,
-        ),
+        batch_size=1,
+        dataset_name="glue",
+        dataset_config_name="sst2",
+        max_length=512,
+        limit_test_samples=64,
+        limit_val_samples=64,
+        limit_train_samples=64,
+        cache_dir=hf_cache_path,
         tokenizer=tokenizer,
     )
     model = TextClassificationTransformer(pretrained_model_name_or_path="prajjwal1/bert-tiny")
@@ -40,17 +37,15 @@ def test_smoke_train(hf_cache_path):
 def test_smoke_predict_with_trainer(hf_cache_path):
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     dm = TextClassificationDataModule(
-        cfg=TextClassificationDataConfig(
-            batch_size=1,
-            dataset_name="glue",
-            dataset_config_name="sst2",
-            max_length=512,
-            limit_test_samples=64,
-            limit_val_samples=64,
-            limit_train_samples=64,
-            cache_dir=hf_cache_path,
-            predict_subset_name="test",  # Use the "test" split of the dataset as our prediction subset
-        ),
+        batch_size=1,
+        dataset_name="glue",
+        dataset_config_name="sst2",
+        max_length=512,
+        limit_test_samples=64,
+        limit_val_samples=64,
+        limit_train_samples=64,
+        cache_dir=hf_cache_path,
+        predict_subset_name="test",  # Use the "test" split of the dataset as our prediction subset
         tokenizer=tokenizer,
     )
     model = TextClassificationTransformer(pretrained_model_name_or_path="prajjwal1/bert-tiny")
@@ -76,10 +71,9 @@ def test_model_has_correct_cfg():
     assert isinstance(model.model, transformers.BertForSequenceClassification)
 
 
-def test_datamodule_has_correct_cfg():
+def test_datamodule_has_tokenizer():
     tokenizer = MagicMock()
     dm = TextClassificationDataModule(tokenizer)
-    assert isinstance(dm.cfg, TextClassificationDataConfig)
     assert dm.tokenizer is tokenizer
 
 
@@ -87,16 +81,14 @@ def test_datamodule_has_correct_cfg():
 def test_huggingface_checkpoint_train(hf_cache_path, tmpdir):
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     dm = TextClassificationDataModule(
-        cfg=TextClassificationDataConfig(
-            batch_size=1,
-            dataset_name="glue",
-            dataset_config_name="sst2",
-            max_length=512,
-            limit_test_samples=64,
-            limit_val_samples=64,
-            limit_train_samples=64,
-            cache_dir=hf_cache_path,
-        ),
+        batch_size=1,
+        dataset_name="glue",
+        dataset_config_name="sst2",
+        max_length=512,
+        limit_test_samples=64,
+        limit_val_samples=64,
+        limit_train_samples=64,
+        cache_dir=hf_cache_path,
         tokenizer=tokenizer,
     )
     ckpt_path = os.path.join(tmpdir, "checkpoints")
diff --git a/tests/task/nlp/test_token_classification.py b/tests/task/nlp/test_token_classification.py
index dd7d662a..7d68e19d 100644
--- a/tests/task/nlp/test_token_classification.py
+++ b/tests/task/nlp/test_token_classification.py
@@ -7,7 +7,6 @@
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.token_classification import (
-    TokenClassificationDataConfig,
     TokenClassificationDataModule,
     TokenClassificationTransformer,
 )
@@ -17,18 +16,16 @@
 def test_smoke_train(hf_cache_path):
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     dm = TokenClassificationDataModule(
-        cfg=TokenClassificationDataConfig(
-            batch_size=1,
-            task_name="ner",
-            dataset_name="conll2003",
-            preprocessing_num_workers=1,
-            label_all_tokens=False,
-            revision="master",
-            limit_test_samples=64,
-            limit_val_samples=64,
-            limit_train_samples=64,
-            cache_dir=hf_cache_path,
-        ),
+        batch_size=1,
+        task_name="ner",
+        dataset_name="conll2003",
+        preprocessing_num_workers=1,
+        label_all_tokens=False,
+        revision="master",
+        limit_test_samples=64,
+        limit_val_samples=64,
+        limit_train_samples=64,
+        cache_dir=hf_cache_path,
         tokenizer=tokenizer,
     )
     model = TokenClassificationTransformer(pretrained_model_name_or_path="prajjwal1/bert-tiny", labels=dm.num_classes)
@@ -59,8 +56,7 @@ def test_model_has_correct_cfg():
     assert model.num_labels == 2
 
 
-def test_datamodule_has_correct_cfg():
+def test_datamodule_has_tokenizer():
     tokenizer = MagicMock()
     dm = TokenClassificationDataModule(tokenizer)
-    assert isinstance(dm.cfg, TokenClassificationDataConfig)
     assert dm.tokenizer is tokenizer
diff --git a/tests/task/nlp/test_translation.py b/tests/task/nlp/test_translation.py
index 5991cff2..704b7895 100644
--- a/tests/task/nlp/test_translation.py
+++ b/tests/task/nlp/test_translation.py
@@ -7,8 +7,6 @@
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.translation import (
-    TranslationConfig,
-    TranslationDataConfig,
     TranslationDataModule,
     TranslationTransformer,
     WMT16TranslationDataModule,
@@ -20,19 +18,16 @@ def test_smoke_train(hf_cache_path):
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="patrickvonplaten/t5-tiny-random")
     model = TranslationTransformer(pretrained_model_name_or_path="patrickvonplaten/t5-tiny-random")
     dm = WMT16TranslationDataModule(
-        cfg=TranslationDataConfig(
-            batch_size=1,
-            dataset_name="wmt16",
-            dataset_config_name="ro-en",
-            source_language="en",
-            target_language="ro",
-            cache_dir=hf_cache_path,
-            limit_train_samples=16,
-            limit_val_samples=16,
-            limit_test_samples=16,
-            max_source_length=32,
-            max_target_length=32,
-        ),
+        batch_size=1,
+        dataset_config_name="ro-en",
+        source_language="en",
+        target_language="ro",
+        cache_dir=hf_cache_path,
+        limit_train_samples=16,
+        limit_val_samples=16,
+        limit_test_samples=16,
+        max_source_length=32,
+        max_target_length=32,
         tokenizer=tokenizer,
     )
     trainer = pl.Trainer(fast_dev_run=True)
@@ -51,14 +46,12 @@ def test_smoke_predict():
     assert isinstance(y[0]["translation_text"], str)
 
 
-def test_model_has_correct_cfg():
+def test_model_has_correct_class():
     model = TranslationTransformer(pretrained_model_name_or_path="patrickvonplaten/t5-tiny-random")
     assert isinstance(model.model, transformers.T5ForConditionalGeneration)
-    assert type(model.cfg) is TranslationConfig
 
 
-def test_datamodule_has_correct_cfg():
+def test_datamodule_has_tokenizer():
     tokenizer = MagicMock()
     dm = TranslationDataModule(tokenizer)
-    assert isinstance(dm.cfg, TranslationDataConfig)
     assert dm.tokenizer is tokenizer
diff --git a/tests/task/vision/test_image_classification.py b/tests/task/vision/test_image_classification.py
index a170b5ae..dea9ca8e 100644
--- a/tests/task/vision/test_image_classification.py
+++ b/tests/task/vision/test_image_classification.py
@@ -7,7 +7,6 @@
 from transformers import AutoFeatureExtractor
 
 from lightning_transformers.task.vision.image_classification import (
-    ImageClassificationDataConfig,
     ImageClassificationDataModule,
     ImageClassificationTransformer,
 )
@@ -18,7 +17,8 @@
 def test_smoke_train(hf_cache_path):
     feature_extractor = AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path="nateraw/tiny-vit-random")
     dm = ImageClassificationDataModule(
-        cfg=ImageClassificationDataConfig(batch_size=1, dataset_name="beans"),
+        batch_size=1,
+        dataset_name="beans",
         feature_extractor=feature_extractor,
     )
     model = ImageClassificationTransformer(
@@ -51,8 +51,7 @@ def test_model_has_correct_cfg():
     assert isinstance(model.model, transformers.ViTForImageClassification)
 
 
-def test_datamodule_has_correct_cfg():
+def test_datamodule_has_extractor():
     feature_extractor = MagicMock()
     dm = ImageClassificationDataModule(feature_extractor)
-    assert isinstance(dm.cfg, ImageClassificationDataConfig)
     assert dm.tokenizer is feature_extractor