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

Issue:
Should we get rid of Hydra?
## Motivation

This is motivated by recent development of [LightningCLI](https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_cli.html) which IMO is simpler, and better supported primarily because it exists within Lightning. 

Over time I've noticed that due to the changes upstream in Lightning, hydra default configs are going out of sync every so often.

I also notice that there are fundamental issues such as this https://github.com/PyTorchLightning/lightning-transformers/issues/149 where the confs exist within the package. We assume that the user will always clone the repo and make their modifications there. I would be curious if this works for users currently.

### Pitch

Remove Hydra, treating this library more as a set of components people can import into their own scripts. If we'd like to keep the ability to run training and testing from the library, we should use the LightningCLI instead of Hydra.

### Alternatives

Keep as is :) 

cc @mathemusician @Borda @carmocca 


Discussion / hints:
Playing devil's advocate:

> Over time I've noticed that due to the changes upstream in Lightning, hydra default configs are going out of sync every so often.

The CLI suffers from exactly this issue too. There's 2 proposals in jsonargparse to improve this a bit but they haven't been worked on at the moment:

https://github.com/omni-us/jsonargparse/issues/79
https://github.com/omni-us/jsonargparse/issues/91

> I also notice that there are fundamental issues such as this #149 where the confs exist within the package. We assume that the user will always clone the repo and make their modifications there.

The CLI would also have this problem if we wanted to distribute its config files. But I guess we wouldn't need to as much if we made the switch.

> treating this library more as a set of components people can import into their own scripts.

This is good regardless of the parsing library used. The choice of Hydra or the LightningCLI will come down to what the users are doing and how.
@SeanNaren @Borda Could you give a quick example of how people would use this library without hydra? If we can come up with a list of use cases where this would make the life of future users easier/simpler, then it would be a lot easier moving this issue forward.

When I started learning about this library, Hydra was the hardest thing to learn since it had the least documentation. Getting rid of it may increase usability, because it would be one less thing to learn. I had initially thought Hydra would make things easier, but it was only easier after learning about the entire library (which took me a month). :/ But that's probably because I'm a slow learner.

I propose an experiment, albeit, an expensive one. We replace Hydra with LightningCLI, and if more people use it, we'll stick with it. If not, we can always go back :0 Converting this library into a `set of components` may be challenging but worthwhile. It might even be merged into Bolts.
Thank you both for your comments! 

Addressing @carmocca comments:

> The CLI suffers from exactly this issue too. There's 2 proposals in jsonargparse to improve this a bit but they haven't been worked on at the moment:

I'm unsure how these are related, it seems the second one is about trying to reproduce a saved config that was run on an earlier version of Lightning? I think this arguably could be expected behaviour. For Lightning-transformers specifically, I meant the fact that we rely on hardcoded parameters within the Trainer [config](https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/trainer/default.yaml). This has been outdated a few times as Lightning deprecates parameters/introduces parameters meaning users couldn't run the code out of the box unless they changed the configs manually.

 > This is good regardless of the parsing library used. The choice of Hydra or the LightningCLI will come down to what the users are doing and how.

Agreed, and currently this is supported (props to @carmocca who ensured that this remained). It can be seen from the test [here](https://github.com/PyTorchLightning/lightning-transformers/blob/master/tests/examples/test_custom_translation.py#L10). 

The objective of this issue is to determine if we believe that making this experience a first class citizen would be beneficial for users and reduce complication by allowing users to pick what they'd like, defaulting to LightningCLI as it's the simpler of the two.

Addressing @mathemusician comments:

> @SeanNaren @Borda Could you give a quick example of how people would use this library without hydra? If we can come up with a list of use cases where this would make the life of future users easier/simpler, then it would be a lot easier moving this issue forward.

The test [here](https://github.com/PyTorchLightning/lightning-transformers/blob/master/tests/examples/test_custom_translation.py#L10) gives us a good sense, however still contains some ugliness due to Hydra. I imagine:

```python
tokenizer = AutoTokenizer.from_pretrained(backbone="patrickvonplaten/t5-tiny-random")
model = TranslationTransformer(
    backbone="patrickvonplaten/t5-tiny-random"
)
dm = TranslationDataModule(
    batch_size=1,
    dataset_name="wmt16",
    dataset_config_name="ro-en",
    source_language="en",
    target_language="ro",
)
trainer = pl.Trainer()

trainer.fit(model, dm)
```

In this sense, we offer just LightningModules and LightningDataModules for the individuals tasks. the trainer can then be swapped out for the LightningCLI. Making this change also removes some of the boilerplate around instantiation etc. If the user wants to swap the optimizer/scheduler, they can do so via overriding `configure_optimizers`.

Personally I'm unsure which direction is better. I believe that simplicity of a library is important, and building frameworks rarely provide long term sustainability unless supported well (and the support for Lightning Transformers has been wishwashy, mostly my fault). 

I like your idea of trying it out, seeing how the code looks like and then making a consensus whether to swap. But right now I think it would be good to ask a few more people to pitch in!

I don't want to sidetrack this conversation but to clarify:

>  I meant the fact that we rely on hardcoded parameters within the Trainer config. This has been outdated a few times as Lightning deprecates parameters/introduces parameters meaning users couldn't run the code out of the box unless they changed the configs manually.

It's the same with the CLI, if you want to save the config inside the repository.
`python script.py --print_config > your_config.yaml` will include all the current arguments for the trainer etc, if they are removed, the config will also become stale.
https://github.com/omni-us/jsonargparse/issues/79 helps by allowing us to specify metadata for this config, most importantly the version that generated it. This could be parsed later to improve the UX when the versions do not match.
https://github.com/omni-us/jsonargparse/issues/91 helps by printing only the arguments that have been customized, which reduces the chances of it breaking due to changes in other non-changed arguments.

Now, back to the main discussion.

I'd say, pick a few common or popular use cases, then we try implementing them with both options and see which is most natural/intuitive/simple, etc. This will also show which of the core components are inflexible.
Thanks for clarifying @carmocca!

Personally would like to give this is a go. I much prefer simpler packages, and I feel we could cut down complexity for elegancy and sustainability :)
Hey guys, I'm considering writing a tiny demo combining transformers with PL CLI. The Doc of Cli didn't explain too much about model inheriting. E.g. I can't use a config Dict anymore and have to write all the arguments in an explicit way. When it comes to multi-level inheriting, I got to write all the parameters manually. This is against simplicity. So is there any "best practice" or working-around of this problem? 

Fix patch:
diff --git a/README.md b/README.md
index 68a8888e..89e91555 100644
--- a/README.md
+++ b/README.md
@@ -64,9 +64,7 @@ dm = TextClassificationDataModule(
     ),
     tokenizer=tokenizer,
 )
-model = TextClassificationTransformer(
-    pretrained_model_name_or_path="bert-base-cased", num_labels=dm.num_classes
-)
+model = TextClassificationTransformer(pretrained_model_name_or_path="bert-base-cased")
 
 trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
 
diff --git a/docs/source/advanced/custom_task.rst b/docs/source/advanced/custom_task.rst
index 262a7cb2..55ece53e 100644
--- a/docs/source/advanced/custom_task.rst
+++ b/docs/source/advanced/custom_task.rst
@@ -11,7 +11,6 @@ through the model, and the loss calculation for a specific task. Below are the s
 
 1. Inherit from Lightning Transformers Base Class
 2. Add custom task logic
-3. Create Hydra config
 
 1. Inherit from Lightning Transformers Base Class
 -------------------------------------------------
@@ -54,40 +53,3 @@ The ``LMHeadAutoModel`` task provides separate keys for the backbone and the ful
                 for param in self.model.parameters():
                     param.add_(torch.randn(param.size()) * 0.1)
             return loss
