Whole Word Masking › Empirical Comparison: WWM vs Random Masking

import torch
from transformers import BertForMaskedLM

# Load a pre-trained BERT model
model = BertForMaskedLM.from_pretrained("bert-base-uncased")
model.eval()


def get_mask_predictions(sentence, mask_positions, tokenizer, model):
    """Get model predictions for masked positions."""
    tokens = tokenizer.tokenize(sentence)
    tokens_with_special = ["[CLS]"] + tokens + ["[SEP]"]

    # Adjust mask positions for [CLS]
    adjusted_positions = [p + 1 for p in mask_positions]

    # Apply masking
    masked_tokens = tokens_with_special.copy()
    for pos in adjusted_positions:
        masked_tokens[pos] = "[MASK]"

    # Get predictions
    input_ids = tokenizer.convert_tokens_to_ids(masked_tokens)
    input_tensor = torch.tensor([input_ids])

    with torch.no_grad():
        outputs = model(input_tensor)
        predictions = outputs.logits[0]

    results = []
    for orig_pos, adj_pos in zip(mask_positions, adjusted_positions):
        probs = torch.softmax(predictions[adj_pos], dim=0)
        true_token_id = tokenizer.convert_tokens_to_ids([tokens[orig_pos]])[0]
        true_token_prob = probs[true_token_id].item()

        top5_ids = torch.topk(probs, 5).indices.tolist()
        top5_tokens = tokenizer.convert_ids_to_tokens(top5_ids)

        results.append(
            {
                "position": orig_pos,
                "true_token": tokens[orig_pos],
                "true_prob": true_token_prob,
                "top5": top5_tokens,
            }
        )

    return results