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? 
hey @ElderWanng super late response.

Could you have a look at the latest code? Things have changed drastically, moving towards a class based approach :) 

Fix patch:
diff --git a/conf/__init__.py b/conf/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/conf/backbone/nlp/default.yaml b/conf/backbone/nlp/default.yaml
deleted file mode 100644
index 4b57be4e..00000000
--- a/conf/backbone/nlp/default.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-# @package backbone
-defaults:
-  - /tokenizer@_group_: autotokenizer # use AutoTokenizer by default
-pretrained_model_name_or_path: bert-base-cased
diff --git a/conf/backbone/nlp/seq2seq.yaml b/conf/backbone/nlp/seq2seq.yaml
deleted file mode 100644
index 459cbb37..00000000
--- a/conf/backbone/nlp/seq2seq.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-# @package backbone
-defaults:
-  - /tokenizer@_group_: autotokenizer # use AutoTokenizer by default
-pretrained_model_name_or_path: t5-base
diff --git a/conf/config.yaml b/conf/config.yaml
deleted file mode 100644
index d86aacc9..00000000
--- a/conf/config.yaml
+++ /dev/null
@@ -1,10 +0,0 @@
-defaults: # loads default configs
-  - task: default
-  - optimizer: adamw
-  - scheduler: linear_schedule_with_warmup
-  - training: default
-  - trainer: default
-
-experiment_name: ${now:%Y-%m-%d}_${now:%H-%M-%S}
-log: False
-ignore_warnings: True # todo: check warnings before release
diff --git a/conf/dataset/default.yaml b/conf/dataset/default.yaml
deleted file mode 100644
index 6849389e..00000000
--- a/conf/dataset/default.yaml
+++ /dev/null
@@ -1,6 +0,0 @@
-# @package dataset
-_target_: lightning_transformers.core.data.TransformerDataModule
-cfg:
-  # torch data-loader specific arguments
-  batch_size: ${training.batch_size}
-  num_workers: ${training.num_workers}
diff --git a/conf/dataset/nlp/default.yaml b/conf/dataset/nlp/default.yaml
deleted file mode 100644
index ec208d7e..00000000
--- a/conf/dataset/nlp/default.yaml
+++ /dev/null
@@ -1,21 +0,0 @@
-# @package dataset
-defaults:
-  - /dataset/default
-_target_: lightning_transformers.core.nlp.HFDataModule
-cfg:
-  dataset_name: null
-  dataset_config_name: null
-  train_file: null
-  validation_file: null
-  test_file: null
-  train_val_split: null
-  max_samples: null
-  cache_dir: null
-  padding: 'max_length'
-  truncation: 'only_first'
-  preprocessing_num_workers: 1
-  load_from_cache_file: True
-  max_length: 128
-  limit_train_samples: null
-  limit_val_samples: null
-  limit_test_samples: null
diff --git a/conf/dataset/nlp/language_modeling/default.yaml b/conf/dataset/nlp/language_modeling/default.yaml
deleted file mode 100644
index 3582fd51..00000000
--- a/conf/dataset/nlp/language_modeling/default.yaml
+++ /dev/null
@@ -1,6 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.language_modeling.LanguageModelingDataModule
-cfg:
-  block_size: null
diff --git a/conf/dataset/nlp/language_modeling/wikitext.yaml b/conf/dataset/nlp/language_modeling/wikitext.yaml
deleted file mode 100644
index 8ea28d3e..00000000
--- a/conf/dataset/nlp/language_modeling/wikitext.yaml
+++ /dev/null
@@ -1,8 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.language_modeling.LanguageModelingDataModule
-cfg:
-  dataset_name: wikitext
-  dataset_config_name: wikitext-2-raw-v1
-  block_size: 512
diff --git a/conf/dataset/nlp/masked_language_modeling/default.yaml b/conf/dataset/nlp/masked_language_modeling/default.yaml
deleted file mode 100644
index 8f762bb0..00000000
--- a/conf/dataset/nlp/masked_language_modeling/default.yaml
+++ /dev/null
@@ -1,8 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.masked_language_modeling.MaskedLanguageModelingDataModule
-cfg:
-  mlm_probability: null
-  line_by_line: null
-  wwm: null
diff --git a/conf/dataset/nlp/masked_language_modeling/wikitext.yaml b/conf/dataset/nlp/masked_language_modeling/wikitext.yaml
deleted file mode 100644
index aa6b884d..00000000
--- a/conf/dataset/nlp/masked_language_modeling/wikitext.yaml
+++ /dev/null
@@ -1,10 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.masked_language_modeling.MaskedLanguageModelingDataModule
-cfg:
-  dataset_name: wikitext
-  dataset_config_name: wikitext-2-raw-v1
-  mlm_probability: 0.15
-  line_by_line: False
-  wwm: False
diff --git a/conf/dataset/nlp/multiple_choice/default.yaml b/conf/dataset/nlp/multiple_choice/default.yaml
deleted file mode 100644
index bdf6649d..00000000
--- a/conf/dataset/nlp/multiple_choice/default.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.multiple_choice.MultipleChoiceDataModule
diff --git a/conf/dataset/nlp/multiple_choice/race.yaml b/conf/dataset/nlp/multiple_choice/race.yaml
deleted file mode 100644
index e8213c48..00000000
--- a/conf/dataset/nlp/multiple_choice/race.yaml
+++ /dev/null
@@ -1,8 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.multiple_choice.RaceMultipleChoiceDataModule
-cfg:
-  dataset_name: race
-  dataset_config_name: 'all'
-  padding: False
diff --git a/conf/dataset/nlp/multiple_choice/swag.yaml b/conf/dataset/nlp/multiple_choice/swag.yaml
deleted file mode 100644
index aa00ff7e..00000000
--- a/conf/dataset/nlp/multiple_choice/swag.yaml
+++ /dev/null
@@ -1,8 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.multiple_choice.SwagMultipleChoiceDataModule
-cfg:
-  dataset_name: swag
-  dataset_config_name: 'regular'
-  padding: False
diff --git a/conf/dataset/nlp/question_answering/default.yaml b/conf/dataset/nlp/question_answering/default.yaml
deleted file mode 100644
index 1fbea775..00000000
--- a/conf/dataset/nlp/question_answering/default.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.question_answering.QuestionAnsweringDataModule
diff --git a/conf/dataset/nlp/question_answering/squad.yaml b/conf/dataset/nlp/question_answering/squad.yaml
deleted file mode 100644
index fc90fe63..00000000
--- a/conf/dataset/nlp/question_answering/squad.yaml
+++ /dev/null
@@ -1,14 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.question_answering.datasets.SquadDataModule
-cfg:
-  dataset_name: squad
-  dataset_config_name: 'plain_text'
-  max_length: 384
-  version_2_with_negative: False
-  null_score_diff_threshold: .0
-  doc_stride: 128
-  n_best_size: 20
-  max_answer_length: 30
-  output_dir: ${training.output_dir}
diff --git a/conf/dataset/nlp/summarization/cnn_dailymail.yaml b/conf/dataset/nlp/summarization/cnn_dailymail.yaml
deleted file mode 100644
index c123a678..00000000
--- a/conf/dataset/nlp/summarization/cnn_dailymail.yaml
+++ /dev/null
@@ -1,9 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.summarization.CNNDailyMailSummarizationDataModule
-cfg:
-  dataset_name: cnn_dailymail
-  dataset_config_name: 3.0.0
-  max_source_length: 1024
-  max_target_length: 128
diff --git a/conf/dataset/nlp/summarization/default.yaml b/conf/dataset/nlp/summarization/default.yaml
deleted file mode 100644
index 4f455aee..00000000
--- a/conf/dataset/nlp/summarization/default.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.summarization.SummarizationDataModule
diff --git a/conf/dataset/nlp/summarization/xsum.yaml b/conf/dataset/nlp/summarization/xsum.yaml
deleted file mode 100644
index 4a5bb7c0..00000000
--- a/conf/dataset/nlp/summarization/xsum.yaml
+++ /dev/null
@@ -1,8 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.summarization.XsumSummarizationDataModule
-cfg:
-  dataset_name: xsum
-  max_source_length: 128
-  max_target_length: 128
diff --git a/conf/dataset/nlp/text_classification/default.yaml b/conf/dataset/nlp/text_classification/default.yaml
deleted file mode 100644
index 4693c1af..00000000
--- a/conf/dataset/nlp/text_classification/default.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.text_classification.TextClassificationDataModule
diff --git a/conf/dataset/nlp/text_classification/emotion.yaml b/conf/dataset/nlp/text_classification/emotion.yaml
deleted file mode 100644
index 2100bf5e..00000000
--- a/conf/dataset/nlp/text_classification/emotion.yaml
+++ /dev/null
@@ -1,6 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.text_classification.TextClassificationDataModule
-cfg:
-  dataset_name: emotion
diff --git a/conf/dataset/nlp/text_classification/glue.yaml b/conf/dataset/nlp/text_classification/glue.yaml
deleted file mode 100644
index f621889f..00000000
--- a/conf/dataset/nlp/text_classification/glue.yaml
+++ /dev/null
@@ -1,8 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.text_classification.TextClassificationDataModule
-cfg:
-  dataset_name: glue
-  # ['cola', 'sst2', 'mrpc', 'qqp', 'stsb', 'mnli', 'mnli_mismatched', 'mnli_matched', 'qnli', 'rte', 'wnli', 'ax']
-  dataset_config_name: sst2
diff --git a/conf/dataset/nlp/text_classification/xnli.yaml b/conf/dataset/nlp/text_classification/xnli.yaml
deleted file mode 100644
index 8821296e..00000000
--- a/conf/dataset/nlp/text_classification/xnli.yaml
+++ /dev/null
@@ -1,8 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.text_classification.TextClassificationDataModule
-cfg:
-  dataset_name: xnli
-   # ['ar', 'bg', 'de', 'el', 'en', 'es', 'fr', 'hi', 'ru', 'sw', 'th', 'tr', 'ur', 'vi', 'zh', 'all_languages']
-  dataset_config_name: en
diff --git a/conf/dataset/nlp/token_classification/conll.yaml b/conf/dataset/nlp/token_classification/conll.yaml
deleted file mode 100644
index 03851819..00000000
--- a/conf/dataset/nlp/token_classification/conll.yaml
+++ /dev/null
@@ -1,10 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.token_classification.TokenClassificationDataModule
-cfg:
-  task_name: ner
-  dataset_name: conll2003
-  preprocessing_num_workers: 1
-  label_all_tokens: False
-  pad_to_max_length: False
diff --git a/conf/dataset/nlp/token_classification/default.yaml b/conf/dataset/nlp/token_classification/default.yaml
deleted file mode 100644
index a10397b8..00000000
--- a/conf/dataset/nlp/token_classification/default.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.token_classification.TokenClassificationDataModule
diff --git a/conf/dataset/nlp/translation/default.yaml b/conf/dataset/nlp/translation/default.yaml
deleted file mode 100644
index fb435d59..00000000
--- a/conf/dataset/nlp/translation/default.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.translation.TranslationDataModule
diff --git a/conf/dataset/nlp/translation/wmt16.yaml b/conf/dataset/nlp/translation/wmt16.yaml
deleted file mode 100644
index b22cfd2b..00000000
--- a/conf/dataset/nlp/translation/wmt16.yaml
+++ /dev/null
@@ -1,11 +0,0 @@
-# @package dataset
-defaults:
-  - nlp/default
-_target_: lightning_transformers.task.nlp.translation.WMT16TranslationDataModule
-cfg:
-  dataset_name: wmt16
-  dataset_config_name: 'ro-en' # WMT translation datasets: ['cs-en', 'de-en', 'fi-en', 'ro-en', 'ru-en', 'tr-en']
-  source_language: 'en'
-  target_language: 'ro'
-  max_source_length: 128
-  max_target_length: 128
diff --git a/conf/hydra/job_logging/custom.yaml b/conf/hydra/job_logging/custom.yaml
deleted file mode 100644
index 92cf00c4..00000000
--- a/conf/hydra/job_logging/custom.yaml
+++ /dev/null
@@ -1,20 +0,0 @@
-hydra:
-    job_logging:
-        formatters:
-            simple:
-                format: "%(message)s"
-        root:
-            handlers: [debug_console_handler, file_handler]
-        version: 1
-        handlers:
-            debug_console_handler:
-                level: DEBUG
-                formatter: simple
-                class: logging.StreamHandler
-                stream: ext://sys.stdout
-            file_handler:
-                level: DEBUG
-                formatter: simple
-                class: logging.FileHandler
-                filename: train.log
-        disable_existing_loggers: False
diff --git a/conf/hydra/output/custom.yaml b/conf/hydra/output/custom.yaml
deleted file mode 100644
index f7d90993..00000000
--- a/conf/hydra/output/custom.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-# @package _global_
-hydra:
-    run:
-        dir: ./outputs/${job_name}/${job_name}-${model_name}-${now:%Y%m%d_%H%M%S}
diff --git a/conf/optimizer/adadelta.yaml b/conf/optimizer/adadelta.yaml
deleted file mode 100644
index 809aecb5..00000000
--- a/conf/optimizer/adadelta.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-_target_: torch.optim.Adadelta
-lr: ${training.lr}
diff --git a/conf/optimizer/adagrad.yaml b/conf/optimizer/adagrad.yaml
deleted file mode 100644
index 200566f1..00000000
--- a/conf/optimizer/adagrad.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-_target_: torch.optim.Adagrad
-lr: ${training.lr}
diff --git a/conf/optimizer/adam.yaml b/conf/optimizer/adam.yaml
deleted file mode 100644
index fd87db0b..00000000
--- a/conf/optimizer/adam.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-_target_: torch.optim.Adam
-lr: ${training.lr}
diff --git a/conf/optimizer/adamax.yaml b/conf/optimizer/adamax.yaml
deleted file mode 100644
index 6af949d6..00000000
--- a/conf/optimizer/adamax.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-_target_: torch.optim.Adamax
-lr: ${training.lr}
diff --git a/conf/optimizer/adamw.yaml b/conf/optimizer/adamw.yaml
deleted file mode 100644
index 385f80de..00000000
--- a/conf/optimizer/adamw.yaml
+++ /dev/null
@@ -1,3 +0,0 @@
-_target_: torch.optim.AdamW
-lr: ${training.lr}
-weight_decay: 0.001
diff --git a/conf/optimizer/asgd.yaml b/conf/optimizer/asgd.yaml
deleted file mode 100644
index df90e0af..00000000
--- a/conf/optimizer/asgd.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-_target_: torch.optim.ASGD
-lr: ${training.lr}
diff --git a/conf/optimizer/deepspeed/adam.yaml b/conf/optimizer/deepspeed/adam.yaml
deleted file mode 100644
index 0a96eb7a..00000000
--- a/conf/optimizer/deepspeed/adam.yaml
+++ /dev/null
@@ -1,6 +0,0 @@
-_target_: deepspeed.ops.adam.FusedAdam
-lr: ${training.lr}
-betas: [ 0.8, 0.999 ]
-eps: 1e-8
-weight_decay: 3e-7
-adam_w_mode: False
diff --git a/conf/optimizer/deepspeed/adamw.yaml b/conf/optimizer/deepspeed/adamw.yaml
deleted file mode 100644
index f4d47bfd..00000000
--- a/conf/optimizer/deepspeed/adamw.yaml
+++ /dev/null
@@ -1,5 +0,0 @@
-_target_: deepspeed.ops.adam.FusedAdam
-lr: ${training.lr}
-betas: [ 0.8, 0.999 ]
-eps: 1e-8
-weight_decay: 3e-7
diff --git a/conf/optimizer/deepspeed/cpu_adam.yaml b/conf/optimizer/deepspeed/cpu_adam.yaml
deleted file mode 100644
index 49accef1..00000000
--- a/conf/optimizer/deepspeed/cpu_adam.yaml
+++ /dev/null
@@ -1,5 +0,0 @@
-_target_: deepspeed.ops.adam.DeepSpeedCPUAdam
-lr: ${training.lr}
-betas: [ 0.8, 0.999 ]
-eps: 1e-8
-weight_decay: 3e-7
diff --git a/conf/optimizer/deepspeed/onebitadam.yaml b/conf/optimizer/deepspeed/onebitadam.yaml
deleted file mode 100644
index 957028ca..00000000
--- a/conf/optimizer/deepspeed/onebitadam.yaml
+++ /dev/null
@@ -1,7 +0,0 @@
-_target_: deepspeed.ops.adam.OnebitAdam
-lr: ${training.lr}
-betas: [ 0.8, 0.999 ]
-eps: 1e-8
-weight_decay: 3e-7
-freeze_step: 400
-cuda_aware: True
diff --git a/conf/optimizer/rmsprop.yaml b/conf/optimizer/rmsprop.yaml
deleted file mode 100644
index 17babb5b..00000000
--- a/conf/optimizer/rmsprop.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-_target_: torch.optim.RMSprop
-lr: ${training.lr}
diff --git a/conf/optimizer/rprop.yaml b/conf/optimizer/rprop.yaml
deleted file mode 100644
index a01f73a5..00000000
--- a/conf/optimizer/rprop.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-_target_: torch.optim.Rprop
-lr: ${training.lr}
diff --git a/conf/optimizer/sgd.yaml b/conf/optimizer/sgd.yaml
deleted file mode 100644
index c0f4fef8..00000000
--- a/conf/optimizer/sgd.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-_target_: torch.optim.SGD
-lr: ${training.lr}
-momentum: 0.9
-weight_decay: 0.0005
diff --git a/conf/scheduler/constant_schedule.yaml b/conf/scheduler/constant_schedule.yaml
deleted file mode 100644
index 5c114b48..00000000
--- a/conf/scheduler/constant_schedule.yaml
+++ /dev/null
@@ -1,1 +0,0 @@
-_target_: transformers.get_constant_schedule
diff --git a/conf/scheduler/constant_schedule_with_warmup.yaml b/conf/scheduler/constant_schedule_with_warmup.yaml
deleted file mode 100644
index d17e0359..00000000
--- a/conf/scheduler/constant_schedule_with_warmup.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-_target_: transformers.get_constant_schedule_with_warmup
-num_warmup_steps: 0.1  # float values determines percentage of training steps to use as warmup
diff --git a/conf/scheduler/cosine_schedule_with_warmup.yaml b/conf/scheduler/cosine_schedule_with_warmup.yaml
deleted file mode 100644
index 6906a78e..00000000
--- a/conf/scheduler/cosine_schedule_with_warmup.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-_target_: transformers.get_cosine_schedule_with_warmup
-num_training_steps: -1  # -1 specifies to infer number of training steps
-num_warmup_steps: 0.1  # float values determines percentage of training steps to use as warmup
-num_cycles: 0.5
diff --git a/conf/scheduler/cosine_with_hard_restarts_schedule_with_warmup.yaml b/conf/scheduler/cosine_with_hard_restarts_schedule_with_warmup.yaml
deleted file mode 100644
index c6d1cd74..00000000
--- a/conf/scheduler/cosine_with_hard_restarts_schedule_with_warmup.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-_target_: transformers.get_cosine_with_hard_restarts_schedule_with_warmup
-num_training_steps: -1  # -1 specifies to infer number of training steps
-num_warmup_steps: 0.1  # float values determines percentage of training steps to use as warmup
-num_cycles: 1
diff --git a/conf/scheduler/linear_schedule_with_warmup.yaml b/conf/scheduler/linear_schedule_with_warmup.yaml
deleted file mode 100644
index 9637ea27..00000000
--- a/conf/scheduler/linear_schedule_with_warmup.yaml
+++ /dev/null
@@ -1,3 +0,0 @@
-_target_: transformers.get_linear_schedule_with_warmup
-num_training_steps: -1  # -1 specifies to infer number of training steps
-num_warmup_steps: 0.1  # float values determines percentage of training steps to use as warmup
diff --git a/conf/scheduler/polynomial_decay_schedule_with_warmup.yaml b/conf/scheduler/polynomial_decay_schedule_with_warmup.yaml
deleted file mode 100644
index 0aa0dff8..00000000
--- a/conf/scheduler/polynomial_decay_schedule_with_warmup.yaml
+++ /dev/null
@@ -1,5 +0,0 @@
-_target_: transformers.get_polynomial_decay_schedule_with_warmup
-num_training_steps: -1  # -1 specifies to infer number of training steps
-num_warmup_steps: 0.1  # float values determines percentage of training steps to use as warmup
-lr_end: 1e07
-power: 1.0
diff --git a/conf/task/default.yaml b/conf/task/default.yaml
deleted file mode 100644
index 6fdf1333..00000000
--- a/conf/task/default.yaml
+++ /dev/null
@@ -1,10 +0,0 @@
-# @package task
-defaults:
-  - /dataset@_group_: default
-
-# By default we turn off recursive instantiation, allowing the user to instantiate themselves at the appropriate times.
-_recursive_: false
-
-_target_: lightning_transformers.core.model.TaskTransformer
-optimizer: ${optimizer}
-scheduler: ${scheduler}
diff --git a/conf/task/nlp/default.yaml b/conf/task/nlp/default.yaml
deleted file mode 100644
index 2aad5a6a..00000000
--- a/conf/task/nlp/default.yaml
+++ /dev/null
@@ -1,6 +0,0 @@
-# @package task
-defaults:
-  - /task/default
-  - /backbone@_group_: nlp/default # use AutoModel backbone by default
-  - override /dataset@_group_: nlp/default
-backbone: ${backbone}
diff --git a/conf/task/nlp/language_modeling.yaml b/conf/task/nlp/language_modeling.yaml
deleted file mode 100644
index 15955cdf..00000000
--- a/conf/task/nlp/language_modeling.yaml
+++ /dev/null
@@ -1,6 +0,0 @@
-# @package task
-defaults:
-  - nlp/default
-  - override /dataset@_group_: nlp/language_modeling/default
-_target_: lightning_transformers.task.nlp.language_modeling.LanguageModelingTransformer
-downstream_model_type: transformers.AutoModelForCausalLM
diff --git a/conf/task/nlp/masked_language_modeling.yaml b/conf/task/nlp/masked_language_modeling.yaml
deleted file mode 100644
index 2b93135b..00000000
--- a/conf/task/nlp/masked_language_modeling.yaml
+++ /dev/null
@@ -1,6 +0,0 @@
-# @package task
-defaults:
-  - nlp/default
-  - override /dataset@_group_: nlp/masked_language_modeling/default
-_target_: lightning_transformers.task.nlp.masked_language_modeling.MaskedLanguageModelingTransformer
-downstream_model_type: transformers.AutoModelForMaskedLM
diff --git a/conf/task/nlp/multiple_choice.yaml b/conf/task/nlp/multiple_choice.yaml
deleted file mode 100644
index 35788095..00000000
--- a/conf/task/nlp/multiple_choice.yaml
+++ /dev/null
@@ -1,6 +0,0 @@
-# @package task
-defaults:
-  - nlp/default
-  - override /dataset@_group_: nlp/multiple_choice/default
-_target_: lightning_transformers.task.nlp.multiple_choice.MultipleChoiceTransformer
-downstream_model_type: transformers.AutoModelForMultipleChoice
diff --git a/conf/task/nlp/question_answering.yaml b/conf/task/nlp/question_answering.yaml
deleted file mode 100644
index ab29f330..00000000
--- a/conf/task/nlp/question_answering.yaml
+++ /dev/null
@@ -1,6 +0,0 @@
-# @package task
-defaults:
-  - nlp/default
-  - override /dataset@_group_: nlp/question_answering/default
-_target_: lightning_transformers.task.nlp.question_answering.QuestionAnsweringTransformer
-downstream_model_type: transformers.AutoModelForQuestionAnswering
diff --git a/conf/task/nlp/summarization.yaml b/conf/task/nlp/summarization.yaml
deleted file mode 100644
index 35fb09a5..00000000
--- a/conf/task/nlp/summarization.yaml
+++ /dev/null
@@ -1,12 +0,0 @@
-# @package task
-defaults:
-  - nlp/default
-  - override /backbone@_group_: nlp/seq2seq # Swap to a default seq2seq compatible backbone
-  - override /dataset@_group_: nlp/summarization/default
-_target_: lightning_transformers.task.nlp.summarization.SummarizationTransformer
-downstream_model_type: transformers.AutoModelForSeq2SeqLM
-cfg:
-  use_stemmer: True
-  val_target_max_length: 142
-  num_beams: null
-  compute_generate_metrics: True
diff --git a/conf/task/nlp/text_classification.yaml b/conf/task/nlp/text_classification.yaml
deleted file mode 100644
index dda2260d..00000000
--- a/conf/task/nlp/text_classification.yaml
+++ /dev/null
@@ -1,6 +0,0 @@
-# @package task
-defaults:
-  - nlp/default
-  - override /dataset@_group_: nlp/text_classification/default
-_target_: lightning_transformers.task.nlp.text_classification.TextClassificationTransformer
-downstream_model_type: transformers.AutoModelForSequenceClassification
diff --git a/conf/task/nlp/token_classification.yaml b/conf/task/nlp/token_classification.yaml
deleted file mode 100644
index d851272d..00000000
--- a/conf/task/nlp/token_classification.yaml
+++ /dev/null
@@ -1,6 +0,0 @@
-# @package task
-defaults:
-  - nlp/default
-  - override /dataset@_group_: nlp/token_classification/default
-_target_: lightning_transformers.task.nlp.token_classification.TokenClassificationTransformer
-downstream_model_type: transformers.AutoModelForTokenClassification
diff --git a/conf/task/nlp/translation.yaml b/conf/task/nlp/translation.yaml
deleted file mode 100644
index 2a530248..00000000
--- a/conf/task/nlp/translation.yaml
+++ /dev/null
@@ -1,13 +0,0 @@
-# @package task
-defaults:
-  - nlp/default
-  - override /backbone@_group_: nlp/seq2seq # Swap to a default seq2seq compatible backbone
-  - override /dataset@_group_: nlp/translation/default
-_target_: lightning_transformers.task.nlp.translation.TranslationTransformer
-downstream_model_type: transformers.AutoModelForSeq2SeqLM
-cfg:
-  n_gram: 4
-  smooth: False
-  val_target_max_length: 142
-  num_beams: null
-  compute_generate_metrics: True
diff --git a/conf/tokenizer/autotokenizer.yaml b/conf/tokenizer/autotokenizer.yaml
deleted file mode 100644
index fef7cd65..00000000
--- a/conf/tokenizer/autotokenizer.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-# @package tokenizer
-_target_: transformers.AutoTokenizer.from_pretrained
-pretrained_model_name_or_path: ${backbone.pretrained_model_name_or_path}
-use_fast: true
diff --git a/conf/trainer/benchmark.yaml b/conf/trainer/benchmark.yaml
deleted file mode 100644
index dace23d3..00000000
--- a/conf/trainer/benchmark.yaml
+++ /dev/null
@@ -1,10 +0,0 @@
-defaults:
-  - default # add args from default trainer conf
-devices: 1
-accelerator: ddp
-precision: 16
-checkpoint_callback: False
-limit_val_batches: 0
-limit_test_batches: 0
-callbacks:
-  _target_: lightning_transformers.core.callback.CUDACallback
diff --git a/conf/trainer/callbacks/lottery_ticket.yaml b/conf/trainer/callbacks/lottery_ticket.yaml
deleted file mode 100644
index cdd303d9..00000000
--- a/conf/trainer/callbacks/lottery_ticket.yaml
+++ /dev/null
@@ -1,3 +0,0 @@
-_target_: pytorch_lightning.callbacks.ModelPruning
-pruning_fn: 'l1_unstructured'
-use_lottery_ticket_hypothesis: True
diff --git a/conf/trainer/callbacks/model_pruning.yaml b/conf/trainer/callbacks/model_pruning.yaml
deleted file mode 100644
index a0791e4f..00000000
--- a/conf/trainer/callbacks/model_pruning.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-_target_: pytorch_lightning.callbacks.ModelPruning
-pruning_fn: 'l1_unstructured'
diff --git a/conf/trainer/callbacks/sparseml.yaml b/conf/trainer/callbacks/sparseml.yaml
deleted file mode 100644
index c8c1f6cc..00000000
--- a/conf/trainer/callbacks/sparseml.yaml
+++ /dev/null
@@ -1,3 +0,0 @@
-_target_: lightning_transformers.core.callback.TransformerSparseMLCallback
-output_dir: /content/MODELS
-recipe_path: /content/recipe.yaml
diff --git a/conf/trainer/deepspeed.yaml b/conf/trainer/deepspeed.yaml
deleted file mode 100644
index 7eaa014e..00000000
--- a/conf/trainer/deepspeed.yaml
+++ /dev/null
@@ -1,5 +0,0 @@
-defaults:
-  - default # inherit from default trainer conf
-  - plugins: deepspeed
-accelerator: ddp
-precision: 16
diff --git a/conf/trainer/default.yaml b/conf/trainer/default.yaml
deleted file mode 100644
index 4a875b6c..00000000
--- a/conf/trainer/default.yaml
+++ /dev/null
@@ -1,47 +0,0 @@
-_target_: pytorch_lightning.Trainer
-logger: tensorboard
-checkpoint_callback: True
-callbacks: null
-default_root_dir: null
-gradient_clip_val: 0.0
-process_position: 0
-num_nodes: 1
-num_processes: 1
-devices: null
-auto_select_gpus: False
-tpu_cores: null
-log_gpu_memory: null
-progress_bar_refresh_rate: 1
-overfit_batches: 0.0
-track_grad_norm: -1
-check_val_every_n_epoch: 1
-fast_dev_run: False
-accumulate_grad_batches: 1
-max_epochs: 1
-min_epochs: 1
-max_steps: null
-min_steps: null
-limit_train_batches: 1.0
-limit_val_batches: 1.0
-limit_test_batches: 1.0
-val_check_interval: 1.0
-flush_logs_every_n_steps: 100
-log_every_n_steps: 50
-accelerator: ddp
-sync_batchnorm: False
-precision: 32
-weights_summary: 'top'
-weights_save_path: null
-num_sanity_val_steps: 2
-resume_from_checkpoint: null
-profiler: null
-benchmark: False
-deterministic: False
-auto_lr_find: False
-replace_sampler_ddp: True
-terminate_on_nan: False
-auto_scale_batch_size: False
-prepare_data_per_node: True
-plugins: null
-amp_backend: 'native'
-move_metrics_to_cpu: False
diff --git a/conf/trainer/logger/sparsewandb.yaml b/conf/trainer/logger/sparsewandb.yaml
deleted file mode 100644
index a4a8f4ab..00000000
--- a/conf/trainer/logger/sparsewandb.yaml
+++ /dev/null
@@ -1,1 +0,0 @@
-_target_: lightning_transformers.core.loggers.WABLogger
diff --git a/conf/trainer/logger/tensorboard.yaml b/conf/trainer/logger/tensorboard.yaml
deleted file mode 100644
index febeaec0..00000000
--- a/conf/trainer/logger/tensorboard.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-_target_: pytorch_lightning.loggers.TensorBoardLogger
-save_dir: logs/
diff --git a/conf/trainer/logger/wandb.yaml b/conf/trainer/logger/wandb.yaml
deleted file mode 100644
index 261780dd..00000000
--- a/conf/trainer/logger/wandb.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-_target_: pytorch_lightning.loggers.WandbLogger
-project: lightning-transformers
diff --git a/conf/trainer/plugins/deepspeed.yaml b/conf/trainer/plugins/deepspeed.yaml
deleted file mode 100644
index 4ab62604..00000000
--- a/conf/trainer/plugins/deepspeed.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-_target_: pytorch_lightning.plugins.DeepSpeedPlugin
-stage: 2
diff --git a/conf/trainer/plugins/deepspeed_offload.yaml b/conf/trainer/plugins/deepspeed_offload.yaml
deleted file mode 100644
index 08c527c6..00000000
--- a/conf/trainer/plugins/deepspeed_offload.yaml
+++ /dev/null
@@ -1,3 +0,0 @@
-_target_: pytorch_lightning.plugins.DeepSpeedPlugin
-stage: 2
-cpu_offload: True
diff --git a/conf/trainer/plugins/deepspeed_offload_stage_3.yaml b/conf/trainer/plugins/deepspeed_offload_stage_3.yaml
deleted file mode 100644
index fcae1055..00000000
--- a/conf/trainer/plugins/deepspeed_offload_stage_3.yaml
+++ /dev/null
@@ -1,3 +0,0 @@
-_target_: pytorch_lightning.plugins.DeepSpeedPlugin
-stage: 3
-cpu_offload: True
diff --git a/conf/trainer/plugins/sharded.yaml b/conf/trainer/plugins/sharded.yaml
deleted file mode 100644
index de3e98f6..00000000
--- a/conf/trainer/plugins/sharded.yaml
+++ /dev/null
@@ -1,1 +0,0 @@
-_target_: pytorch_lightning.plugins.DDPShardedPlugin
diff --git a/conf/trainer/sharded.yaml b/conf/trainer/sharded.yaml
deleted file mode 100644
index 9ddd882e..00000000
--- a/conf/trainer/sharded.yaml
+++ /dev/null
@@ -1,5 +0,0 @@
-defaults:
-  - default # inherit from default trainer conf
-  - plugins: sharded
-accelerator: ddp
-precision: 16
diff --git a/conf/trainer/sparseml.yaml b/conf/trainer/sparseml.yaml
deleted file mode 100644
index 0c8cd4e9..00000000
--- a/conf/trainer/sparseml.yaml
+++ /dev/null
@@ -1,4 +0,0 @@
-defaults:
-  - default # inherit from default trainer conf
-  - callbacks: sparseml
-  - logger: sparsewandb
diff --git a/conf/trainer/tpus.yaml b/conf/trainer/tpus.yaml
deleted file mode 100644
index bc1454f8..00000000
--- a/conf/trainer/tpus.yaml
+++ /dev/null
@@ -1,3 +0,0 @@
-defaults:
-  - default # add args from default trainer conf
-tpu_cores: 8
diff --git a/conf/training/default.yaml b/conf/training/default.yaml
deleted file mode 100644
index 49b29ead..00000000
--- a/conf/training/default.yaml
+++ /dev/null
@@ -1,7 +0,0 @@
-run_test_after_fit: True
-lr: 5e-5
-output_dir: '.'
-
-# read in dataset
-batch_size: 16
-num_workers: 16
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 7e6ddebd..72f1b0b1 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -130,7 +130,6 @@ def _package_list_from_file(pfile):
     "scikit-learn": "sklearn",
     "Pillow": "PIL",
     "PyYAML": "yaml",