-
-
-3. Create Hydra Config
-----------------------
-
-Finally to use the Hydra CLI and configs, we would add our own custom yaml file containing the necessary code to run using our task.
-
-We create a file at ``conf/task/nlp/my_language_modeling.yaml`` containing the below config.
-
-.. code-block:: yaml
-
-    # @package task
-    defaults:
-      - nlp/default # Use the defaults from the default config found at `conf/task/nlp/default.yaml`
-    _target_: examples.custom_language_modeling.model.MyLanguageModelingTransformer # path to the class we'd like to instantiate
-    downstream_model_type: transformers.AutoModelForCausalLM
-
-Hydra supports config inheritence, so we could inherit from the language modeling task directly, simplifying our config a bit:
-
-.. code-block:: yaml
-
-    # @package task
-    defaults:
-      - nlp/language_modeling # Use the defaults from the config found at `conf/task/nlp/language_modeling.yaml`
-    _target_: examples.custom_language_modeling.model.MyLanguageModelingTransformer # path to the class we'd like to instantiate
-
-With this in place you can now train using pre-made HuggingFace datasets:
-
-.. code-block:: python
-
-    python train.py task=nlp/my_language_modeling dataset=nlp/language_modeling/wikitext dataset.train_file=train.csv dataset.validation_file=valid.csv
-
-Or with your own files:
-
-.. code-block:: python
-
-    python train.py task=nlp/my_language_modeling dataset.train_file=train.csv dataset.validation_file=valid.csv
diff --git a/docs/source/advanced/nlp/language_modeling_data.rst b/docs/source/advanced/nlp/language_modeling_data.rst
index a9b6130a..f1945b65 100644
--- a/docs/source/advanced/nlp/language_modeling_data.rst
+++ b/docs/source/advanced/nlp/language_modeling_data.rst
@@ -1,18 +1,10 @@
 Language Modeling using Custom Data Processing
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-Below we show an example of how to override data processing logic. In this example, we add a prefix to each block of text used in the language modeling task.
+Below we show how to override data processing logic.
 
-This reflects the idea of passing a conditional term that is used to give the language model context. Check :doc:`/tasks/nlp/language_modeling` for more information around the task.
-
-Ultimately to create your own custom data processing the flow is like this:
-
-1. Extend the ``LanguageModelingDataModule`` base class, Override hooks with your own logic
-2. (Optional) Keep file in the specific task directory
-3. Add a hydra config object to use your new dataset
-
-1. Extend the ``LanguageModelingDataModule`` base class
-"""""""""""""""""""""""""""""""""""""""""""""""""""""""
+Extend the ``LanguageModelingDataModule`` base class
+""""""""""""""""""""""""""""""""""""""""""""""""""""
 
 The base data module can be used to modify this code, and follows a simple pattern. Internally the dataset is loaded via HuggingFace Datasets, which returns an `Apache Arrow Parquet <https://arrow.apache.org/docs/python/generated/pyarrow.parquet.ParquetDataset.html>`_ Dataset. This data format is easy to transform and modify using map functions, which you'll see within the class.
 
@@ -105,33 +97,3 @@ Below we have the pseudo code version to show where most of the changes happened
             }
             result["labels"] = result["input_ids"].copy()
             return result
-
-
-To see the full example, see ``examples/custom/dataset/language_modeling/custom_dataset.py``
-
-2. (Optional) Keep file in the specific task directory
-""""""""""""""""""""""""""""""""""""""""""""""""""""""
-
-This makes tracking of files easier. Our example is stored in ``examples/`` however in reality we would store our DataModule in ``lightning_transformers/task/nlp/language_modeling/custom_dataset.py``.
-
-3. Add a hydra config object to use your new dataset
-""""""""""""""""""""""""""""""""""""""""""""""""""""
-
-Finally to use the Hydra CLI and configs, we would add our own custom yaml file containing the necessary code to run using our dataset.
-
-We create a file at ``conf/datasets/nlp/language_modeling/my_dataset.yaml`` containing the below config.
-
-.. code-block:: yaml
-
-    # @package dataset
-    defaults:
-      - nlp/default # Use the defaults from the default config found at `conf/dataset/nlp/default.yaml`
-    _target_: lightning_transformers.custom_language_modeling.dataset.MyLanguageModelingDataModule # path to the class we'd like to instantiate
-    cfg:
-      block_size: 512 # any parameters you'd like from the inherited config object.
-
-With this in place you can now train using either HuggingFace Datasets or your own custom files.
-
-.. code-block:: bash
-
-    python train.py task=nlp/language_modeling dataset=nlp/language_modeling/my_dataset dataset.cfg.train_file=train.csv dataset.cfg.validation_file=valid.csv
diff --git a/docs/source/advanced/nlp/translation_data.rst b/docs/source/advanced/nlp/translation_data.rst
index f913deba..f2dd0615 100644
--- a/docs/source/advanced/nlp/translation_data.rst
+++ b/docs/source/advanced/nlp/translation_data.rst
@@ -1,16 +1,10 @@
 Translation using Custom Data Processing
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-Below we show an example of how overriding data processing logic, by adding a prefix to the source language sample in translation. Check :doc:`/tasks/nlp/translation` for more information around the task.
+Below we show how to override data processing logic.
 
-Ultimately to create your own custom data processing the flow is like this:
-
-1. Extend the ``TranslationDataModule`` base class, Override hooks with your own logic
-2. (Optional) Keep file in the specific task directory
-3. Add a hydra config object to use your new dataset
-
-1. Extend the ``TranslationDataModule`` base class
-""""""""""""""""""""""""""""""""""""""""""""""""""
+Extend the ``TranslationDataModule`` base class
+"""""""""""""""""""""""""""""""""""""""""""""""
 
 The base data module can be used to modify this code, and follows a simple pattern. Internally the dataset is loaded via HuggingFace Datasets, which returns an `Apache Arrow Parquet <https://arrow.apache.org/docs/python/generated/pyarrow.parquet.ParquetDataset.html>`_ Dataset. This data format is easy to transform and modify using map functions, which you'll see within the class.
 
@@ -63,32 +57,3 @@ Extend ``TranslationDataModule``, like this.
         ...
 
 Make any changes you'd like to the dataset processing via the hooks.
-
-To see the full example, see ``examples/custom/dataset/translation/custom_dataset.py``
-
-2. (Optional) Keep file in the specific task directory
-""""""""""""""""""""""""""""""""""""""""""""""""""""""
-
-This makes tracking of files easier. Our example is stored in ``examples/`` however in reality we would store our DataModule in ``lightning_transformers/task/nlp/translation/datasets/custom_dataset.py``.
-
-3. Add a hydra config object to use your new dataset
-""""""""""""""""""""""""""""""""""""""""""""""""""""
-
-Finally to use the Hydra CLI and configs, we would add our own custom yaml file containing the necessary code to run using our dataset.
-
-We create a file at ``conf/datasets/nlp/translation/my_dataset.yaml`` containing the below config.
-
-.. code-block:: yaml
-
-    # @package dataset
-    defaults:
-      - nlp/default # Use the defaults from the default config found at `conf/dataset/nlp/default.yaml`
-    _target_: examples.custom_translation.dataset.MyTranslationDataModule # path to the class we'd like to instantiate
-    cfg:
-      max_source_length: 128 # any parameters you'd like from the inherited config object.
-
-With this in place you can now train using either HuggingFace Datasets or your own custom files.
-
-.. code-block:: bash
-
-    python train.py task=nlp/translation dataset=nlp/translation/my_dataset dataset.cfg.train_file=train.csv dataset.cfg.validation_file=valid.csv
diff --git a/docs/source/datasets/nlp/custom_subset_names.rst b/docs/source/datasets/nlp/custom_subset_names.rst
index 8924a80b..96853ab4 100644
--- a/docs/source/datasets/nlp/custom_subset_names.rst
+++ b/docs/source/datasets/nlp/custom_subset_names.rst
@@ -4,15 +4,24 @@ Custom Subset Names (Edge Cases such as MNLI)
 Some datasets, such as MNLI when loaded from the Huggingface `datasets` library, have special subset names that don't match the standard train/validation/test convention.
 Specifically, MNLI has two validation and two test sets, with flavors 'matched' and 'mismatched'.
 When using such datasets, you must manually indicate which subset names you want to use for each of train/validation/text.
