Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

BERT Pre-training › Putting It All Together

Michael BrenndoerferUpdated July 19, 20252 min read

Language AI HandbookAll chapters

Part 24 of 67BERT and Variants

Chapter 170BERT Pre-training › Putting It All Together

Reading section 8 of 12 · Putting It All Together · 2 minSection 8 of 12 · 2 min

Section 8 of 12 · 2 min read

BERT Pre-training › Putting It All TogetherOne section · Sections in this chapter8 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Let's implement a simplified but complete pre-training loop that demonstrates all the components:

In[38]:

Code

from torch.utils.data import Dataset


class BertPretrainingDataset(Dataset):
    """Dataset for BERT pre-training."""

    def __init__(self, documents, tokenizer, max_length=128, mask_prob=0.15):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.mask_prob = mask_prob

        # Create sentence pairs
        self.pairs = []
        all_sentences = [s for doc in documents for s in doc]

        for doc in documents:
            for i in range(len(doc) - 1):
                # 50% positive, 50% negative
                if random.random() < 0.5:
                    self.pairs.append((doc[i], doc[i + 1], True))
                else:
                    random_sent = random.choice(all_sentences)
                    self.pairs.append((doc[i], random_sent, False))

    def __len__(self):
        return len(self.pairs)

    def __getitem__(self, idx):
        sent_a, sent_b, is_next = self.pairs[idx]

        # Tokenize
        encoding = self.tokenizer(
            sent_a,
            sent_b,
            padding="max_length",
            truncation=True,
            max_length=self.max_length,
            return_tensors="pt",
        )

        input_ids = encoding["input_ids"].squeeze()
        attention_mask = encoding["attention_mask"].squeeze()
        token_type_ids = encoding["token_type_ids"].squeeze()

        # Apply MLM masking
        masked_ids, mlm_labels = apply_bert_masking(
            input_ids, self.tokenizer, self.mask_prob
        )

        return {
            "input_ids": masked_ids,
            "attention_mask": attention_mask,
            "token_type_ids": token_type_ids,
            "mlm_labels": mlm_labels,
            "nsp_labels": torch.tensor(1 if is_next else 0),
        }

In[39]:

Code

from torch.utils.data import DataLoader

# Create a small dataset for demonstration
demo_documents = [
    [
        "The sun was setting behind the mountains.",
        "Golden light painted the clouds.",
        "Birds flew home for the evening.",
        "The temperature began to drop.",
    ],
    [
        "Machine learning algorithms learn from data.",
        "They identify patterns automatically.",
        "This enables predictions on new inputs.",
        "Deep learning uses neural networks.",
    ],
    [
        "The chef prepared a delicious meal.",
        "Fresh ingredients make the difference.",
        "She seasoned everything perfectly.",
        "The guests enjoyed every bite.",
    ],
]

dataset = BertPretrainingDataset(demo_documents, tokenizer, max_length=64)
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

Out[40]:

Console

Dataset size: 9 sentence pairs
Batch count: 3

Sample batch shapes:
  input_ids: torch.Size([64])
  attention_mask: torch.Size([64])
  token_type_ids: torch.Size([64])
  mlm_labels: torch.Size([64])
  nsp_labels: torch.Size([])

This simplified implementation captures the essential pre-training workflow. A production system would add:

Whole word masking instead of token-level masking
Dynamic masking (fresh masks each epoch)
Distributed training across multiple GPUs/TPUs
Gradient accumulation for larger effective batch sizes
Mixed precision training for memory efficiency
Checkpointing and resumption