-    "hydra-core": "hydra",
     "rouge-score": "rouge_score",
 }
 MOCK_PACKAGES = ["omegaconf"]
diff --git a/docs/source/datasets/nlp/language_modeling_data.rst b/docs/source/datasets/nlp/language_modeling_data.rst
index b001b78b..751a2156 100644
--- a/docs/source/datasets/nlp/language_modeling_data.rst
+++ b/docs/source/datasets/nlp/language_modeling_data.rst
@@ -14,8 +14,6 @@ Below we have defined a csv file to use as our input data.
     this is the second sentence,
 
 
-When specifying the file path with hydra, it is important to use the absolute path to the file.
-
 .. code-block:: python
 
     from lightning_transformers.task.nlp.language_modeling import (
diff --git a/docs/source/index.rst b/docs/source/index.rst
index ac3472e0..ebf00a5e 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -34,13 +34,6 @@ Lightning Transformers
 
    optimizations/sparseml
 
-.. toctree::
-   :maxdepth: 1
-   :name: trainer
-   :caption: Trainer Features
-
-   trainer/logging
-
 .. toctree::
    :maxdepth: 1
    :name: advanced
diff --git a/docs/source/tasks/nlp/translation.rst b/docs/source/tasks/nlp/translation.rst
index 710d9ea4..eaa2f75f 100644
--- a/docs/source/tasks/nlp/translation.rst
+++ b/docs/source/tasks/nlp/translation.rst
@@ -69,8 +69,6 @@ Translation Inference Pipeline
 
 By default we use the translation pipeline, which requires a source text 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
 
     from transformers import AutoTokenizer
diff --git a/docs/source/trainer/logging.rst b/docs/source/trainer/logging.rst
deleted file mode 100644
index 3bb37ad1..00000000
--- a/docs/source/trainer/logging.rst
+++ /dev/null
@@ -1,45 +0,0 @@
-.. _logging:
-
-Logging
--------
-
-Lightning Transformers offers pre-built logging configurations out of the box to quickly get training visualizations.
-
-.. note::
-
-    If the logger saves to a directory, this will be located within the hydra default save directory. The configuration for this can be found in the `hydra run config <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/hydra/output/custom.yaml>`__ and modified on the CLI, e.g ``hydra.run.dir=./``.
-
-TensorBoard
-^^^^^^^^^^^
-
-To enable `TensorBoard <https://pytorch-lightning.readthedocs.io/en/latest/common/loggers.html#tensorboard>`__ logging to a directory, append ``+trainer/logger=tensorboard`` to your command. This appends the `tensorboard configuration <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/trainer/logger/tensorboard.yaml>`__.
-
-You can also modify the save directory by passing ``trainer.logger.save_dir=my_dir/``.
-
-.. code-block:: bash
-
-    python train.py task=nlp/multiple_choice dataset=nlp/multiple_choice/race +trainer/logger=tensorboard trainer.logger.save_dir=my_dir/
-
-
-Weights & Biases
-^^^^^^^^^^^^^^^^
-
-To enable `Weights & Biases <https://pytorch-lightning.readthedocs.io/en/latest/common/loggers.html#weights-and-biases>`__ logging to a directory, append ``+trainer/logger=wandb`` to your command. This appends the `wandb configuration <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/trainer/logger/wandb.yaml>`__.
-
-You can also modify the project by passing ``trainer.logger.project=my_lightning_transformers_exp``.
-
-.. code-block:: bash
-
-    python train.py task=nlp/multiple_choice dataset=nlp/multiple_choice/race +trainer/logger=wandb trainer.logger.project=my_lightning_transformers_exp/
-
-
-TestTube
-^^^^^^^^
-
-To enable `TestTube <https://pytorch-lightning.readthedocs.io/en/latest/common/loggers.html#test-tube>`__ logging to a directory, append ``+trainer/logger=testtube`` to your command. This appends the `testtube configuration <https://github.com/PyTorchLightning/lightning-transformers/blob/master/conf/trainer/logger/testtube.yaml>`__.
-
-You can also modify the save directory by passing ``trainer.logger.save_dir=my_dir/``.
-
-.. code-block:: bash
-
-    python train.py task=nlp/multiple_choice dataset=nlp/multiple_choice/race +trainer/logger=testtube trainer.logger.save_dir=my_dir/
diff --git a/lightning_transformers/cli/__init__.py b/lightning_transformers/cli/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/lightning_transformers/cli/predict.py b/lightning_transformers/cli/predict.py
deleted file mode 100644
index 226eb4c7..00000000
--- a/lightning_transformers/cli/predict.py
+++ /dev/null
@@ -1,74 +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 typing import Any, Dict, List, Mapping, Optional
-
-import hydra
-from hydra.utils import get_class
-from omegaconf import DictConfig, OmegaConf
-from pytorch_lightning.utilities.distributed import rank_zero_info
-
-from lightning_transformers.core import TaskTransformer
-from lightning_transformers.core.config import TaskConfig, TokenizerConfig
-from lightning_transformers.core.instantiator import HydraInstantiator, Instantiator
-
-
-def run(
-    x: Any,
-    instantiator: Instantiator,
-    checkpoint_path: Optional[str] = None,
-    task: TaskConfig = TaskConfig(),
-    model_data_kwargs: Optional[Dict[str, Any]] = None,
-    tokenizer: Optional[TokenizerConfig] = None,
-    pipeline_kwargs: Optional[dict] = None,  # mostly for the device
-    predict_kwargs: Optional[dict] = None,
-) -> List[Dict[str, Any]]:
-    model: TaskTransformer
-    if checkpoint_path:
-        model = get_class(task._target_).load_from_checkpoint(checkpoint_path)
-    else:
-        model = instantiator.model(
-            task, model_data_kwargs=model_data_kwargs, tokenizer=tokenizer, pipeline_kwargs=pipeline_kwargs
-        )
-
-    predict_kwargs = predict_kwargs or {}
-    if isinstance(x, Mapping):
-        return model.hf_predict(**x, **predict_kwargs)
-    else:
-        return model.hf_predict(x, **predict_kwargs)
-
-
-def main(cfg: DictConfig) -> Any:
-    rank_zero_info(OmegaConf.to_yaml(cfg))
-    instantiator = HydraInstantiator()
-    y = run(
-        cfg.x,
-        instantiator,
-        checkpoint_path=cfg.get("checkpoint_path"),
-        task=cfg.task,
-        model_data_kwargs=cfg.get("model_data_kwargs"),
-        tokenizer=cfg.get("tokenizer"),
-        pipeline_kwargs=cfg.get("pipeline_kwargs", {}),
-        predict_kwargs=cfg.get("predict_kwargs", {}),
-    )
-    rank_zero_info(y)
-    return y
-
-
-@hydra.main(config_path="../../conf", config_name="config")
-def hydra_entry(cfg: DictConfig) -> None:
-    main(cfg)
-
-
-if __name__ == "__main__":
-    hydra_entry()
diff --git a/lightning_transformers/cli/train.py b/lightning_transformers/cli/train.py
deleted file mode 100644
index 7a614d25..00000000
--- a/lightning_transformers/cli/train.py
+++ /dev/null
@@ -1,86 +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 typing import Any, Optional
-
-import hydra
-from omegaconf import DictConfig, OmegaConf
-from pytorch_lightning import LightningDataModule
-from pytorch_lightning.utilities.distributed import rank_zero_info
-
-from lightning_transformers.core import TaskTransformer, TransformerDataModule
-from lightning_transformers.core.config import TaskConfig, TokenizerConfig, TrainerConfig, TransformerDataConfig
-from lightning_transformers.core.instantiator import HydraInstantiator, Instantiator
-from lightning_transformers.core.utils import set_ignore_warnings
-
-
-def run(
-    instantiator: Instantiator,
-    ignore_warnings: bool = True,
-    run_test_after_fit: bool = True,
-    dataset: TransformerDataConfig = TransformerDataConfig(),
-    task: TaskConfig = TaskConfig(),
-    trainer: TrainerConfig = TrainerConfig(),
-    tokenizer: Optional[TokenizerConfig] = None,
-    logger: Optional[Any] = None,
-) -> None:
-    if ignore_warnings:
-        set_ignore_warnings()
-
-    data_module_kwargs = {}
-    if tokenizer is not None:
-        data_module_kwargs["tokenizer"] = tokenizer
-
-    data_module: TransformerDataModule = instantiator.data_module(dataset, **data_module_kwargs)
-    if data_module is None:
-        raise ValueError("No dataset found. Hydra hint: did you set `dataset=...`?")
-    if not isinstance(data_module, LightningDataModule):
-        raise ValueError(
-            "The instantiator did not return a DataModule instance." " Hydra hint: is `dataset._target_` defined?`"
-        )
-    data_module.setup("fit")
-
-    model: TaskTransformer = instantiator.model(task, model_data_kwargs=getattr(data_module, "model_data_kwargs", None))
-    trainer = instantiator.trainer(
-        trainer,
-        logger=logger,
-    )
-
-    trainer.fit(model, datamodule=data_module)
-    if run_test_after_fit:
-        trainer.test(model, datamodule=data_module)
-
-
-def main(cfg: DictConfig) -> None:
-    rank_zero_info(OmegaConf.to_yaml(cfg))
-    instantiator = HydraInstantiator()
-    logger = instantiator.logger(cfg)
-    run(
-        instantiator,
-        ignore_warnings=cfg.get("ignore_warnings"),
-        run_test_after_fit=cfg.get("training").get("run_test_after_fit"),
-        dataset=cfg.get("dataset"),
-        tokenizer=cfg.get("tokenizer"),
-        task=cfg.get("task"),
-        trainer=cfg.get("trainer"),
-        logger=logger,
-    )
-
-
-@hydra.main(config_path="../../conf", config_name="config")
-def hydra_entry(cfg: DictConfig) -> None:
-    main(cfg)
-
-
-if __name__ == "__main__":
-    hydra_entry()
diff --git a/lightning_transformers/core/config.py b/lightning_transformers/core/config.py
index 3438d6be..5a2739db 100644
--- a/lightning_transformers/core/config.py
+++ b/lightning_transformers/core/config.py
@@ -45,40 +45,3 @@ class TransformerDataConfig:
     test_subset_name: Optional[str] = None
     predict_subset_name: Optional[str] = None
     streaming: bool = False
-
-
-@dataclass
-class OptimizerConfig:
-    lr: float = 1e-5
-    weight_decay: float = 0.0
-
-
-@dataclass
-class SchedulerConfig:
-    num_training_steps: int = -1
-    num_warmup_steps: float = 0.1
-
-
-@dataclass
-class TrainerConfig:
-    ...
-
-
-@dataclass
-class BackboneConfig:
-    pretrained_model_name_or_path: Optional[str] = None
-
-
-@dataclass
-class TaskConfig:
-    optimizer: OptimizerConfig = OptimizerConfig()
-    scheduler: SchedulerConfig = SchedulerConfig()
-    downstream_model_type: Optional[str] = None
-    backbone: BackboneConfig = BackboneConfig()
-
-
-@dataclass
-class TokenizerConfig:
-    downstream_model_type: Optional[str] = None
-    pretrained_model_name_or_path: Optional[str] = None
-    use_fast: bool = True
diff --git a/lightning_transformers/core/instantiator.py b/lightning_transformers/core/instantiator.py
deleted file mode 100644
index 8e9868a6..00000000
--- a/lightning_transformers/core/instantiator.py
+++ /dev/null
@@ -1,113 +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.
-import logging
-from typing import TYPE_CHECKING, Optional
-
-import hydra
-import pytorch_lightning as pl
-import torch
-from omegaconf import DictConfig
-
-from lightning_transformers.core import TransformerDataModule
-
-if TYPE_CHECKING:
-    # avoid circular imports
-    from lightning_transformers.core import TaskTransformer
-
-
-class Instantiator:
-    def model(self, *args, **kwargs):
-        raise NotImplementedError("Child class must implement method")
-
-    def optimizer(self, *args, **kwargs):
-        raise NotImplementedError("Child class must implement method")
-
-    def scheduler(self, *args, **kwargs):
-        raise NotImplementedError("Child class must implement method")
-
-    def data_module(self, *args, **kwargs):
-        raise NotImplementedError("Child class must implement method")
-
-    def logger(self, *args, **kwargs):
-        raise NotImplementedError("Child class must implement method")
-
-    def trainer(self, *args, **kwargs):
-        raise NotImplementedError("Child class must implement method")
-
-    def instantiate(self, *args, **kwargs):
-        raise NotImplementedError("Child class must implement method")
-
-
-class HydraInstantiator(Instantiator):
-    def model(
-        self,
-        cfg: DictConfig,
-        model_data_kwargs: Optional[DictConfig] = None,
-        tokenizer: Optional[DictConfig] = None,
-        pipeline_kwargs: Optional[DictConfig] = None,
-    ) -> "TaskTransformer":
-        if model_data_kwargs is None:
-            model_data_kwargs = {}
-        model_data_kwargs = dict(model_data_kwargs)  # avoid ConfigKeyError: Key 'tokenizer' is not in struct`
-
-        # use `model_data_kwargs` to pass `tokenizer` and `pipeline_kwargs`
-        # as not all models might contain these parameters.
-        if tokenizer:
-            model_data_kwargs["tokenizer"] = self.instantiate(tokenizer)
-        if pipeline_kwargs:
-            model_data_kwargs["pipeline_kwargs"] = pipeline_kwargs
-
-        return self.instantiate(cfg, instantiator=self, **model_data_kwargs)
-
-    def optimizer(self, model: torch.nn.Module, cfg: DictConfig) -> torch.optim.Optimizer:
-        no_decay = ["bias", "LayerNorm.weight"]
-        grouped_parameters = [
-            {
-                "params": [
-                    p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay) and p.requires_grad
-                ],
-                "weight_decay": cfg.weight_decay,
-            },
-            {
-                "params": [
-                    p for n, p in model.named_parameters() if any(nd in n for nd in no_decay) and p.requires_grad
-                ],
-                "weight_decay": 0.0,
-            },
-        ]
-        return self.instantiate(cfg, grouped_parameters)
-
-    def scheduler(self, cfg: DictConfig, optimizer: torch.optim.Optimizer) -> torch.optim.lr_scheduler._LRScheduler:
-        return self.instantiate(cfg, optimizer=optimizer)
-
-    def data_module(
-        self,
-        cfg: DictConfig,
-        tokenizer: Optional[DictConfig] = None,
-    ) -> TransformerDataModule:
-        if tokenizer:
-            return self.instantiate(cfg, tokenizer=self.instantiate(tokenizer))
-        return self.instantiate(cfg)
-
-    def logger(self, cfg: DictConfig) -> Optional[logging.Logger]:
-        if cfg.get("log"):
-            if isinstance(cfg.trainer.logger, bool):
-                return cfg.trainer.logger
-            return self.instantiate(cfg.trainer.logger)
-
-    def trainer(self, cfg: DictConfig, **kwargs) -> pl.Trainer:
-        return self.instantiate(cfg, **kwargs)
-
-    def instantiate(self, *args, **kwargs):
-        return hydra.utils.instantiate(*args, **kwargs)
diff --git a/lightning_transformers/core/model.py b/lightning_transformers/core/model.py
index d4f81631..b18c60c6 100644
--- a/lightning_transformers/core/model.py
+++ b/lightning_transformers/core/model.py
@@ -16,14 +16,10 @@
 import pytorch_lightning as pl
 import torch
 import transformers