-For this, you can set the config variables `dataset.cfg.train_subset_name`, `dataset.cfg.validation_subset_name` and `dataset.cfg.test_subset_name`.
 
 An example for how to train and validate on MNLI would the the following:
 
 .. code-block:: python
 
-    python train.py task=nlp/text_classification dataset=nlp/text_classification/glue dataset.cfg.dataset_config_name=mnli ++dataset.cfg.validation_subset_name=validation_matched
+    from lightning_transformers.task.nlp.text_classification import (
+        TextClassificationDataConfig,
+        TextClassificationDataModule,
+        TextClassificationTransformer,
+    )
 
-It also works for train and test subsets, like so:
-
-++dataset.cfg.train_subset_name=name_of_subset
-++dataset.cfg.test_subset_name=name_of_subset
+    dm = TextClassificationDataModule(
+        cfg=TextClassificationDataConfig(
+            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 225114d8..b001b78b 100644
--- a/docs/source/datasets/nlp/language_modeling_data.rst
+++ b/docs/source/datasets/nlp/language_modeling_data.rst
@@ -18,4 +18,16 @@ When specifying the file path with hydra, it is important to use the absolute pa
 
 .. code-block:: python
 
-    python train.py task=nlp/language_modeling dataset.cfg.train_file=abs/path/train.csv dataset.cfg.validation_file=abs/path/valid.csv
+    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"
+        ),
+        tokenizer=tokenizer,
+    )
diff --git a/docs/source/datasets/nlp/multiple_choice_data.rst b/docs/source/datasets/nlp/multiple_choice_data.rst
index fd65b397..aaa6023f 100644
--- a/docs/source/datasets/nlp/multiple_choice_data.rst
+++ b/docs/source/datasets/nlp/multiple_choice_data.rst
@@ -21,4 +21,19 @@ We override the dataset files, allowing us to still use the data transforms defi
 
 .. code-block:: python
 
-    python train.py task=nlp/multiple_choice dataset=language_modeling/race dataset.cfg.train_file=train.json dataset.cfg.validation_file=valid.json
+    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"
+        ),
+        tokenizer=tokenizer,
+    )
diff --git a/docs/source/datasets/nlp/question_answering_data.rst b/docs/source/datasets/nlp/question_answering_data.rst
index 3f7a14d2..fc270ff0 100644
--- a/docs/source/datasets/nlp/question_answering_data.rst
+++ b/docs/source/datasets/nlp/question_answering_data.rst
@@ -22,4 +22,25 @@ We override the dataset files, allowing us to still use the data transforms defi
 
 .. code-block:: python
 
-    python train.py task=nlp/question_answering dataset=nlp/question_answering/squad dataset.cfg.train_file=train.json dataset.cfg.validation_file=valid.json
+
+    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"
+        ),
+        tokenizer=tokenizer,
+    )
diff --git a/docs/source/datasets/nlp/summarization_data.rst b/docs/source/datasets/nlp/summarization_data.rst
index 686b7bcc..388ce27c 100644
--- a/docs/source/datasets/nlp/summarization_data.rst
+++ b/docs/source/datasets/nlp/summarization_data.rst
@@ -14,4 +14,20 @@ We override the dataset files, allowing us to still use the data transforms defi
 
 .. code-block:: python
 
-    python train.py task=nlp/summarization dataset.cfg.train_file=train.json dataset.cfg.validation_file=valid.json
+    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"
+        ),
+        tokenizer=tokenizer,
+    )
diff --git a/docs/source/datasets/nlp/text_classification_data.rst b/docs/source/datasets/nlp/text_classification_data.rst
index 25891a47..529ff6a8 100644
--- a/docs/source/datasets/nlp/text_classification_data.rst
+++ b/docs/source/datasets/nlp/text_classification_data.rst
@@ -14,4 +14,18 @@ The label mapping is automatically generated from the training dataset labels if
 
 .. code-block:: python
 
-    python train.py task=nlp/text_classification dataset.cfg.train_file=train.json dataset.cfg.validation_file=valid.json
+    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"
+        ),
+        tokenizer=tokenizer,
+    )
diff --git a/docs/source/datasets/nlp/token_classification_data.rst b/docs/source/datasets/nlp/token_classification_data.rst
index 20b48536..5835334a 100644
--- a/docs/source/datasets/nlp/token_classification_data.rst
+++ b/docs/source/datasets/nlp/token_classification_data.rst
@@ -12,4 +12,21 @@ To use custom text files, the files should contain new line delimited json objec
 
 .. code-block:: python
 
-    python train.py task=nlp/token_classification dataset.cfg.train_file=train.json dataset.cfg.validation_file=valid.json
+    from lightning_transformers.task.nlp.token_classification import (
+        TokenClassificationDataConfig,
+        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"
+        ),
+        tokenizer=tokenizer,
+    )
diff --git a/docs/source/datasets/nlp/translation_data.rst b/docs/source/datasets/nlp/translation_data.rst
index a63abbe7..a8e51a9a 100644
--- a/docs/source/datasets/nlp/translation_data.rst
+++ b/docs/source/datasets/nlp/translation_data.rst
@@ -14,4 +14,22 @@ We override the dataset files, allowing us to still use the data transforms defi
 
 .. code-block:: python
 
-    python train.py task=nlp/translation dataset.cfg.train_file=train.json dataset.cfg.validation_file=valid.json
+    from lightning_transformers.task.nlp.translation import (
+        TranslationDataConfig,
+        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"
+        ),
+        tokenizer=tokenizer,
+    )
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 5fd926df..06bf4606 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -14,7 +14,6 @@ Lightning Transformers
    :caption: Get started
 
    quickstart
-   structure/conf
 
 .. toctree::
    :maxdepth: 1
@@ -32,11 +31,8 @@ Lightning Transformers
 .. toctree::
    :maxdepth: 1
    :name: optimization
-   :caption: Training Optimizations
+   :caption: Transformer Optimizations
 
-   optimizations/lightning
-   optimizations/deepspeed
-   optimizations/sharded
    optimizations/sparseml
 
 .. toctree::
