Table of Contents
Problem Setup and What the Pipeline Produces
A basic text classification pipeline takes raw text examples and produces a trained model that maps each text to a label such as positive or negative sentiment, topic categories, or spam versus not spam. In practice you will build a chain of steps that turns strings into tensors, batches them correctly, feeds them into a sequence model, computes a classification loss, and evaluates accuracy or related metrics.
The core output of the pipeline is a PyTorch nn.Module whose forward accepts tokenized sequences and returns class logits of shape (batch_size, num_classes).
From Raw Text to Integer Tokens
Neural networks do not consume strings directly, they consume numbers. For a beginner friendly pipeline, start with a simple vocabulary that maps tokens to integer ids. Tokenization can be as simple as splitting on whitespace and lowercasing, as long as you apply the same rules at training and inference time. You will also reserve a few special tokens such as padding and unknown tokens.
You must use the same tokenization rules and the same vocabulary mapping at training time and at inference time. If these differ, the model will see different ids for the same words and predictions will be unreliable.
A minimal vocabulary typically includes "<pad>" and "<unk>". Every token not in the vocabulary maps to "<unk>". Each example becomes a Python list of integers.
Padding, Truncation, and the Collate Function
Text examples have different lengths, but batches must be rectangular tensors. The standard approach is to pad shorter sequences up to the length of the longest sequence in the batch, and optionally truncate very long sequences to a maximum length for speed and stability.
In PyTorch, the clean way to do this is in a custom collate_fn passed to DataLoader. The collate function receives a list of samples and returns tensors such as input_ids, lengths, and labels.
A typical batch output is input_ids of shape (batch_size, seq_len) with dtype torch.long, plus lengths of shape (batch_size,) that records the true lengths before padding. Many sequence models need lengths so they can ignore padding, or so you can pool only over real tokens.
Padding id must be consistent everywhere. If you use pad_id = 0 in batching, you must also configure your embedding layer with padding_idx=pad_id so the padding token does not receive meaningful gradients.
Model Choice: Embedding Plus Sequence Encoder Plus Classifier
A simple and effective baseline for text classification is an embedding layer followed by a sequence encoder and then a classifier head.
The embedding layer maps token ids to dense vectors. If V is vocabulary size and d is embedding dimension, the embedding produces a tensor of shape (batch_size, seq_len, d).
The sequence encoder can be an nn.RNN, nn.GRU, or nn.LSTM. For classification you need to convert the sequence of hidden states into a single vector per example. Common beginner friendly choices are using the final hidden state, or using a masked mean over time. Which pooling strategy you use is part of your model definition, but your pipeline must provide whatever additional tensors it requires, such as lengths for masking.
The classifier head is typically a linear layer that maps the pooled vector to logits. Logits are unnormalized scores. For C classes, logits shape is (batch_size, C).
Loss Function and Target Format
For standard single label classification, use nn.CrossEntropyLoss. Your model outputs logits, not probabilities, and your labels are integer class indices.
Targets must have shape (batch_size,) and dtype torch.long. If your labels are one hot vectors, convert them to class indices before computing loss.
nn.CrossEntropyLoss expects raw logits and integer class targets. Do not apply softmax before the loss. Do not pass one hot labels.
A Minimal Training Step for Text Classification
A single optimization step in this pipeline usually looks like this in conceptual terms. You take a batch from the DataLoader, move tensors to the device, run the forward pass to get logits, compute cross entropy loss against labels, backpropagate, and update parameters.
The key text specific detail is that you pass input_ids and often lengths into the model, and you ensure dtypes are correct. Token ids and labels are torch.long. Most other tensors are torch.float32.
Evaluation: Turning Logits Into Predictions
During evaluation you typically compute predicted classes with argmax over the class dimension. If logits is (batch_size, C), predictions are pred = logits.argmax(dim=1).
Accuracy can then be computed by comparing pred with labels. If you need probabilities for interpretability, you can compute probs = torch.softmax(logits, dim=1) after the forward pass, but keep the training loss on logits.
Handling Padding Correctly in Pooling
If you use mean pooling over token representations, you must not average over padding tokens. This is where the lengths tensor or an attention mask becomes essential. The usual approach is to create a boolean mask where real tokens are True and padded positions are False, then compute a masked sum divided by the number of real tokens.
If you average over padded positions, the representation of short sentences will be biased toward the padding embedding and training will be noisier and less stable.
Inference: Reusing the Same Preprocessing
At inference time you repeat the same preprocessing, tokenize, map to ids with the same vocabulary, pad or truncate using the same rules, then feed the tensor through the model in model.eval() mode with gradients disabled. The output logits are converted to predicted labels with argmax.
For deployment readiness, save both the model weights and the vocabulary object or mapping, because the vocabulary is part of the model’s input contract.
A Practical “First Working” Checklist
A basic text classification pipeline is correct when token ids are torch.long, padding uses a fixed pad_id, the embedding uses padding_idx=pad_id, labels are class indices with dtype torch.long, the model returns logits of shape (batch_size, num_classes), and training loss decreases on the training set while validation metrics improve.
If you can not overfit a tiny subset of your data, for example 20 to 100 examples, there is usually a bug in preprocessing, batching, label format, or the model output shape. Fix that before scaling up.