-from hydra.utils import get_class
-from pytorch_lightning.utilities import rank_zero_info, rank_zero_warn
+from pytorch_lightning.utilities import rank_zero_warn
 from transformers import PreTrainedTokenizerBase
 from transformers import pipeline as hf_transformers_pipeline
 
-from lightning_transformers.core.config import BackboneConfig, OptimizerConfig, SchedulerConfig
-from lightning_transformers.core.instantiator import Instantiator
-
 if TYPE_CHECKING:
     from transformers import AutoModel, Pipeline
 
@@ -37,13 +33,7 @@ class TaskTransformer(pl.LightningModule):
     Args:
         downstream_model_type: The AutoModel downstream model type.
             See https://huggingface.co/transformers/model_doc/auto.html
-        backbone: Config containing backbone specific arguments.
         pretrained_model_name_or_path: Huggingface model to use if backbone config not passed.
-        optimizer: Config containing optimizer specific arguments.
-        scheduler: Config containing scheduler specific arguments.
-        instantiator: Used to instantiate objects (when using Hydra).
-            If Hydra is not being used the instantiator is not required,
-            and functions that use instantiation such as ``configure_optimizers`` has been overridden.
         tokenizer: The pre-trained tokenizer.
         pipeline_kwargs: Arguments required for the HuggingFace inference pipeline class.
         **model_data_kwargs: Arguments passed from the data module to the class.
