<jupyter_start><jupyter_text>Introduction This tutorial will demonstrate how to fine-tune a pretrained HuggingFace transformer using the composer library! Composer provides a highly optimized and functional training loop and the ability to compose several methods that can accelerate training.We will focus on fine-tuning a pretrained BERT-base model on the Stanford Sentiment Treebank v2 (SST-2) dataset. After fine-tuning, the BERT model should be able to determine if a setence has positive or negative sentiment.Let's do this 🚀  Defining a Composer Model  The first task is to create a composer model. A composer model defines four components for the composer trainer:- `forward()` - parses the dataloader output for the model's forward function and extracts the necessary components of the model's output for loss calculation.- `loss()` - computes the loss for the current batch using the model and dataloader outputs.- `validate()` - parses the dataloader and model output for the torchmetrics.- `metrics()` - defines the torchmetrics to use during training/validation.<jupyter_code>import transformers
from torchmetrics import Accuracy
from torchmetrics.collections import MetricCollection
from composer.models.base import ComposerModel
from composer.models.nlp_metrics import LanguageCrossEntropyLoss

# Define a Composer Model
class ComposerBERT(ComposerModel):
  def __init__(self, model):
    super().__init__()
    self.module = model

    # Metrics
    self.train_loss = LanguageCrossEntropyLoss()
    self.val_loss = LanguageCrossEntropyLoss()

    self.train_acc = Accuracy()
    self.val_acc = Accuracy()

  def forward(self, batch):
    output = self.module(**batch)
    return output

  def loss(self, outputs, batch):
    return outputs['loss']

  def validate(self, batch):
    labels = batch.pop('labels')
    output = self.forward(batch)
    output = output['logits']

    return (output, labels)
  
  def metrics(self, train: bool = False):
    return MetricCollection([self.train_loss, self.train_acc]) if train else MetricCollection([self.val_loss, self.val_acc])

# Create a BERT sequence classification model using HuggingFace transformsers
model = transformers.AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2) # in BERT hparams

# Package as a composer model
composer_model = ComposerBERT(model)<jupyter_output><empty_output><jupyter_text>Creating dataloaders Next, we will download and tokenize the SST-2 datasets.<jupyter_code>import datasets
from multiprocessing import cpu_count

# Create BERT tokenizer
tokenizer = transformers.AutoTokenizer.from_pretrained('bert-base-uncased') # from transformer_shared
def tokenize_function(sample):
  return tokenizer(text=sample['sentence'], padding="max_length", max_length=256, truncation=True)

# Tokenize SST-2
sst2_dataset = datasets.load_dataset("glue", "sst2")
tokenized_sst2_dataset = sst2_dataset.map(tokenize_function,
                                          batched=True, 
                                          num_proc=cpu_count(),
                                          batch_size=1000,
                                          remove_columns=['idx', 'sentence'])

# Split dataset into train and validation sets
train_dataset = tokenized_sst2_dataset["train"]
eval_dataset = tokenized_sst2_dataset["validation"]<jupyter_output><empty_output><jupyter_text>Here, we will create a PyTorch `Dataloader` for each of the datasets generated in the previous block.<jupyter_code>from torch.utils.data import DataLoader
data_collator = transformers.data.data_collator.default_data_collator
train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=False, drop_last=False, collate_fn=data_collator)
eval_dataloader = DataLoader(eval_dataset,batch_size=16, shuffle=False, drop_last=False, collate_fn=data_collator)<jupyter_output><empty_output><jupyter_text>To use the composer `Trainer`, we need to define a `split_batch` function. This function defines how to split the dataloader output into several "microbatches". Microbatchs are chunks of the batch that were divided based on the amount of gradient accumulation used.<jupyter_code>from composer.core import DataSpec

def split_batch_dict(batch, n_microbatches: int):
      chunked = {k: v.chunk(n_microbatches) for k, v in batch.items()}
      num_chunks = len(list(chunked.values())[0])
      return [{k: v[idx] for k, v in chunked.items()} for idx in range(num_chunks)]

train_dataspec = DataSpec(dataloader=train_dataloader,
                          split_batch=split_batch_dict)
eval_dataspec = DataSpec(dataloader=eval_dataloader,
                         split_batch=split_batch_dict)<jupyter_output><empty_output><jupyter_text>Optimizers and Learning Rate Schedulers The last setup step is to create an optimizer and a learning rate scheduler. We will use PyTorch's AdamW optimizer and linear learning rate scheduler since these are typically used to fine-tune BERT on tasks such as SST-2.<jupyter_code>from torch.optim import AdamW
from torch.optim.lr_scheduler import LinearLR

optimizer = AdamW(params=composer_model.parameters(), lr=3e-5, betas=(0.9, 0.98), eps=1e-6, weight_decay=3e-6)
linear_lr_decay = LinearLR(optimizer, start_factor=1.0, end_factor=0, total_iters=150)<jupyter_output><empty_output><jupyter_text>Composer Trainer We will now specify a composer `Trainer` object and run our training! `Trainer` has many arguments that are described in our [documentation](https://docs.mosaicml.com/en/stable/trainer.htmlcomposer.Trainer), but let's discuss the less obvious arguments used below:- `max_duration` - a string specifying how long to train either in terms of batches (e.g. '10ba' is 10 batches) or epochs (e.g. '1ep' is 1 epoch).- `schedulers` - a list of PyTorch learning rate schedulers that will be composed together.- `device` - specifies if the training will be done on CPU or GPU by using 'cpu' or 'gpu', respectively.- `train_subset_num_batches` - specifies the number of training batches to use for each epoch. This is not a necessary argument, but is useful for quickly testing code.- `precision` - whether to do the training in full precision 'fp32' or mixed precision 'amp'. Mixed precision provides almost 2x speedup in training time on certain hardware. If you get a P100, try `precision='amp'`!- `seed` - sets the random seed for the training run, so the results are reproducible!<jupyter_code>from composer import Trainer

# Create Trainer Object
trainer = Trainer(model=composer_model, 
                  train_dataloader=train_dataspec,
                  eval_dataloader=eval_dataspec,
                  max_duration="1ep",
                  optimizers=optimizer,
                  schedulers=[linear_lr_decay],
                  device='gpu',
                  train_subset_num_batches=150,
                  precision='fp32',
                  seed=17
                  )
# Start training
trainer.fit()<jupyter_output><empty_output><jupyter_text>Visualizing Results Our model reaches almost 86% accuracy with only 100 iterations of training! Let's visualize a few samples from the validation set to see how our model performs.<jupyter_code>import torch
eval_batch = next(iter(eval_dataloader))

# Move batch to gpu
eval_batch = {k: v.cuda() for k, v in eval_batch.items()}
with torch.no_grad():
  predictions = composer_model(eval_batch)["logits"].argmax(dim=1)

# Visualize only 5 samples
predictions = predictions[:6]

label = ['negative', 'positive']
for i, prediction in enumerate(predictions[:6]):
  sentence = sst2_dataset["validation"][i]["sentence"]
  correct_label = label[sst2_dataset["validation"][i]["label"]]
  prediction_label = label[prediction]
  print(f"Sample: {sentence}")
  print(f"Label: {correct_label}")
  print(f"Prediction: {prediction_label}")
  print()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08354591836734694, 'fraction_numerical': 0.009183673469387756, 'mean_word_length': 4.479385045422782, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 29, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '61679', 'license': 'Apache-2.0', 'n_tokens_mistral': 2318, 'n_tokens_neox': 2209, 'n_words': 722, 'path': 'notebooks/nlp_notebook_tutorial.ipynb', 'repo_name': 'dblalock/composer', 'size': '7840'}