Quick Start


pip install lightning-transformers

Install bleeding-edge:

pip install git+ --upgrade

Install all optional dependencies as well:

pip install lightning-transformers["extra"]

Using Lightning-Transformers

Lightning Transformers has a collections of tasks for common NLP problems such as Language Modeling, Translation and more. To use, simply:

  1. Pick a task to train (passed to as task=)

  2. Pick a dataset (passed to as dataset=)

  3. Customize the backbone, optimizer, or any component within the config

  4. Add any Lightning supported parameters and optimizations

python \
     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


In this example we will finetune the Text Classification task on the emotion dataset.

To fine-tune using the Text Classification default bert-based-cased model, use the following:

python task=nlp/text_classification dataset=nlp/text_classification/emotion

Change the Backbone

Swap to the RoBERTa model:

python 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/ for all supported optimizers):

python task=nlp/text_classification dataset=nlp/text_classification/emotion optimizer=rmsprop

For more info on how to override configuration, see The Config Structure.

Lightning Trainer Options

We expose all Pytorch Lightning Trainer 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 by setting the parameter via the CLI:

python task=nlp/text_classification dataset=nlp/text_classification/emotion trainer.precision=16

Setting the maximum epochs:

python task=nlp/translation dataset=nlp/translation/wmt16 trainer.max_epochs=4

Using multiple GPUs:

python task=nlp/translation dataset=nlp/translation/wmt16 trainer.gpus=4

Using TPUs:

python task=nlp/translation dataset=nlp/translation/wmt16 trainer.tpu_cores=8

See the Pytorch Lightning Trainer or conf/trainer/default for all parameters.


Run inference once model trained (experimental):

python task=nlp/text_classification +checkpoint_path=/path/to/model.ckpt +x="Classify this sentence."

# Returns [{'label': 'LABEL_0', 'score': 0.545...}]

You can also run prediction using the default HuggingFace pre-trained model:

python task=nlp/text_classification +x="Classify this sentence."

Or run prediction on a specified HuggingFace pre-trained model:

python task=nlp/text_classification backbone.pretrained_model_name_or_path=bert-base-cased +x="Classify this sentence."
Read the Docs v: stable
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.