diff --git a/docs/source/optimizations/deepspeed.rst b/docs/source/optimizations/deepspeed.rst
deleted file mode 100644
index 56a3204a..00000000
--- a/docs/source/optimizations/deepspeed.rst
+++ /dev/null
@@ -1,39 +0,0 @@
-.. _deepspeed:
-
-DeepSpeed
-=========
-
-`DeepSpeed <https://github.com/microsoft/DeepSpeed>`__ is a deep learning training optimization library, providing the means to train massive billion parameter models at scale.
-This allows us to train large transformer models optimizing for compute. For more details, see `the DeepSpeed PyTorch Lightning docs <https://pytorch-lightning.readthedocs.io/en/latest/advanced/multi_gpu.html#deepspeed>`__.
-
-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>`_.
-
-DeepSpeed ZeRO Stage 2
-----------------------
-
-We provide out of the box configs to use the DeepSpeed plugin. Below is an example of how you can swap to the default trainer config for DeepSpeed when using the translation task.
-
-.. code-block:: python
-
-   python train.py task=nlp/translation dataset=nlp/translation/wmt16 trainer=deepspeed
-
-All options can be found in `conf/trainer/plugins/deepspeed.yaml <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/trainer/plugins/deepspeed.yaml>`__. We suggest referring to `the DeepSpeed PyTorch Lightning docs <https://pytorch-lightning.readthedocs.io/en/latest/advanced/multi_gpu.html#deepspeed>`__ for more information on the parameters.
-
-DeepSpeed ZeRO Stage 2 Offload
-------------------------------
-
-ZeRO-Offload enables large model training by being efficient with memory usage. ZeRO-Offload leverages the host CPU to execute the optimizer.
-
-We provide default trainer configurations to enable ZeRO-Offload:
-
-.. code-block:: python
-
-   python train.py task=nlp/translation dataset=nlp/translation/wmt16 trainer=deepspeed_offload
-
-To see the configuration settings see `conf/trainer/plugins/deepspeed_offload.yaml <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/trainer/plugins/deepspeed_offload.yaml>`__  for the parameters. Within the config file, you can modify the ZeRO parameters as described in `the DeepSpeed PyTorch Lightning docs <https://pytorch-lightning.readthedocs.io/en/latest/advanced/multi_gpu.html#deepspeed>`__.
-
-.. code-block:: yaml
-
-    _target_: pytorch_lightning.plugins.DeepSpeedPlugin
-    stage: 2
-    cpu_offload: True
diff --git a/docs/source/optimizations/lightning.rst b/docs/source/optimizations/lightning.rst
deleted file mode 100644
index d86c57b1..00000000
--- a/docs/source/optimizations/lightning.rst
+++ /dev/null
@@ -1,35 +0,0 @@
-.. _lightning:
-
-Lightning Trainer
-=================
-
-The Lightning Trainer offers many advanced features and training optimization without changing your models.
-
-Any Trainer flag can be used in Lightning transformers by passing ``trainer.FLAG=VALUE`` to ``train.py``.
-
-Max epochs
-""""""""""
-
-.. code-block:: python
-
-   python train.py task=nlp/text_classification dataset=nlp/text_classification/emotion trainer.max_epochs=4
-
-16-bit precision
-""""""""""""""""
-.. code-block:: python
-
-   python train.py task=nlp/text_classification dataset=nlp/text_classification/emotion trainer.precision=16
-
-Multi GPU
-"""""""""
-.. code-block:: python
-
-    python train.py task=nlp/translation dataset=nlp/translation/wmt16 trainer.gpus=4
-
-TPUs
-""""
-.. code-block:: python
-
-    python train.py task=nlp/translation dataset=nlp/translation/wmt16 trainer.tpu_cores=8
-
-See the `Pytorch Lightning Trainer <https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html>`_  or `conf/trainer/default <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/trainer/default.yaml>`_ for all the Lightning Trainer options.
diff --git a/docs/source/optimizations/sharded.rst b/docs/source/optimizations/sharded.rst
deleted file mode 100644
index 41f47d32..00000000
--- a/docs/source/optimizations/sharded.rst
+++ /dev/null
@@ -1,23 +0,0 @@
-.. _fairscale:
-
-Sharded Training
-================
-
-Sharded Training, provided by FairScale and Lightning shards optimizer states and gradients across multiple GPUs. This means the memory overhead per GPU is lower, as each GPU only has to maintain a partition of your optimizer state and gradients.
-
-This is particularly useful when training larger models, and reducing the memory requirement so that you can train with larger batch sizes or fit larger models into memory.
-For more information see `the Sharded Training PyTorch Lightning docs <https://pytorch-lightning.readthedocs.io/en/latest/advanced/multi_gpu.html#sharded-training>`_.
-
-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>`_.
-
-Install FairScale
-
-.. code-block:: bash
-
-    pip install fairscale
-
-Here we enable sharded training for the translation task.
-
-.. code-block:: bash
-
-    python train.py task=nlp/translation dataset=nlp/translation/wmt16 trainer=sharded
diff --git a/docs/source/quickstart.rst b/docs/source/quickstart.rst
index f9a39bf8..f09ee651 100644
--- a/docs/source/quickstart.rst
+++ b/docs/source/quickstart.rst
@@ -12,105 +12,118 @@ Using Lightning-Transformers
 
 Lightning Transformers has a collections of tasks for common NLP problems such as :ref:`language_modeling`, :ref:`translation` and more. To use, simply:
 
-1. Pick a task to train (passed to ``train.py`` as ``task=``)
+1. Pick a task to train (`LightningModule`)
 
-2. Pick a dataset (passed to ``train.py`` as ``dataset=``)
+2. Pick a dataset (`LightningDataModule`)
 
-3. Customize the backbone, optimizer, or any component within the config
+3. Use any `PyTorch Lightning parameters and optimizations <https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html>`_
 
-4. Add any `Lightning supported parameters and optimizations <https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html>`_
+Here is an example of training `bert-base-cased <https://huggingface.co/bert-base-cased>`__ on the `CARER <https://huggingface.co/datasets/emotion>`__ emotion dataset using the Text Classification task.
 
 .. code-block:: python
 
-   python train.py \
-        task=<TASK> \
-        dataset=<DATASET> \
-        backbone.pretrained_model_name_or_path=<BACKBONE> # Optionally change the HF backbone
-        optimizer=<OPTIMIZER> # Optionally specify optimizer (Default AdamW)
-        trainer.<ANY_TRAINER_FLAGS> # Optionally specify Lightning trainer arguments
+    import pytorch_lightning as pl
+    from transformers import AutoTokenizer
 
-Finetuning
-----------
+    from lightning_transformers.task.nlp.text_classification import (
+        TextClassificationDataModule,
+        TextClassificationTransformer,
+        TextClassificationDataConfig,
+    )
 
-In this example we will finetune the :ref:`text_classification` task on the `emotion <https://huggingface.co/datasets/emotion>`_ dataset.
+    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,
+        ),
+        tokenizer=tokenizer,
+    )
+    model = TextClassificationTransformer(pretrained_model_name_or_path="bert-base-cased")
 
-To fine-tune using the Text Classification default `bert-based-cased <https://huggingface.co/bert-base-cased>`_ model, use the following:
+    trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
 
-.. code-block:: python
+    trainer.fit(model, dm)
 
-   python train.py task=nlp/text_classification dataset=nlp/text_classification/emotion
 
-Change the Backbone
--------------------
+Changing the Optimizer
+----------------------
 
-Swap to the `RoBERTa <https://huggingface.co/roberta-base>`_ model:
+Swapping to the RMSProp optimizer:
 
 .. code-block:: python
 
-   python train.py task=nlp/text_classification dataset=nlp/text_classification/emotion backbone.pretrained_model_name_or_path=roberta-base
-
-Change the Optimizer
---------------------
-
-Swap to using RMSProp optimizer (see `conf/optimizer/ <https://github.com/PyTorchLightning/lightning-transformers/tree/master/conf/optimizer>`_ for all supported optimizers):
+    import pytorch_lightning as pl
+    import torch
+    import transformers
+    from transformers import AutoTokenizer
 
-.. code-block:: python
-
-   python train.py task=nlp/text_classification dataset=nlp/text_classification/emotion optimizer=rmsprop
+    from lightning_transformers.task.nlp.text_classification import (
+        TextClassificationDataModule,
+        TextClassificationTransformer,
+        TextClassificationDataConfig,
+    )
 
-For more info on how to override configuration, see :ref:`conf`.
 