@@ -51,73 +41,38 @@ class TaskTransformer(pl.LightningModule):
 
     def __init__(
         self,
-        downstream_model_type: str,
-        optimizer: OptimizerConfig = OptimizerConfig(),
-        scheduler: SchedulerConfig = SchedulerConfig(),
+        downstream_model_type: Type["AutoModel"],
         pretrained_model_name_or_path: Optional[str] = None,
-        backbone: Optional[BackboneConfig] = None,
-        instantiator: Optional[Instantiator] = None,
         tokenizer: Optional[PreTrainedTokenizerBase] = None,
         pipeline_kwargs: Optional[dict] = None,
         **model_data_kwargs,
     ) -> None:
         super().__init__()
         self.save_hyperparameters()
-        model_cls: Type["AutoModel"] = get_class(downstream_model_type)
-        model_path = backbone.pretrained_model_name_or_path if backbone else pretrained_model_name_or_path
-        self.model = model_cls.from_pretrained(model_path, **model_data_kwargs)
+        self.model = downstream_model_type.from_pretrained(pretrained_model_name_or_path, **model_data_kwargs)
         self._tokenizer = tokenizer  # necessary for hf_pipeline
         self._hf_pipeline = None
         self._hf_pipeline_kwargs = pipeline_kwargs or {}
