1import random
2
3MASK, VOCAB = "[MASK]", ["meal", "song", "cake", "plan", "dish"] # toy vocab
4
5
6def corrupt(tokens: list[str], mask_rate: float = 0.15):
7 """BERT's 80/10/10 masking recipe. Returns (corrupted, targets)."""
8 corrupted, targets = [], {}
9 for i, tok in enumerate(tokens):
10 if random.random() < mask_rate:
11 targets[i] = tok # we must predict the ORIGINAL
12 roll = random.random()
13 if roll < 0.8:
14 corrupted.append(MASK) # 80%: hide it
15 elif roll < 0.9:
16 corrupted.append(random.choice(VOCAB)) # 10%: wrong word
17 else:
18 corrupted.append(tok) # 10%: leave it (stay suspicious!)
19 else:
20 corrupted.append(tok)
21 return corrupted, targets
22
23
24def training_step(encoder, tokens):
25 corrupted, targets = corrupt(tokens)
26
27 # The encoder sees the FULL corrupted sentence — both sides of
28 # every blank. No causal mask. This is the bidirectional part.
29 hidden = encoder(corrupted) # (seq_len, d_model)
30
31 loss = 0.0
32 for pos, original in targets.items():
33 # Predict the original token from its two-sided context
34 probs = softmax(hidden[pos] @ encoder.vocab_matrix.T)
35 loss += -log(probs[vocab_id(original)])
36 return loss / max(1, len(targets))
37
38
39sentence = "the chef cooked a delicious meal in the kitchen".split()
40# corrupt() might yield: the chef cooked a delicious [MASK] in the kitchen
41# ...and the model learns that both "delicious" AND "kitchen" point to "meal".