-Lightning Trainer Options
---------------------------
+    class RMSPropTransformer(TextClassificationTransformer):
+        def configure_optimizers(self):
+            optimizer = torch.optim.AdamW(self.parameters(), lr=1e-5)
+            # automatically find the total number of steps we need!
+            num_training_steps, num_warmup_steps = self.compute_warmup(self.num_training_steps, num_warmup_steps=0.1)
+            scheduler = transformers.get_linear_schedule_with_warmup(
+                self.optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps
+            )
+            return {
+                "optimizer": optimizer,
+                "lr_scheduler": {"scheduler": scheduler, "interval": "step", "frequency": 1},
+            }
 
-We expose all `Pytorch Lightning Trainer <https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html>`_ parameters via the ``trainer`` config. This makes it easy to configure the Lightning Trainer without touching the code.
 
-Below we enable `Pytorch Lightning Native 16bit precision <https://pytorch-lightning.readthedocs.io/en/latest/amp.html#gpu-16-bit>`_ by setting the parameter via the CLI:
+    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,
+        ),
+        tokenizer=tokenizer,
+    )
+    model = RMSPropTransformer(pretrained_model_name_or_path="bert-base-cased")
 
-.. code-block:: python
+    trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
 
-   python train.py task=nlp/text_classification dataset=nlp/text_classification/emotion trainer.precision=16
+    trainer.fit(model, dm)
 
 
-Setting the maximum epochs:
+Enabling DeepSpeed/Sharded/Mixed Precision and more is super simple through the Lightning Trainer.
 
 .. code-block:: python
 
-    python train.py task=nlp/translation dataset=nlp/translation/wmt16 trainer.max_epochs=4
-
-Using multiple GPUs:
-
-.. code-block:: python
-
-    python train.py task=nlp/translation dataset=nlp/translation/wmt16 trainer.gpus=4
-
-Using TPUs:
+    # enable DeepSpeed with 16bit precision
+    trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1, strategy='deepspeed', precision=16)
 
-.. code-block:: python
-
-    python train.py task=nlp/translation dataset=nlp/translation/wmt16 trainer.tpu_cores=8
+    # enable DeepSpeed ZeRO Stage 3 with BFLOAT16 precision
+    trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1, strategy='deepspeed_stage_3_offload', precision="bf16")
 
-See the `Pytorch Lightning Trainer <https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html>`_  or `conf/trainer/default <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/trainer/default.yaml>`_ for all parameters.
+    # enable Sharded Training with 16bit precision
+    trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1, strategy='ddp_sharded', precision=16)
 
 Inference
 ---------
 
-Run inference once model trained (experimental):
-
 .. code-block:: python
 
-   python predict.py task=nlp/text_classification +checkpoint_path=/path/to/model.ckpt +x="Classify this sentence."
+    from transformers import AutoTokenizer
+    from lightning_transformers.task.nlp.text_classification import TextClassificationTransformer
 
+    model = TextClassificationTransformer(
+        pretrained_model_name_or_path="bert-base-uncased",
+        tokenizer=AutoTokenizer.from_pretrained(pretrained_model_name_or_path="bert-base-uncased"),
+    )
+    model.hf_predict("Lightning rocks!")
    # Returns [{'label': 'LABEL_0', 'score': 0.545...}]
-
-You can also run prediction using the default HuggingFace pre-trained model:
-
-.. code-block:: python
-
-   python predict.py task=nlp/text_classification +x="Classify this sentence."
-
-Or run prediction on a specified HuggingFace pre-trained model:
-
-.. code-block:: python
-
-   python predict.py task=nlp/text_classification backbone.pretrained_model_name_or_path=bert-base-cased +x="Classify this sentence."
diff --git a/docs/source/structure/conf.rst b/docs/source/structure/conf.rst
deleted file mode 100644
index 1242d11e..00000000
--- a/docs/source/structure/conf.rst
+++ /dev/null
@@ -1,85 +0,0 @@
-.. _conf:
-
-The Config Structure
-====================
-
-Lightning Transformers relies on Hydra config composition, meaning it is very easy to swap/modify components for training without having to touch the code.
-
-In practice, this means you can easily experiment and try different training procedures, parameters, optimizers/schedulers extremely quickly.
-
-Below is a diagram of the config structure, and how it fits together through the command line interface.
-
-.. code-block:: bash
-
-    conf/
-    ┣ backbone/  # Configs defining the backbone of the model/pre-trained model if any
-    ┣ dataset/ # Configs defining datasets
-    ┣ optimizer/ # Configs for optimizers
-    ┣ scheduler/ # Configs for schedulers
-    ┣ task/ # Configs defining the task, and any task specific parameters
-    ┣ tokenizer/ # Configs defining tokenizers, if any.
-    ┣ trainer/ # Configs defining PyTorch Lightning Trainers, with different configurations
-    ┣ training/ # Configs defining training specific parameters, such as batch size.
-    ┗ config.yaml # The main entrypoint containing all our chosen config components
-
-
-Hydra requires understanding of some syntactic sugar to use different configurations, which we explain below. For more details, we suggest looking at the `Hydra Documentation <https://hydra.cc/docs/next/advanced/override_grammar/basic>`_.
-
-By default ``train.py`` loads all the defined defaults in the `conf/config.yaml <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/config.yaml>`__ file in the `conf/ <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/>`__ directory.
-Task and Dataset have been marked as required within the ``config.yaml``, thus have to be specified on the command line.
-
-.. code-block:: bash
-
-    python train.py task=nlp/text_classification dataset=nlp/text_classification/emotion
-
-
-``task`` and ``dataset`` allow us to append configs from the conf folder to our training config. These config files can be found in `task/nlp/text_classification  <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/task/nlp/text_classification.yaml>`__ and `dataset/nlp/text_classification/emotion  <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/dataset/nlp/text_classification/emotion.yaml>`__ respectively.
-
-Overriding Configs
-******************
-All default configs can be overridden with custom configs, or pre-built configs found in the `conf/ <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/>`__ directory.
-
-For example, we can swap to different optimizers like below.
-
-.. code-block:: python
-
-    python train.py optimizer=sgd
-
-
-We can also modify the parameters in the config object from the commandline.
-
-.. code-block:: python
-
-    python train.py optimizer=sgd optimizer.momentum=0.99
-
-
-We use Hydras' `instantiation <https://hydra.cc/docs/next/patterns/instantiate_objects/overview>`_ to abstract the optimizer from the code, meaning it is very simple to plug in your own custom optimizers or schedulers via configs rather than code.
-
-.. code-block:: yaml
-
-    # Contents of conf/optimizer/sgd
-
-    _target_: torch.optim.SGD # _target_ class we'd like to instantiate our optimizer
-    lr: ${training.lr}
-    momentum: 0.9
-    weight_decay: 0.0005
-
-
-Config Inheritance
-******************
-
-Just like code, configs inherit from other configs.
-
-For example when defining a task, there is a default config object at `task/nlp/default.yaml <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/task/nlp/default.yaml>`__ that contains a set of default configurations for all tasks.
-This is useful, as our task configs do not need to define these parameters and instead can import them, like below.
-
-.. code-block:: yaml
-
-    # Contents of conf/task/nlp/language_modeling
-
-    defaults:
-      - nlp/default # Import our default configurations
-    _target_: lightning_transformers.task.nlp.language_modeling.LanguageModelingTransformer
-    downstream_model_type: transformers.AutoModelForCausalLM
-
-These are just a subset of Hydras' features. We offer full Hydra support to provide a robust and flexible API, thus we suggest reading the `Hydra tutorials <https://hydra.cc/docs/next/tutorials/intro>`_ for advanced use cases.
diff --git a/docs/source/tasks/nlp/language_modeling.rst b/docs/source/tasks/nlp/language_modeling.rst
index 450720c5..c2076341 100644
--- a/docs/source/tasks/nlp/language_modeling.rst
+++ b/docs/source/tasks/nlp/language_modeling.rst
@@ -24,37 +24,46 @@ Language Models pre-trained or fine-tuned to the Causal Language Modeling task c
 
 .. code-block:: python
 
