Text classification

Public LLMs are great generalists, but they can't always meet your domain, cost, speed, or privacy needs. Lightning makes training custom text models fast, scalable, and production-ready - without reinventing the wheel in PyTorch.

Note: This tutorial assumes basic familiarity with PyTorch

LLMs can't solve all problems

Public LLMs are great generalists, but they can’t always meet your domain, cost, speed, or privacy needs. When you need custom models, training from scratch in PyTorch gets messy fast - juggling tokenization, batching, attention masks, device placement, and scaling by hand.

Train without reinventing

Lightning makes training custom text models fast, scalable, and production-ready - without reinventing the wheel. With LightningModule and LightningDataModule, you organize model, data, and training infrastructure into modular, reusable pieces - making it easy to swap models, update datasets, and scale across GPUs.

Full code

Here's the full code with explanations for each step. This PyTorch Lightning module trains a BERT model to classify movie reviews as positive or negative.

Modularize your modelWrap a Hugging Face BERT model inside a LightningModule for clean, reusable training loops.Simplify training logicDefine a clear training step - Lightning handles device placement, optimization, and logging for you.Swap optimizers easilyConfigure optimizers independently to make experimentation and fine-tuning faster.Separate your data pipelineLoad, tokenize, batch, and shuffle your dataset cleanly with a LightningDataModule - no cluttering your model code.Launch scalable trainingInitialize your modular model and data pipeline - ready for GPU scaling, mixed precision, and production.Productionize with one lineTrain your model with automatic scaling - and upload to the Lightning AI model hub for easy deployment and sharing.

1/6

2/6

3/6

4/6

5/6

6/6

import torch
from transformers import BertTokenizer, BertForSequenceClassification
from datasets import load_dataset
import lightning as L

class LitTextClassification(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = BertForSequenceClassification.from_pretrained("bert-base-uncased").train()

    def training_step(self, batch):
        output = self.model(
            input_ids=batch["input_ids"],
            attention_mask=batch["attention_mask"],
            labels=batch["label"],
        )
        self.log("train_loss", output.loss)
        return output.loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.model.parameters(), lr=0.001)

class TextClassificationData(L.LightningDataModule):
    def prepare_data(self):
        load_dataset("imdb")

    def train_dataloader(self):
        tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
        dataset = load_dataset("imdb")["train"]
        dataset = dataset.map(
            lambda sample: tokenizer(sample["text"], padding="max_length", truncation=True))
        dataset.set_format(type="torch")
        return torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True, num_workers=2)

if __name__ == "__main__":
    model = LitTextClassification()
    data = TextClassificationData()

    trainer = L.Trainer(model_registry="LitTextClassification")
    trainer.fit(model, data)

Studio templateOpen Template →