🤗 and 3LC example on the IMDb dataset

Open in GitHub

This notebook demonstrates fine-tuning a pretrained DistilBERT model from transformers on the IMDb dataset, using the 3LC integrations with Trainer and datasets from Hugging Face. 3LC metrics are collected before and after one epoch of training.

The notebook covers:

• Creating a Table from a datasets dataset.

• Fine-tuning a pretrained transformers model on the IMDb dataset with TLCTrainer.

• Using a custom function for metrics collection.

Project Setup

EPOCHS = 1
TRAIN_BATCH_SIZE = 16
EVAL_BATCH_SIZE = 256
TRAIN_DATASET_NAME = "hf-imdb-train"
EVAL_DATASET_NAME = "hf-imdb-test"
TRANSIENT_DATA_PATH = "../transient_data"
DEVICE = None
PROJECT_NAME = "hf-imdb"
RUN_NAME = "Train DistilBERT on IMDB"
DESCRIPTION = "Example notebook for training a DistilBERT model on the IMDB dataset"
TLC_PUBLIC_EXAMPLES_DEVELOPER_MODE = True
INSTALL_DEPENDENCIES = False

%%capture
if INSTALL_DEPENDENCIES:
    %pip --quiet install torch --index-url https://download.pytorch.org/whl/cu118
    %pip --quiet install torchvision --index-url https://download.pytorch.org/whl/cu118
    %pip --quiet install accelerate
    %pip --quiet install scikit-learn
    %pip --quiet install 3lc[huggingface]

Imports

import os

import datasets
import evaluate
import numpy as np
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, DataCollatorWithPadding, TrainingArguments

import tlc

os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"  # Removing DistilBertTokenizerFast tokenizer warning

datasets.utils.logging.disable_progress_bar()

if DEVICE is None:
    if torch.cuda.is_available():
        DEVICE = "cuda"
    elif torch.backends.mps.is_available():
        DEVICE = "mps"
    else:
        DEVICE = "cpu"

Initialize a 3LC Run

We initialize a Run with a call to tlc.init, and add the configuration to the Run object.

config = {
    "epochs": EPOCHS,
    "train_batch_size": TRAIN_BATCH_SIZE,
    "eval_batch_size": EVAL_BATCH_SIZE,
}

run = tlc.init(
    project_name=PROJECT_NAME,
    run_name=RUN_NAME,
    description=DESCRIPTION,
    parameters=config,
    if_exists="overwrite",
)

With the 3LC integration, you can use tlc.Table.from_hugging_face() as a drop-in replacement for datasets.load_dataset() to create a tlc.Table. Notice .latest(), which gets the latest version of the 3LC dataset.

train_dataset = tlc.Table.from_hugging_face(
    "imdb",
    split="train",
    project_name=PROJECT_NAME,
    dataset_name=TRAIN_DATASET_NAME,
    description="IMDB train dataset",
    if_exists="overwrite",
).latest()

eval_dataset = tlc.Table.from_hugging_face(
    "imdb",
    split="test",
    project_name=PROJECT_NAME,
    dataset_name=EVAL_DATASET_NAME,
    description="IMDB test dataset",
    if_exists="overwrite",
).latest()

You can use the data produced by these Tables like you would with a 🤗 dataset.

train_dataset_hf = datasets.load_dataset("imdb", split="train")
train_dataset_hf[0]

train_dataset[0]

Table provides a method map to apply both preprocessing and on-the-fly transforms to your data before it is sent to the model.

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased", model_max_length=512)


def tokenize(sample):
    return {**sample, **tokenizer(sample["text"], truncation=True)}

train_tokenized = train_dataset.map(tokenize)
eval_tokenized = eval_dataset.map(tokenize)

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

id2label = {0: "neg", 1: "pos"}
label2id = {"neg": 0, "pos": 1}

# For demonstration purposes, we use the distilbert-base-uncased model with a different set of labels than
# it was trained on. As a result, there will be a warning about the inconsistency of the classifier and
# pre_classifier weights. This is expected and can be ignored.
model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2, id2label=id2label, label2id=label2id
)

Setup Metrics Collection

Computing metrics is done by implementing a function which returns per-sample metrics you would like to see in the 3LC Dashboard.

We keep the metrics function in Hugging Face to see the intermediate aggregate metrics.

For special metrics such as the predicted category we specify that we would like this to be shown as a CategoricalLabel.

accuracy = evaluate.load("accuracy")


def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=labels)


def compute_tlc_metrics(logits, labels):
    probabilities = torch.nn.functional.softmax(logits, dim=-1)

    predictions = logits.argmax(dim=-1)
    loss = torch.nn.functional.cross_entropy(logits, labels, reduction="none")
    confidence = probabilities.gather(dim=-1, index=predictions.unsqueeze(-1)).squeeze()

    return {
        "predicted": predictions,
        "loss": loss,
        "confidence": confidence,
    }


compute_tlc_metrics.column_schemas = {
    "predicted": tlc.CategoricalLabelSchema(
        display_name="Predicted Label", class_names=id2label.values(), display_importance=4005
    ),
    "loss": tlc.Schema(display_name="Loss", writable=False, value=tlc.Float32Value()),
    "confidence": tlc.Schema(display_name="Confidence", writable=False, value=tlc.Float32Value()),
}

# Add references to the input datasets used by the Run.
run.add_input_table(train_dataset)
run.add_input_table(eval_dataset)

Train the model with TLCTrainer

To perform model training, we replace the usual Trainer with TLCTrainer and provide the per-sample metrics collection function. We also specify that we would like to collect metrics prior to training.

from tlc.integration.hugging_face import TLCTrainer

training_args = TrainingArguments(
    output_dir=TRANSIENT_DATA_PATH,
    learning_rate=2e-5,
    per_device_train_batch_size=TRAIN_BATCH_SIZE,
    per_device_eval_batch_size=EVAL_BATCH_SIZE,
    num_train_epochs=EPOCHS,
    weight_decay=0.01,
    report_to="none",  # Disable wandb logging
    eval_strategy="epoch",
    use_cpu=DEVICE == "cpu",
    disable_tqdm=True,
)

trainer = TLCTrainer(
    model=model,
    args=training_args,
    train_dataset=train_tokenized,
    eval_dataset=eval_tokenized,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_hf_metrics=compute_metrics,
    compute_tlc_metrics=compute_tlc_metrics,
    compute_tlc_metrics_on_train_begin=True,
)

trainer.train()