-    python train.py task=nlp/language_modeling dataset=nlp/language_modeling/wikitext
+    import pytorch_lightning as pl
+    from transformers import AutoTokenizer
 
-Swap to GPT backbone:
+    from lightning_transformers.task.nlp.language_modeling import (
+        LanguageModelingDataConfig,
+        LanguageModelingDataModule,
+        LanguageModelingTransformer,
+    )
 
-.. code-block:: python
+    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",
+        ),
+        tokenizer=tokenizer,
+    )
+    trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
+
+    trainer.fit(model, dm)
 
-    python train.py task=nlp/language_modeling dataset=nlp/language_modeling/wikitext backbone.pretrained_model_name_or_path=gpt2
 
-We report the Cross Entropy Loss for validation. Find all options available for the task `here <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/task/nlp/language_modeling.yaml>`_.
+We report the Cross Entropy Loss for validation.
 
 .. include:: /datasets/nlp/language_modeling_data.rst
 
-Language Modeling Inference Pipeline (experimental)
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Language Modeling Inference Pipeline
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 By default we use the text generation pipeline, which requires a conditional input string and generates an output string.
 
-For Hydra to correctly parse your input argument, if your input contains any special characters you must either wrap the entire call in single quotes like `'+x="my, sentence"'` or escape special characters. See `Escaped characters in unquoted values <https://hydra.cc/docs/advanced/override_grammar/basic/#escaped-characters-in-unquoted-values>`_.
-
-.. code-block:: python
-
-    python predict.py task=nlp/language_modeling +checkpoint_path=/path/to/model.ckpt +x="Condition sentence for the language model"
-
-You can also run prediction using a default HuggingFace pre-trained model:
-
 .. code-block:: python
 
-   python predict.py task=nlp/language_modeling +x="Condition sentence for the language model"
-
-Or run prediction on a specified HuggingFace pre-trained model:
-
-.. code-block:: python
+    from transformers import AutoTokenizer
+    from lightning_transformers.task.nlp.language_modeling import LanguageModelingTransformer
 
-   python predict.py task=nlp/language_modeling backbone.pretrained_model_name_or_path=bert-base-cased +x="Condition sentence for the language model"
+    model = LanguageModelingTransformer(
+            pretrained_model_name_or_path="prajjwal1/bert-tiny",
+            tokenizer=AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny"),
+        )
+    model.hf_predict("The house:")
diff --git a/docs/source/tasks/nlp/multiple_choice.rst b/docs/source/tasks/nlp/multiple_choice.rst
index 588d1df3..a1c34d86 100644
--- a/docs/source/tasks/nlp/multiple_choice.rst
+++ b/docs/source/tasks/nlp/multiple_choice.rst
@@ -30,15 +30,31 @@ Training
 
 .. code-block:: python
 
-    python train.py task=nlp/multiple_choice dataset=nlp/multiple_choice/race # can use the swag dataset instead
-
-Swap to GPT backbone:
-
-.. code-block:: python
-
-    python train.py task=nlp/multiple_choice dataset=nlp/multiple_choice/race backbone.pretrained_model_name_or_path=gpt2
-
-We report Cross Entropy Loss, Precision, Recall and Accuracy for validation. Find all options available for the task `here <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/task/nlp/multiple_choice.yaml>`_.
+    import pytorch_lightning as pl
+    from transformers import AutoTokenizer
+
+    from lightning_transformers.task.nlp.multiple_choice import (
+        MultipleChoiceDataConfig,
+        MultipleChoiceTransformer,
+        SwagMultipleChoiceDataModule,
+    )
+
+    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,
+        ),
+        tokenizer=tokenizer,
+    )
+    trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
+
+    trainer.fit(model, dm)
+
+We report Cross Entropy Loss, Precision, Recall and Accuracy for validation.
 
 .. include:: /datasets/nlp/multiple_choice_data.rst
 
diff --git a/docs/source/tasks/nlp/question_answering.rst b/docs/source/tasks/nlp/question_answering.rst
index a79ebb23..f47413f4 100644
--- a/docs/source/tasks/nlp/question_answering.rst
+++ b/docs/source/tasks/nlp/question_answering.rst
@@ -27,35 +27,49 @@ Training
 
 .. code-block:: python
 
-    python train.py task=nlp/question_answering dataset=nlp/question_answering/squad
-
-Swap to GPT backbone:
-
-.. code-block:: python
-
-    python train.py task=nlp/question_answering dataset=nlp/question_answering/squad backbone.pretrained_model_name_or_path=gpt2
+    import pytorch_lightning as pl
+    from transformers import AutoTokenizer
+
+    from lightning_transformers.task.nlp.question_answering import (
+        QuestionAnsweringDataConfig,
+        QuestionAnsweringTransformer,
+        SquadDataModule,
+    )
+
+    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,
+        ),
+        tokenizer=tokenizer,
+    )
+    trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
+
+    trainer.fit(model, dm)
 
 .. include:: /datasets/nlp/question_answering_data.rst
 
-Question Answering Inference Pipeline (experimental)
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Question Answering Inference Pipeline
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 By default we use the question answering pipeline, which requires a context and a question as input.
 
-For Hydra to correctly parse your input argument, if your input contains any special characters you must either wrap the entire call in single quotes like `'+x="my, sentence"'` or escape special characters. See `escaped characters in unquoted values <https://hydra.cc/docs/advanced/override_grammar/basic/#escaped-characters-in-unquoted-values>`_.
-
-.. code-block:: python
-
-    python predict.py task=nlp/question_answering +checkpoint_path=/path/to/model.ckpt +x='{context: "The ground is black, the sky is blue and the car is red.", question: "What color is the sky?"}'
-
-You can also run prediction using a default HuggingFace pre-trained model:
-
 .. code-block:: python
 
-   python predict.py task=nlp/question_answering +x='{context: "The ground is black, the sky is blue and the car is red.", question: "What color is the sky?"}'
-
-Or run prediction on a specified HuggingFace pre-trained model:
-
-.. code-block:: python
+    from transformers import AutoTokenizer
+    from lightning_transformers.task.nlp.question_answering import QuestionAnsweringTransformer
 
-   python predict.py task=nlp/question_answering backbone.pretrained_model_name_or_path=bert-base-cased +x='{context: "The ground is black, the sky is blue and the car is red.", question: "What color is the sky?"}'
+    model = QuestionAnsweringTransformer(
+        pretrained_model_name_or_path="sshleifer/tiny-gpt2",
+        tokenizer=AutoTokenizer.from_pretrained(pretrained_model_name_or_path="sshleifer/tiny-gpt2"),
+    )
+    model.hf_predict(dict(context="Lightning is great", question="What is great?"))
diff --git a/docs/source/tasks/nlp/summarization.rst b/docs/source/tasks/nlp/summarization.rst
index a50a9707..e39b2027 100644
--- a/docs/source/tasks/nlp/summarization.rst
+++ b/docs/source/tasks/nlp/summarization.rst
@@ -22,33 +22,63 @@ In the CNN/Daily Mail dataset, this involves taking long articles and summarizin
 Training
 ^^^^^^^^
 
-To use this task, we must select a Seq2Seq Encoder/Decoder based model, such as `T5 <https://huggingface.co/transformers/model_doc/t5.html>`_ or `BART <https://huggingface.co/transformers/model_doc/bart.html>`_. Conventional models like `GPT <https://huggingface.co/transformers/model_doc/gpt.html>`_/`BERT <https://huggingface.co/transformers/model_doc/bert.html>`_ will not work as they are encoder only.
+To use this task, we must select a Seq2Seq Encoder/Decoder based model, such as `T5 <https://huggingface.co/transformers/model_doc/t5.html>`_ or `BART <https://huggingface.co/transformers/model_doc/bart.html>`_. Encoder only models like `GPT <https://huggingface.co/transformers/model_doc/gpt.html>`_/`BERT <https://huggingface.co/transformers/model_doc/bert.html>`_ will not work as they are encoder only.
 
 .. code-block:: python
 