-        self.instantiator = instantiator
-        self.optimizer_cfg = optimizer
-        self.scheduler_cfg = scheduler
 
     def configure_optimizers(self) -> Dict:
-        if self.instantiator is None:
-            rank_zero_warn(
-                "You haven't specified an optimizer or lr scheduler. "
-                "Defaulting to AdamW with an lr of 1e-5 and linear warmup for 10% of steps. "
-                "To change this, override ``configure_optimizers`` in the TransformerModule."
-            )
-            optimizer = torch.optim.AdamW(self.parameters(), lr=1e-5)
-            num_training_steps, num_warmup_steps = self.compute_warmup(
-                num_training_steps=-1,
-                num_warmup_steps=0.1,
-            )
-            scheduler = transformers.get_linear_schedule_with_warmup(
-                optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps
-            )
-            return {
-                "optimizer": optimizer,
-                "lr_scheduler": {"scheduler": scheduler, "interval": "step", "frequency": 1},
-            }
-
-        optimizer = self.instantiator.optimizer(self.model, self.optimizer_cfg)
-        # compute_warmup needs the datamodule to be available when `self.num_training_steps`
-        # is called that is why this is done here and not in the __init__
-        self.scheduler_cfg.num_training_steps, self.scheduler_cfg.num_warmup_steps = self.compute_warmup(
-            num_training_steps=self.scheduler_cfg.num_training_steps,
-            num_warmup_steps=self.scheduler_cfg.num_warmup_steps,
+        rank_zero_warn(
+            "You haven't specified an optimizer or lr scheduler. "
+            "Defaulting to AdamW with an lr of 1e-5 and linear warmup for 10% of steps. "
+            "To change this, override ``configure_optimizers`` in the TransformerModule."
+        )
+        optimizer = torch.optim.AdamW(self.parameters(), lr=1e-5)
+        num_training_steps, num_warmup_steps = self.compute_warmup(
+            num_training_steps=-1,
+            num_warmup_steps=0.1,
+        )
+        scheduler = transformers.get_linear_schedule_with_warmup(
+            optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps
         )
-        rank_zero_info(f"Inferring number of training steps, set to {self.scheduler_cfg.num_training_steps}")
-        rank_zero_info(f"Inferring number of warmup steps from ratio, set to {self.scheduler_cfg.num_warmup_steps}")
-        scheduler = self.instantiator.scheduler(self.scheduler_cfg, optimizer)
         return {
             "optimizer": optimizer,
             "lr_scheduler": {"scheduler": scheduler, "interval": "step", "frequency": 1},
         }
 
-    def on_save_checkpoint(self, checkpoint: Dict[str, Any]):
-        # Save tokenizer from datamodule for predictions
-        if self.instantiator:
-            checkpoint["instantiator"] = self.instantiator
-
-    def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
-        if "tokenizer" in checkpoint:
-            self.tokenizer = checkpoint["tokenizer"]
-        self.instantiator = checkpoint.get("instantiator")
-
     @property
     def num_training_steps(self) -> int:
         return self.trainer.estimated_stepping_batches
diff --git a/lightning_transformers/task/nlp/language_modeling/model.py b/lightning_transformers/task/nlp/language_modeling/model.py
index a6fa31b8..ad68db3a 100644
--- a/lightning_transformers/task/nlp/language_modeling/model.py
+++ b/lightning_transformers/task/nlp/language_modeling/model.py
@@ -11,8 +11,15 @@
 # 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 typing import TYPE_CHECKING, Type
+
+import transformers
+
 from lightning_transformers.core import TaskTransformer
 
+if TYPE_CHECKING:
+    from transformers import AutoModel
+
 
 class LanguageModelingTransformer(TaskTransformer):
     """Defines ``LightningModule`` for the Language Modeling Task.
@@ -23,7 +30,9 @@ class LanguageModelingTransformer(TaskTransformer):
         **kwargs: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
     """
 
-    def __init__(self, *args, downstream_model_type: str = "transformers.AutoModelForCausalLM", **kwargs) -> None:
+    def __init__(
+        self, *args, downstream_model_type: Type["AutoModel"] = transformers.AutoModelForCausalLM, **kwargs
+    ) -> None:
         super().__init__(downstream_model_type, *args, **kwargs)
 
     def on_fit_start(self):
diff --git a/lightning_transformers/task/nlp/masked_language_modeling/model.py b/lightning_transformers/task/nlp/masked_language_modeling/model.py
index 1baf7200..06f0080e 100644
--- a/lightning_transformers/task/nlp/masked_language_modeling/model.py
+++ b/lightning_transformers/task/nlp/masked_language_modeling/model.py
@@ -11,8 +11,15 @@
 # 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 typing import TYPE_CHECKING, Type
+
+import transformers
+
 from lightning_transformers.core import TaskTransformer
 
+if TYPE_CHECKING:
+    from transformers import AutoModel
+
 
 class MaskedLanguageModelingTransformer(TaskTransformer):
     """Defines ``LightningModule`` for the Masked Language Modeling Task.
@@ -23,7 +30,9 @@ class MaskedLanguageModelingTransformer(TaskTransformer):
         **kwargs: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
     """
 
-    def __init__(self, *args, downstream_model_type: str = "transformers.AutoModelForMaskedLM", **kwargs) -> None:
+    def __init__(
+        self, *args, downstream_model_type: Type["AutoModel"] = transformers.AutoModelForMaskedLM, **kwargs
+    ) -> None:
         super().__init__(downstream_model_type, *args, **kwargs)
 
     def on_fit_start(self):
diff --git a/lightning_transformers/task/nlp/multiple_choice/model.py b/lightning_transformers/task/nlp/multiple_choice/model.py
index 1092ad19..175698af 100644
--- a/lightning_transformers/task/nlp/multiple_choice/model.py
+++ b/lightning_transformers/task/nlp/multiple_choice/model.py
@@ -11,11 +11,17 @@
 # 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 typing import TYPE_CHECKING, Type
+
 import torch
+import transformers
 from torchmetrics import Accuracy, Precision, Recall
 
 from lightning_transformers.core import TaskTransformer
 
+if TYPE_CHECKING:
+    from transformers import AutoModel
+
 
 class MultipleChoiceTransformer(TaskTransformer):
     """Defines ``LightningModule`` for the Multiple Choice Task.
@@ -27,7 +33,9 @@ class MultipleChoiceTransformer(TaskTransformer):
         **kwargs: :class:`lightning_transformers.core.nlp.HFTransformer` arguments.
     """
 
-    def __init__(self, *args, downstream_model_type: str = "transformers.AutoModelForMultipleChoice", **kwargs) -> None:
+    def __init__(
+        self, *args, downstream_model_type: Type["AutoModel"] = transformers.AutoModelForMultipleChoice, **kwargs
+    ) -> None:
         super().__init__(downstream_model_type, *args, **kwargs)
 
     def training_step(self, batch, batch_idx):
diff --git a/lightning_transformers/task/nlp/question_answering/model.py b/lightning_transformers/task/nlp/question_answering/model.py
index 0071bcac..4a635870 100644
--- a/lightning_transformers/task/nlp/question_answering/model.py
+++ b/lightning_transformers/task/nlp/question_answering/model.py
@@ -12,14 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from functools import partial
-from typing import Any
+from typing import TYPE_CHECKING, Any, Type
 
 import torch
+import transformers
 
 from lightning_transformers.core import TaskTransformer
 from lightning_transformers.task.nlp.question_answering import QuestionAnsweringDataModule
 from lightning_transformers.task.nlp.question_answering.datasets.squad.metric import SquadMetric
 
+if TYPE_CHECKING:
+    from transformers import AutoModel
+
 
 class QuestionAnsweringTransformer(TaskTransformer):
     """Defines ``LightningModule`` for the Question Answering Task.
@@ -32,7 +36,7 @@ class QuestionAnsweringTransformer(TaskTransformer):
     """
 
     def __init__(
-        self, *args, downstream_model_type: str = "transformers.AutoModelForQuestionAnswering", **kwargs
+        self, *args, downstream_model_type: Type["AutoModel"] = transformers.AutoModelForQuestionAnswering, **kwargs
     ) -> None:
         super().__init__(downstream_model_type, *args, **kwargs)
 
diff --git a/lightning_transformers/task/nlp/summarization/model.py b/lightning_transformers/task/nlp/summarization/model.py
index 016b0a71..11041e35 100644
--- a/lightning_transformers/task/nlp/summarization/model.py
+++ b/lightning_transformers/task/nlp/summarization/model.py
@@ -11,11 +11,17 @@
 # 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 typing import TYPE_CHECKING, Type
+
+import transformers
 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
+
 
 class SummarizationTransformer(Seq2SeqTransformer):
     """Defines ``LightningModule`` for the Summarization Task.
@@ -30,7 +36,7 @@ class SummarizationTransformer(Seq2SeqTransformer):
     def __init__(
         self,
         *args,
-        downstream_model_type: str = "transformers.AutoModelForSeq2SeqLM",
+        downstream_model_type: Type["AutoModel"] = transformers.AutoModelForSeq2SeqLM,
         cfg: SummarizationConfig = SummarizationConfig(),
         **kwargs
     ) -> None:
diff --git a/lightning_transformers/task/nlp/text_classification/model.py b/lightning_transformers/task/nlp/text_classification/model.py
index e563b01c..228eac8d 100644
--- a/lightning_transformers/task/nlp/text_classification/model.py
+++ b/lightning_transformers/task/nlp/text_classification/model.py
@@ -11,13 +11,17 @@
 # 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 typing import Any, Dict
+from typing import TYPE_CHECKING, Any, Dict, Type
 
 import torch
+import transformers
 from torchmetrics import Accuracy, Precision, Recall
 
 from lightning_transformers.core import TaskTransformer
 
+if TYPE_CHECKING:
+    from transformers import AutoModel
+
 
 class TextClassificationTransformer(TaskTransformer):
     """Defines ``LightningModule`` for the Text Classification Task.
@@ -30,7 +34,10 @@ class TextClassificationTransformer(TaskTransformer):
     """
 
     def __init__(
-        self, *args, downstream_model_type: str = "transformers.AutoModelForSequenceClassification", **kwargs
+        self,
+        *args,
+        downstream_model_type: Type["AutoModel"] = transformers.AutoModelForSequenceClassification,
+        **kwargs,
     ) -> None:
         super().__init__(downstream_model_type, *args, **kwargs)
         self.metrics = {}
diff --git a/lightning_transformers/task/nlp/token_classification/model.py b/lightning_transformers/task/nlp/token_classification/model.py
index 6db6509c..4a2b0286 100644
--- a/lightning_transformers/task/nlp/token_classification/model.py
+++ b/lightning_transformers/task/nlp/token_classification/model.py
@@ -11,13 +11,17 @@
 # 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 typing import Any, Dict, List, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Type, Union
 
 import torch
+import transformers
 from torchmetrics import Accuracy, F1Score, Precision, Recall
 
 from lightning_transformers.core import TaskTransformer
 
+if TYPE_CHECKING:
+    from transformers import AutoModel
+
 
 class TokenClassificationTransformer(TaskTransformer):
     """Defines ``LightningModule`` for the Text Classification Task.
@@ -33,7 +37,7 @@ def __init__(
         self,
         *args,
         labels: Union[int, List[str]],
-        downstream_model_type: str = "transformers.AutoModelForTokenClassification",
+        downstream_model_type: Type["AutoModel"] = transformers.AutoModelForTokenClassification,
         **kwargs,
     ) -> None:
         num_labels = labels if isinstance(labels, int) else len(labels)
diff --git a/lightning_transformers/task/nlp/translation/model.py b/lightning_transformers/task/nlp/translation/model.py
index 07416909..0793adae 100644
--- a/lightning_transformers/task/nlp/translation/model.py
+++ b/lightning_transformers/task/nlp/translation/model.py
@@ -11,12 +11,18 @@
 # 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 typing import TYPE_CHECKING, Type
+
+import transformers
 from torchmetrics.text.bleu import BLEUScore
 from transformers import MBartTokenizer
 
 from lightning_transformers.core.seq2seq.model import Seq2SeqTransformer
 from lightning_transformers.task.nlp.translation.config import TranslationConfig, TranslationDataConfig
 
+if TYPE_CHECKING:
+    from transformers import AutoModel
+
 
 class TranslationTransformer(Seq2SeqTransformer):
     """Defines ``LightningModule`` for the Translation Task.
@@ -31,7 +37,7 @@ class TranslationTransformer(Seq2SeqTransformer):
     def __init__(
         self,
         *args,
-        downstream_model_type: str = "transformers.AutoModelForSeq2SeqLM",
+        downstream_model_type: Type["AutoModel"] = transformers.AutoModelForSeq2SeqLM,
         cfg: TranslationConfig = TranslationConfig(),
         **kwargs,
     ) -> None:
diff --git a/lightning_transformers/task/vision/image_classification/model.py b/lightning_transformers/task/vision/image_classification/model.py
index 34c81f4f..0954e51b 100644
--- a/lightning_transformers/task/vision/image_classification/model.py
+++ b/lightning_transformers/task/vision/image_classification/model.py
@@ -11,15 +11,16 @@
 # 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 typing import TYPE_CHECKING, Any, Dict
+from typing import TYPE_CHECKING, Any, Dict, Type
 
 import torch
+import transformers
 from torchmetrics import Accuracy, Precision, Recall
 
 from lightning_transformers.core import TaskTransformer
 
 if TYPE_CHECKING:
-    from transformers import Pipeline
+    from transformers import AutoModel, Pipeline
 
 
 class ImageClassificationTransformer(TaskTransformer):
@@ -33,7 +34,7 @@ class ImageClassificationTransformer(TaskTransformer):
     """
 
     def __init__(
-        self, *args, downstream_model_type: str = "transformers.AutoModelForImageClassification", **kwargs
+        self, *args, downstream_model_type: Type["AutoModel"] = transformers.AutoModelForImageClassification, **kwargs
     ) -> None:
         super().__init__(downstream_model_type, *args, **kwargs)
         self.metrics = {}
diff --git a/predict.py b/predict.py
deleted file mode 100644
index f21836a9..00000000
--- a/predict.py
+++ /dev/null
@@ -1,14 +0,0 @@
-"""The shell entry point `$ pl-transformers-predict` is also available."""
-import hydra
-from omegaconf import DictConfig
-
-from lightning_transformers.cli.predict import main
-
-
-@hydra.main(config_path="./conf", config_name="config")
-def hydra_entry(cfg: DictConfig) -> None:
-    main(cfg)
-
-
-if __name__ == "__main__":
-    hydra_entry()
diff --git a/requirements.txt b/requirements.txt
index e1bbd924..5cdd9d26 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,9 +10,6 @@ torchmetrics>=0.7.0
 transformers>=4.11.2
 datasets
 
-# hydra
-hydra-core>=1.1.0
-
 # task-specific
 sentencepiece
 torchmetrics[text]
diff --git a/setup.py b/setup.py
index 4cde5acf..b4a315ec 100755
--- a/setup.py
+++ b/setup.py
@@ -38,12 +38,6 @@
     setup_requires=[],
     extras_require=extras,
     install_requires=setup_tools._load_requirements(_PATH_ROOT),
-    entry_points={
-        "console_scripts": [
-            "pl-transformers-train=lightning_transformers.cli.train:hydra_entry",
-            "pl-transformers-predict=lightning_transformers.cli.predict:hydra_entry",
-        ],
-    },
     project_urls={
         "Bug Tracker": "https://github.com/PyTorchLightning/lightning-transformers/issues",
         "Documentation": "https://lightning-transformers.readthedocs.io/en/stable/",
diff --git a/train.py b/train.py
deleted file mode 100644
index 4947916a..00000000
--- a/train.py
+++ /dev/null
@@ -1,14 +0,0 @@
-"""The shell entry point `$ pl-transformers-train` is also available."""
-import hydra
-from omegaconf import DictConfig
-
-from lightning_transformers.cli.train import main
-
-
-@hydra.main(config_path="./conf", config_name="config")
-def hydra_entry(cfg: DictConfig) -> None:
-    main(cfg)
-
-
-if __name__ == "__main__":
-    hydra_entry()


Test patch:
diff --git a/conf/trainer/logger/testtube.yaml b/conf/trainer/logger/testtube.yaml
deleted file mode 100644
index 0829ba81..00000000
--- a/conf/trainer/logger/testtube.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-_target_: pytorch_lightning.loggers.TestTubeLogger
-save_dir: logs/
diff --git a/tests/cli/test_predict.py b/tests/cli/test_predict.py
deleted file mode 100644
index deec8c02..00000000
--- a/tests/cli/test_predict.py
+++ /dev/null
@@ -1,40 +0,0 @@
-import shutil
-import subprocess
-from unittest import mock
-from unittest.mock import ANY, MagicMock
-
-from omegaconf import OmegaConf
-
-import lightning_transformers.cli.predict as cli
-
-
-def test_predict_run():
-    model = MagicMock()
-    instantiator = MagicMock()
-    instantiator.model.return_value = model
-
-    cli.run("test", instantiator)
-
-    model.hf_predict.assert_called_with("test")
-
-
-@mock.patch("lightning_transformers.cli.predict.run")
-def test_predict_main(run_mock):
-    cfg = OmegaConf.create({"x": "foo", "task": "bar"})
-    cli.main(cfg)
-    run_mock.assert_called_with(
-        "foo",
-        ANY,
-        checkpoint_path=None,
-        task="bar",
-        model_data_kwargs=None,
-        tokenizer=None,
-        pipeline_kwargs={},
-        predict_kwargs={},
-    )
-
-
-def test_predict_entry_point():
-    proc = subprocess.run([shutil.which("pl-transformers-predict"), "--help"], stdout=subprocess.PIPE)
-    help = proc.stdout.decode().strip()
-    assert help.startswith("predict is powered by Hydra")
diff --git a/tests/cli/test_train.py b/tests/cli/test_train.py
deleted file mode 100644
index 5775ba62..00000000
--- a/tests/cli/test_train.py
+++ /dev/null
@@ -1,54 +0,0 @@
-import shutil
-import subprocess
-from unittest import mock
-from unittest.mock import ANY, MagicMock
-
-import pytest
-from omegaconf import OmegaConf
-from pytorch_lightning import LightningDataModule
-
-import lightning_transformers.cli.train as cli
-
-
-def test_train_run_raises():
-    instantiator = MagicMock()
-    instantiator.data_module.return_value = None
-    with pytest.raises(ValueError, match="No dataset found"):
-        cli.run(instantiator)
-
-    with pytest.raises(ValueError, match="did not return a DataModule"):
-        cli.run(MagicMock())
-
-
-def test_train_run():
-    instantiator = MagicMock()
-    instantiator.data_module.return_value = MagicMock(spec=LightningDataModule)
-
-    cli.run(instantiator)
-
-    instantiator.data_module.return_value.setup.assert_called_with("fit")
-    trainer_mock = instantiator.trainer.return_value
-    trainer_mock.fit.assert_called()
-    trainer_mock.test.assert_called()
-
-
-@mock.patch("lightning_transformers.cli.train.run")
-def test_train_main(run_mock):
-    cfg = OmegaConf.create({"training": {"run_test_after_fit": True}})
-    cli.main(cfg)
-    run_mock.assert_called_with(
-        ANY,
-        ignore_warnings=ANY,
-        run_test_after_fit=True,
-        dataset=ANY,
-        task=ANY,
-        trainer=ANY,
-        logger=None,
-        tokenizer=ANY,
-    )
-
-
-def test_train_entry_point():
-    proc = subprocess.run([shutil.which("pl-transformers-train"), "--help"], stdout=subprocess.PIPE)
-    help = proc.stdout.decode().strip()
-    assert help.startswith("train is powered by Hydra")
diff --git a/tests/core/nlp/test_nlp_model.py b/tests/core/nlp/test_nlp_model.py
index 32bb6909..8c85bed8 100644
--- a/tests/core/nlp/test_nlp_model.py
+++ b/tests/core/nlp/test_nlp_model.py
@@ -109,14 +109,10 @@ class TestModel(TaskTransformer):
         def hf_pipeline_task(self):
             return "task_name"
 
-    downstream_model_type = "model_type"
     cls_mock = MagicMock()
-    backbone_config = MagicMock()
-
-    with patch("lightning_transformers.core.model.get_class", return_value=cls_mock) as get_class_mock:
-        model = TestModel(downstream_model_type, backbone=backbone_config, pipeline_kwargs=dict(device=0), foo="bar")
-    get_class_mock.assert_called_once_with(downstream_model_type)
-    cls_mock.from_pretrained.assert_called_once_with(backbone_config.pretrained_model_name_or_path, foo="bar")
+    model = TestModel(
+        cls_mock, pretrained_model_name_or_path="prajjwal1/bert-tiny", pipeline_kwargs=dict(device=0), foo="bar"
+    )
 
     with patch("lightning_transformers.core.model.hf_transformers_pipeline") as pipeline_mock:
         model.hf_pipeline
diff --git a/tests/task/nlp/test_language_modeling.py b/tests/task/nlp/test_language_modeling.py
index a6aed6e7..e442af42 100644
--- a/tests/task/nlp/test_language_modeling.py
+++ b/tests/task/nlp/test_language_modeling.py
@@ -3,6 +3,7 @@
 
 import pytest
 import pytorch_lightning as pl
+import transformers
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.language_modeling import (
@@ -45,7 +46,7 @@ def test_smoke_predict():
 
 def test_model_has_correct_cfg():
     model = LanguageModelingTransformer(pretrained_model_name_or_path="prajjwal1/bert-tiny")
-    assert model.hparams.downstream_model_type == "transformers.AutoModelForCausalLM"
+    assert isinstance(model.model, transformers.BertLMHeadModel)
 
 
 def test_datamodule_has_correct_cfg():
diff --git a/tests/task/nlp/test_masked_language_modeling.py b/tests/task/nlp/test_masked_language_modeling.py
index 89355856..bb316af5 100644
--- a/tests/task/nlp/test_masked_language_modeling.py
+++ b/tests/task/nlp/test_masked_language_modeling.py
@@ -3,6 +3,7 @@
 
 import pytest
 import pytorch_lightning as pl
+import transformers
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.masked_language_modeling import (
@@ -43,7 +44,7 @@ def test_smoke_predict():
 
 def test_model_has_correct_cfg():
     model = MaskedLanguageModelingTransformer(pretrained_model_name_or_path="prajjwal1/bert-tiny")
-    assert model.hparams.downstream_model_type == "transformers.AutoModelForMaskedLM"
+    assert isinstance(model.model, transformers.BertForMaskedLM)
 
 
 def test_datamodule_has_correct_cfg():
diff --git a/tests/task/nlp/test_multiple_choice.py b/tests/task/nlp/test_multiple_choice.py
index f50fd2b3..60c3cb78 100644
--- a/tests/task/nlp/test_multiple_choice.py
+++ b/tests/task/nlp/test_multiple_choice.py
@@ -3,6 +3,7 @@
 
 import pytest
 import pytorch_lightning as pl
+import transformers
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.multiple_choice import (
@@ -34,7 +35,7 @@ def test_smoke_train(hf_cache_path):
 
 def test_model_has_correct_cfg():
     model = MultipleChoiceTransformer(pretrained_model_name_or_path="bert-base-cased")
-    assert model.hparams.downstream_model_type == "transformers.AutoModelForMultipleChoice"
+    assert isinstance(model.model, transformers.BertForMultipleChoice)
 
 
 def test_datamodule_has_correct_cfg():
diff --git a/tests/task/nlp/test_question_answering.py b/tests/task/nlp/test_question_answering.py
index f15a449f..e6d6d4e4 100644
--- a/tests/task/nlp/test_question_answering.py
+++ b/tests/task/nlp/test_question_answering.py
@@ -3,6 +3,7 @@
 
 import pytest
 import pytorch_lightning as pl
+import transformers
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.question_answering import (
@@ -55,7 +56,7 @@ def test_smoke_predict():
 
 def test_model_has_correct_cfg():
     model = QuestionAnsweringTransformer(pretrained_model_name_or_path="bert-base-cased")
-    assert model.hparams.downstream_model_type == "transformers.AutoModelForQuestionAnswering"
+    assert isinstance(model.model, transformers.BertForQuestionAnswering)
 
 
 def test_datamodule_has_correct_cfg():
diff --git a/tests/task/nlp/test_summarization.py b/tests/task/nlp/test_summarization.py
index 3c552934..2f73e13c 100644
--- a/tests/task/nlp/test_summarization.py
+++ b/tests/task/nlp/test_summarization.py
@@ -3,6 +3,7 @@
 
 import pytest
 import pytorch_lightning as pl
+import transformers
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.summarization import (
@@ -63,7 +64,7 @@ def test_smoke_predict():
 
 def test_model_has_correct_cfg():
     model = SummarizationTransformer(pretrained_model_name_or_path="t5-base")
-    assert model.hparams.downstream_model_type == "transformers.AutoModelForSeq2SeqLM"
+    assert isinstance(model.model, transformers.T5ForConditionalGeneration)
     assert type(model.cfg) is SummarizationConfig
 
 
diff --git a/tests/task/nlp/test_text_classification.py b/tests/task/nlp/test_text_classification.py
index c5690b2a..52a51377 100644
--- a/tests/task/nlp/test_text_classification.py
+++ b/tests/task/nlp/test_text_classification.py
@@ -3,6 +3,7 @@
 
 import pytest
 import pytorch_lightning as pl
+import transformers
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.text_classification import (
@@ -69,7 +70,7 @@ def test_smoke_predict():
 
 def test_model_has_correct_cfg():
     model = TextClassificationTransformer(pretrained_model_name_or_path="prajjwal1/bert-tiny")
-    assert model.hparams.downstream_model_type == "transformers.AutoModelForSequenceClassification"
+    assert isinstance(model.model, transformers.BertForSequenceClassification)
 
 
 def test_datamodule_has_correct_cfg():
diff --git a/tests/task/nlp/test_token_classification.py b/tests/task/nlp/test_token_classification.py
index 116d2ac6..dd7d662a 100644
--- a/tests/task/nlp/test_token_classification.py
+++ b/tests/task/nlp/test_token_classification.py
@@ -3,6 +3,7 @@
 
 import pytest
 import pytorch_lightning as pl
+import transformers
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.token_classification import (
@@ -54,7 +55,7 @@ def test_model_has_correct_cfg():
         pretrained_model_name_or_path="bert-base-cased",
         labels=2,
     )
-    assert model.hparams.downstream_model_type == "transformers.AutoModelForTokenClassification"
+    assert isinstance(model.model, transformers.BertForTokenClassification)
     assert model.num_labels == 2
 
 
diff --git a/tests/task/nlp/test_translation.py b/tests/task/nlp/test_translation.py
index 126c71dd..5991cff2 100644
--- a/tests/task/nlp/test_translation.py
+++ b/tests/task/nlp/test_translation.py
@@ -3,6 +3,7 @@
 
 import pytest
 import pytorch_lightning as pl
+import transformers
 from transformers import AutoTokenizer
 
 from lightning_transformers.task.nlp.translation import (
@@ -52,7 +53,7 @@ def test_smoke_predict():
 
 def test_model_has_correct_cfg():
     model = TranslationTransformer(pretrained_model_name_or_path="patrickvonplaten/t5-tiny-random")
-    assert model.hparams.downstream_model_type == "transformers.AutoModelForSeq2SeqLM"
+    assert isinstance(model.model, transformers.T5ForConditionalGeneration)
     assert type(model.cfg) is TranslationConfig
 
 
diff --git a/tests/task/vision/test_image_classification.py b/tests/task/vision/test_image_classification.py
index 179a1507..f01c8246 100644
--- a/tests/task/vision/test_image_classification.py
+++ b/tests/task/vision/test_image_classification.py
@@ -3,6 +3,7 @@
 
 import pytest
 import pytorch_lightning as pl
+import transformers
 from transformers import AutoFeatureExtractor
 
 from lightning_transformers.task.vision.image_classification import (
@@ -46,7 +47,7 @@ def test_model_has_correct_cfg():
     model = ImageClassificationTransformer(
         pretrained_model_name_or_path="nateraw/tiny-vit-random",
     )
-    assert model.hparams.downstream_model_type == "transformers.AutoModelForImageClassification"
+    assert isinstance(model.model, transformers.ViTForImageClassification)
 
 
 def test_datamodule_has_correct_cfg():