Whole Word Masking › Implementation
Language AI HandbookAll chapters Whole Word Masking › ImplementationOne section · Sections in this chapter5 / 13 ↓
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 whole word masking step by step. We will start with the core logic for identifying word boundaries, then build the complete masking function.
Identifying Word Boundaries
The first task is grouping subword tokens into words. WordPiece tokenizers mark continuation tokens with a ## prefix. The grouping logic traverses the token list and creates a new group whenever it encounters a token without the ## prefix, which signals the start of a new word.
from transformers import BertTokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
def identify_word_groups(tokens):
"""
Group tokens into words based on WordPiece conventions.
Returns list of (start_idx, end_idx) tuples for each word.
"""
word_groups = []
current_start = 0
for i, token in enumerate(tokens):
# New word starts when token doesn't have ## prefix
# (and isn't the first token)
if i > 0 and not token.startswith("##"):
word_groups.append((current_start, i))
current_start = i
# Don't forget the last word
if current_start < len(tokens):
word_groups.append((current_start, len(tokens)))
return word_groupsTokens: ['the', 'transformation', 'was', 'und', '##enia', '##bly', 'remarkable', '.'] Word: ['the'] Word: ['transformation'] Word: ['was'] Word: ['und', '##enia', '##bly'] Word: ['remarkable'] Word: ['.']
The function correctly identifies word boundaries by detecting the ## prefix. "The" remains a single-token word, while "transformation" is grouped as ["transform", "##ation"] and "undeniably" as ["undeni", "##ably"]. Each tuple represents the start and end indices in the token list as a half-open interval (start inclusive, end exclusive), making it easy to slice and mask all subwords of a word together using standard Python indexing.
Notice that the function handles both single-token and multi-token words uniformly. A single-token word like "was" produces a group spanning a single index. A three-token word would produce a group spanning three indices. This uniformity simplifies the masking code that follows, since the same loop structure works regardless of word length.
Building the Masking Function
Now we implement the complete WWM function. We will use the greedy word selection approach (shuffle words, accumulate until we hit the token budget) and apply the 80-10-10 rule at the word level.
import random
def whole_word_masking(tokens, tokenizer, mask_prob=0.15, mask_token="[MASK]"):
"""
Apply whole word masking to a token sequence.
Args:
tokens: List of subword tokens
tokenizer: Tokenizer for vocabulary access
mask_prob: Probability of masking (default 15%)
mask_token: Token to use for masking
Returns:
masked_tokens: Tokens with WWM applied
labels: Original tokens for masked positions, -100 elsewhere
"""
# Get word boundaries
word_groups = identify_word_groups(tokens)
# Calculate target number of tokens to mask
num_tokens = len(tokens)
target_masked = int(num_tokens * mask_prob)
# Shuffle word groups and select until we hit target
shuffled_groups = word_groups.copy()
random.shuffle(shuffled_groups)
selected_groups = []
tokens_selected = 0
for group in shuffled_groups:
group_size = group[1] - group[0]
if tokens_selected + group_size <= target_masked + 2:
selected_groups.append(group)
tokens_selected += group_size
if tokens_selected >= target_masked:
break
# Apply masking with 80-10-10 rule
masked_tokens = tokens.copy()
labels = [-100] * num_tokens # -100 = ignore in loss
vocab = list(tokenizer.get_vocab().keys())
for start, end in selected_groups:
for i in range(start, end):
labels[i] = tokenizer.convert_tokens_to_ids([tokens[i]])[0]
rand = random.random()
if rand < 0.8:
masked_tokens[i] = mask_token
elif rand < 0.9:
masked_tokens[i] = random.choice(vocab)
# else: keep original (10% case)
return masked_tokens, labelsOriginal: ['the', 'quick', 'brown', 'fox', 'jumps', 'over', 'the', 'lazy', 'dog', '.'] Masked: ['the', 'quick', 'brown', 'fox', 'jumps', 'over', 'the', '[MASK]', 'dog', '.'] Label IDs: ['_', '_', '_', '_', '_', '_', '_', 13971, '_', '_']
The function masks entire words together. When "jumps" is selected, both the root and any continuation subwords receive the same treatment. The labels array stores the original token IDs for computing the loss, with -100 marking positions to ignore during loss computation. This -100 convention is standard in PyTorch's cross-entropy loss, which skips positions with this sentinel value automatically.
The masking loop applies the 80-10-10 rule independently to each subword token within a selected word. This means there is a small chance that different subwords of the same word get different treatments (one masked, another randomly replaced). The more principled approach applies one dice roll per word, but the difference in practice is negligible and the independent-per-token approach is what many library implementations use.
Comparing with Random Subword Masking
Let's implement random subword masking for comparison. This is the original BERT approach: select individual tokens uniformly at random, with no regard for word boundaries.
def random_subword_masking(tokens, tokenizer, mask_prob=0.15):
"""Standard random subword masking (original BERT)."""
masked_tokens = tokens.copy()
labels = [-100] * len(tokens)
vocab = list(tokenizer.get_vocab().keys())
for i, token in enumerate(tokens):
if random.random() < mask_prob:
labels[i] = tokenizer.convert_tokens_to_ids([token])[0]
rand = random.random()
if rand < 0.8:
masked_tokens[i] = "[MASK]"
elif rand < 0.9:
masked_tokens[i] = random.choice(vocab)
return masked_tokens, labelsOriginal tokens: ['transformation', '##al', 'leadership', 'inspire', '##s', 'organizational', 'change', '.'] Random Subword Masking: Masked: ['[MASK]', '##al', 'leadership', '[MASK]', '##s', 'organizational', 'change', '.'] Info leak: Visible subwords may hint at masked siblings Whole Word Masking: Masked: ['transformation', '##al', 'leadership', 'inspire', '##s', '[MASK]', 'change', '.'] No leak: Complete words are masked together
In random subword masking, you might see "transform" with "##ational" masked, revealing the word structure through the visible prefix. In WWM, the entire word is hidden, forcing the model to reason about what concept fits the context rather than simply completing a visible stem.
The output also reveals another practical difference: random subword masking creates unpredictable patterns where some multi-token words have both subwords masked (by chance), some have none masked, and some have a mix. WWM produces cleaner masking patterns aligned with linguistic units. This alignment makes the masked language modeling task more consistent across different word types and lengths.