-    python train.py task=nlp/summarization dataset=nlp/summarization/cnn_dailymail backbone.pretrained_model_name_or_path=t5-base # dataset can be swapped to xsum
+    import pytorch_lightning as pl
+    from transformers import AutoTokenizer
+
+    from lightning_transformers.task.nlp.summarization import (
+        SummarizationConfig,
+        SummarizationDataConfig,
+        SummarizationTransformer,
+        XsumSummarizationDataModule,
+    )
+
+    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,
+        ),
+    )
+    dm = XsumSummarizationDataModule(
+        cfg=SummarizationDataConfig(
+            batch_size=1,
+            dataset_name="xsum",
+            max_source_length=128,
+            max_target_length=128,
+        ),
+        tokenizer=tokenizer,
+    )
+    trainer = pl.Trainer(accelerator="auto", devices=1, max_epochs=1)
+
+    trainer.fit(model, dm)
+
 
 .. include:: /datasets/nlp/summarization_data.rst
 
-Summarization Inference Pipeline (experimental)
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Summarization Inference Pipeline
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 By default we use the summarization pipeline, which requires an input document as text.
 
-For Hydra to correctly parse your input argument, if your input contains any special characters you must either wrap the entire call in single quotes like `'+x="my, sentence"'` or escape special characters. See `escaped characters in unquoted values <https://hydra.cc/docs/advanced/override_grammar/basic/#escaped-characters-in-unquoted-values>`_.
-
-.. code-block:: bash
-
-    python predict.py task=nlp/summarization +checkpoint_path=/path/to/model.ckpt '+x="The car was racing towards the tunnel, whilst blue lights were flashing behind it. The car entered the tunnel and vanished..."'
-
-You can also run prediction using a default HuggingFace pre-trained model:
-
-.. code-block:: bash
-
-   python predict.py task=nlp/summarization '+x="The car was racing towards the tunnel, whilst blue lights were flashing behind it. The car entered the tunnel and vanished..."'
+.. code-block:: python
 
-Or run prediction on a specified HuggingFace pre-trained model:
+    from transformers import AutoTokenizer
+    from lightning_transformers.task.nlp.summarization import SummarizationTransformer
 
-.. code-block:: bash
+    model = SummarizationTransformer(
+        pretrained_model_name_or_path="patrickvonplaten/t5-tiny-random",
+        tokenizer=AutoTokenizer.from_pretrained(pretrained_model_name_or_path="patrickvonplaten/t5-tiny-random"),
+    )
 
-   python predict.py task=nlp/summarization backbone.pretrained_model_name_or_path=t5-base '+x="The car was racing towards the tunnel, whilst blue lights were flashing behind it. The car entered the tunnel and vanished..."'
+    model.hf_predict(
+        "The results found significant improvements over all tasks evaluated",
+        min_length=2,
+        max_length=12,
+    )
diff --git a/docs/source/tasks/nlp/text_classification.rst b/docs/source/tasks/nlp/text_classification.rst
index 1275563a..68c13174 100644
--- a/docs/source/tasks/nlp/text_classification.rst
+++ b/docs/source/tasks/nlp/text_classification.rst
@@ -24,37 +24,46 @@ For this task, you can rely on most Transformer models as your backbone.
 
 .. code-block:: python
 
-    python train.py task=nlp/text_classification dataset=nlp/text_classification/emotion # can be swapped to xlni or glue
-
-Swap to GPT backbone:
-
-.. code-block:: python
-
-    python train.py task=nlp/text_classification dataset=nlp/text_classification/emotion backbone.pretrained_model_name_or_path=gpt2
-
-We report the Precision, Recall, Accuracy and Cross Entropy Loss for validation. Find all options available for the task `here <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/task/nlp/text_classification.yaml>`_.
+    import pytorch_lightning as pl
+    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,
+        ),
+        tokenizer=tokenizer,
+    )
+    model = TextClassificationTransformer(pretrained_model_name_or_path="bert-base-uncased", num_labels=dm.num_classes)
+    trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
+
+    trainer.fit(model, dm)
+
+We report the Precision, Recall, Accuracy and Cross Entropy Loss for validation.
 
 .. include:: /datasets/nlp/text_classification_data.rst
 
-Text Classification Inference Pipeline (experimental)
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Text Classification Inference Pipeline
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 By default we use the sentiment-analysis pipeline, which requires an input string.
 
-For Hydra to correctly parse your input argument, if your input contains any special characters you must either wrap the entire call in single quotes like `'+x="my, sentence"'` or escape special characters. See `escaped characters in unquoted values <https://hydra.cc/docs/advanced/override_grammar/basic/#escaped-characters-in-unquoted-values>`_.
-
 .. code-block:: python
 
-    python predict.py task=nlp/text_classification +checkpoint_path=/path/to/model.ckpt '+x="I dont like this at all!"'
-
-You can also run prediction using a default HuggingFace pre-trained model:
-
-.. code-block:: python
-
-   python predict.py task=nlp/text_classification '+x="I dont like this at all!"'
-
-Or run prediction on a specified HuggingFace pre-trained model:
-
-.. code-block:: python
+    from transformers import AutoTokenizer
+    from lightning_transformers.task.nlp.text_classification import TextClassificationTransformer
 
-   python predict.py task=nlp/text_classification backbone.pretrained_model_name_or_path=bert-base-cased '+x="I dont like this at all!"'
+    model = TextClassificationTransformer(
+        pretrained_model_name_or_path="prajjwal1/bert-tiny",
+        tokenizer=AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny"),
+    )
+    model.hf_predict("Lightning rocks!")
diff --git a/docs/source/tasks/nlp/token_classification.rst b/docs/source/tasks/nlp/token_classification.rst
index 8b897ecf..9086aa3e 100644
--- a/docs/source/tasks/nlp/token_classification.rst
+++ b/docs/source/tasks/nlp/token_classification.rst
@@ -17,37 +17,49 @@ Training
 
 .. code-block:: python
 
-    python train.py task=nlp/token_classification dataset=nlp/token_classification/conll
-
-Swap to GPT backbone:
-
-.. code-block:: python
-
-    python train.py task=nlp/token_classification dataset=nlp/token_classification/conll backbone.pretrained_model_name_or_path=gpt2
-
-We report the Precision, Recall, Accuracy and Cross Entropy Loss for validation. To see all options available for the task, see `Find all options available for the task `here <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/task/nlp/token_classification.yaml>`_.
+    import pytorch_lightning as pl
+    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",
+        ),
+        tokenizer=tokenizer,
+    )
+    model = TokenClassificationTransformer(pretrained_model_name_or_path="bert-base-uncased", labels=dm.labels)
+    trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
+
+    trainer.fit(model, dm)
+
+We report the Precision, Recall, Accuracy and Cross Entropy Loss for validation.
 
 .. include:: /datasets/nlp/token_classification_data.rst
 
-Token Classification Inference Pipeline (experimental)
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Token Classification Inference Pipeline
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 By default we use the NER pipeline, which requires a an input sequence string and the number of labels.
 
-For Hydra to correctly parse your input argument, if your input contains any special characters you must either wrap the entire call in single quotes like `'+x="my, sentence"'` or escape special characters. See `escaped characters in unquoted values <https://hydra.cc/docs/advanced/override_grammar/basic/#escaped-characters-in-unquoted-values>`_.
-
 .. code-block:: python
 
-    python predict.py task=nlp/token_classification +checkpoint_path=/path/to/model.ckpt +x="London is the capital of the United Kingdom." +model_data_kwargs='{labels: 2}'
-
-You can also run prediction using a default HuggingFace pre-trained model:
-
-.. code-block:: python
-
-   python predict.py task=nlp/token_classification +x="London is the capital of the United Kingdom." +model_data_kwargs='{labels: 2}'
-
-Or run prediction on a specified HuggingFace pre-trained model:
-
-.. code-block:: python
+    from transformers import AutoTokenizer
+    from lightning_transformers.task.nlp.token_classification import TokenClassificationTransformer
 
-   python predict.py task=nlp/token_classification backbone.pretrained_model_name_or_path=bert-base-cased +x="London is the capital of the United Kingdom." +model_data_kwargs='{labels: 2}'
+    model = TokenClassificationTransformer(
+        pretrained_model_name_or_path="prajjwal1/bert-tiny",
+        tokenizer=AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny"),
+        labels=2,
+    )
+    model.hf_predict("Have a good day!")
diff --git a/docs/source/tasks/nlp/translation.rst b/docs/source/tasks/nlp/translation.rst
index 33d49196..710d9ea4 100644
--- a/docs/source/tasks/nlp/translation.rst
+++ b/docs/source/tasks/nlp/translation.rst
@@ -25,18 +25,47 @@ In addition, you also need a tokenizer that has been created on multi-lingual te
 
 .. code-block:: python
 
-    python train.py task=nlp/translation dataset=nlp/translation/wmt16 backbone.pretrained_model_name_or_path=google/mt5-base
-
-By default we fine-tune to translate from English to Romanian. This can be changed by specifying the source/target languages (see more in ``conf/dataset/nlp/translation/wmt16.yaml``):
-
-.. code-block:: python
-
-    python train.py task=nlp/translation dataset=nlp/translation/wmt16 backbone.pretrained_model_name_or_path=google/mt5-base dataset.cfg.dataset_config_name=de-en dataset.cfg.source_language=de dataset.cfg.target_language=en
+    import pytorch_lightning as pl
+    from transformers import AutoTokenizer
+
+    from lightning_transformers.task.nlp.translation import (
+        TranslationConfig,
+        TranslationDataConfig,
+        TranslationTransformer,
+        WMT16TranslationDataModule,
+    )
+
+    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,
+        ),
+    )
+    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,
+        ),
+        tokenizer=tokenizer,
+    )
+    trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
+
+    trainer.fit(model, dm)
 
 .. include:: /datasets/nlp/translation_data.rst
 
-Translation Inference Pipeline (experimental)
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Translation Inference Pipeline
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 By default we use the translation pipeline, which requires a source text string.
 
@@ -44,10 +73,11 @@ For Hydra to correctly parse your input argument, if your input contains any spe
 
 .. code-block:: python
 
-    python predict.py task=nlp/translation +checkpoint_path=/path/to/model.ckpt +x="The ground is black, the sky is blue and the car is red."
-
-You can also run prediction using a HuggingFace pre-trained model:
-
-.. code-block:: python
+    from transformers import AutoTokenizer
+    from lightning_transformers.task.nlp.translation import TranslationTransformer
 
-   python predict.py task=nlp/translation backbone.pretrained_model_name_or_path="Helsinki-NLP/opus-mt-en-de" '+x="The ground is black, the sky is blue and the car is red."'
+    model = TranslationTransformer(
+        pretrained_model_name_or_path="patrickvonplaten/t5-tiny-random",
+        tokenizer=AutoTokenizer.from_pretrained(pretrained_model_name_or_path="patrickvonplaten/t5-tiny-random"),
+    )
+    model.hf_predict("¡Hola Sean!")
diff --git a/examples/text_classification.py b/examples/text_classification.py
index 1f5015fd..05716d4c 100644
--- a/examples/text_classification.py
+++ b/examples/text_classification.py
@@ -18,7 +18,6 @@
         ),
         tokenizer=tokenizer,
     )
-    dm.setup("fit")
     model = TextClassificationTransformer(pretrained_model_name_or_path="bert-base-uncased", num_labels=dm.num_classes)
     trainer = pl.Trainer(accelerator="auto", devices="auto", max_epochs=1)
 
diff --git a/lightning_transformers/core/nlp/config.py b/lightning_transformers/core/nlp/config.py
index 7ba5f6ec..b874a24d 100644
--- a/lightning_transformers/core/nlp/config.py
+++ b/lightning_transformers/core/nlp/config.py
@@ -27,6 +27,9 @@ class HFTransformerDataConfig(TransformerDataConfig):
     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
 
 
 @dataclass
diff --git a/lightning_transformers/core/nlp/data.py b/lightning_transformers/core/nlp/data.py
index ec5ea390..f72b4063 100644
--- a/lightning_transformers/core/nlp/data.py
+++ b/lightning_transformers/core/nlp/data.py
@@ -72,7 +72,7 @@ def load_dataset(self) -> Dataset:
         # Use special subset names if provided, and rename them back to standard ones
         for subset in ("train", "validation", "test"):
             config_attr = f"{subset}_subset_name"
-            if hasattr(self.cfg, config_attr):
+            if getattr(self.cfg, config_attr) is not None:
                 special_subset_name = getattr(self.cfg, config_attr)
                 if special_subset_name not in dataset:
                     raise KeyError(
diff --git a/lightning_transformers/task/nlp/text_classification/data.py b/lightning_transformers/task/nlp/text_classification/data.py
index 0561ced4..a84ccd58 100644
--- a/lightning_transformers/task/nlp/text_classification/data.py
+++ b/lightning_transformers/task/nlp/text_classification/data.py
@@ -14,6 +14,7 @@
 from typing import Any, Dict, List, Optional
 
 from datasets import ClassLabel, Dataset
+from pytorch_lightning.utilities import rank_zero_warn
 from transformers import PreTrainedTokenizerBase
 
 from lightning_transformers.core.nlp import HFDataModule
@@ -27,6 +28,7 @@ class TextClassificationDataModule(HFDataModule):
 
     def __init__(self, *args, cfg: TextClassificationDataConfig = TextClassificationDataConfig(), **kwargs) -> None:
         super().__init__(*args, cfg=cfg, **kwargs)
+        self.labels = None
 
     def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset:
         input_feature_fields = [k for k, v in dataset["train"].features.items() if k not in ["label", "idx"]]
@@ -50,6 +52,9 @@ def process_data(self, dataset: Dataset, stage: Optional[str] = None) -> Dataset
 
     @property
     def num_classes(self) -> int:
+        if self.labels is None:
+            rank_zero_warn("Labels has not been set, calling `setup('fit')`.")
+            self.setup("fit")
         return self.labels.num_classes
 
     @property


Test patch:
diff --git a/tests/task/nlp/test_text_classification.py b/tests/task/nlp/test_text_classification.py
index d859447e..a0e64d0e 100644
--- a/tests/task/nlp/test_text_classification.py
+++ b/tests/task/nlp/test_text_classification.py
@@ -15,7 +15,6 @@
 @pytest.mark.skipif(sys.platform == "win32", reason="Currently Windows is not supported")
 def test_smoke_train(hf_cache_path):
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path="prajjwal1/bert-tiny")
-    model = TextClassificationTransformer(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     dm = TextClassificationDataModule(
         cfg=TextClassificationDataConfig(
             batch_size=1,
@@ -29,6 +28,7 @@ def test_smoke_train(hf_cache_path):
         ),
         tokenizer=tokenizer,
     )
+    model = TextClassificationTransformer(pretrained_model_name_or_path="prajjwal1/bert-tiny")
     trainer = pl.Trainer(fast_dev_run=True)
     trainer.fit(model, dm)