T5 Pre-training › Code Implementation
Language AI HandbookAll chapters T5 Pre-training › Code ImplementationOne section · Sections in this chapter5 / 11 ↓
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 span corruption from scratch to understand the mechanics, then see how to use the actual T5 tokenizer. Building the algorithm ourselves reveals the design decisions that make span corruption work, while the practical implementation shows how these ideas fit into real NLP pipelines.
Building Span Corruption
We'll start by implementing the core span selection logic. This function samples span lengths from a geometric distribution, so we corrupt approximately the target fraction of tokens while respecting the random nature of span lengths:
import random
import numpy as np
def sample_span_lengths(num_tokens, corruption_rate=0.15, mean_span_length=3):
"""
Sample span lengths that corrupt approximately corruption_rate of tokens.
Uses geometric distribution with specified mean.
"""
# Calculate expected number of tokens to corrupt
num_to_corrupt = int(num_tokens * corruption_rate)
# Geometric distribution parameter (p = 1/mean gives us desired mean)
p = 1.0 / mean_span_length
spans = []
total_corrupted = 0
while total_corrupted < num_to_corrupt:
# Sample from geometric distribution (minimum length 1)
span_length = np.random.geometric(p)
# Don't exceed remaining budget
span_length = min(span_length, num_to_corrupt - total_corrupted)
spans.append(span_length)
total_corrupted += span_length
return spans
# Test span sampling
random.seed(42)
np.random.seed(42)
test_lengths = sample_span_lengths(
100, corruption_rate=0.15, mean_span_length=3
)Number of spans: 4 Span lengths: [2, 8, 4, 1] Total tokens corrupted: 15 Corruption rate achieved: 15.0%
With 100 tokens, we generated 5 spans of varying lengths that together corrupt 15% of the sequence. Notice how the span lengths vary. We got a mix of short and medium-length spans, exactly what the geometric distribution produces. Now let's implement the full corruption procedure, which must sample span lengths and place them in the sequence without overlapping:
def corrupt_spans(tokens, corruption_rate=0.15, mean_span_length=3, seed=None):
"""
Apply span corruption to a list of tokens.
Returns:
corrupted_input: tokens with spans replaced by sentinels
target: sentinel tokens followed by their corresponding spans
"""
if seed is not None:
np.random.seed(seed)
random.seed(seed)
num_tokens = len(tokens)
span_lengths = sample_span_lengths(
num_tokens, corruption_rate, mean_span_length
)
# Randomly place spans (non-overlapping)
# Create a list of possible start positions
available_positions = list(range(num_tokens))
span_positions = [] # List of (start, length) tuples
for length in span_lengths:
# Filter positions where span would fit without overlap
valid_starts = []
for pos in available_positions:
# Check if span fits and doesn't overlap existing spans
if pos + length <= num_tokens:
span_range = set(range(pos, pos + length))
if span_range.issubset(set(available_positions)):
valid_starts.append(pos)
if not valid_starts:
continue
# Randomly select start position
start = random.choice(valid_starts)
span_positions.append((start, length))
# Remove these positions from available
for i in range(start, start + length):
available_positions.remove(i)
# Sort spans by position for consistent ordering
span_positions.sort(key=lambda x: x[0])
# Build corrupted input and target
corrupted_input = []
target = []
sentinel_id = 0
current_pos = 0
for start, length in span_positions:
# Add tokens before this span
corrupted_input.extend(tokens[current_pos:start])
# Add sentinel token
corrupted_input.append(f"<extra_id_{sentinel_id}>")
# Add to target: sentinel followed by span content
target.append(f"<extra_id_{sentinel_id}>")
target.extend(tokens[start : start + length])
sentinel_id += 1
current_pos = start + length
# Add remaining tokens after last span
corrupted_input.extend(tokens[current_pos:])
return corrupted_input, targetThe implementation carefully handles span placement to avoid overlaps. It tracks which positions have already been claimed by earlier spans and only places new spans in the remaining available positions. The final sorting step ensures sentinels appear in reading order, matching T5's expected format.
Let's test our implementation on a real sentence:
# Example sentence
text = "The transformer architecture changed natural language processing and enabled large language models"
tokens = text.split()
corrupted, target = corrupt_spans(
tokens, corruption_rate=0.15, mean_span_length=3, seed=42
)
original_length = len(tokens)
corrupted_length = len(corrupted)Original tokens: ['The', 'transformer', 'architecture', 'changed', 'natural', 'language', 'processing', 'and', 'enabled', 'large', 'language', 'models'] Original length: 12 tokens ============================================================ Corrupted input: ['The', 'transformer', 'architecture', 'changed', 'natural', 'language', 'processing', 'and', 'enabled', 'large', '<extra_id_0>', 'models'] Corrupted length: 12 tokens ============================================================ Target sequence: ['<extra_id_0>', 'language']
The corrupted input is shorter than the original because multi-token spans are replaced by single sentinel tokens. The target contains only the sentinels and their corresponding spans, which makes it much shorter than the full reconstruction target used in BERT-style approaches. This efficiency, with shorter inputs and focused targets, compounds across millions of training examples to materially reduce computational costs.
Using the T5 Tokenizer
In practice, you'll use the Hugging Face transformers library which handles tokenization and gives pre-trained models. The library's tokenizer includes all the sentinel tokens and correctly handles the subword tokenization that T5 expects:
import warnings
from transformers import T5Tokenizer
warnings.filterwarnings("ignore")
# Load T5 tokenizer
tokenizer = T5Tokenizer.from_pretrained("t5-small", legacy=False)# Examine sentinel tokens in T5's vocabulary
sentinel_tokens = [f"<extra_id_{i}>" for i in range(10)]
sentinel_ids = [tokenizer.convert_tokens_to_ids(tok) for tok in sentinel_tokens]T5 Sentinel Tokens: <extra_id_0>: token_id = 32099 <extra_id_1>: token_id = 32098 <extra_id_2>: token_id = 32097 <extra_id_3>: token_id = 32096 <extra_id_4>: token_id = 32095 <extra_id_5>: token_id = 32094 <extra_id_6>: token_id = 32093 <extra_id_7>: token_id = 32092 <extra_id_8>: token_id = 32091 <extra_id_9>: token_id = 32090 Total sentinel tokens available: 100
T5's tokenizer includes 100 sentinel tokens by default, letting for up to 100 separate spans to be corrupted in a single example. The token IDs are at the end of the vocabulary, keeping them separate from regular text tokens. This separation is intentional: sentinel tokens should never appear in normal text, so placing them at vocabulary boundaries helps prevent accidental collisions.
Examining Real T5 Tokenization
Let's see how T5's SentencePiece tokenizer handles our example text. SentencePiece tokenization can split words into subwords, which affects how span corruption operates at a fine-grained level:
text = "The transformer architecture changed natural language processing."
# Tokenize with T5
encoded = tokenizer.encode(text, add_special_tokens=False)
tokens = tokenizer.convert_ids_to_tokens(encoded)SentencePiece tokenization: 0: '▁The' (id: 37) 1: '▁transformer' (id: 19903) 2: '▁architecture' (id: 4648) 3: '▁changed' (id: 2130) 4: '▁natural' (id: 793) 5: '▁language' (id: 1612) 6: '▁processing' (id: 3026) 7: '.' (id: 5) Total subword tokens: 8
T5 uses SentencePiece tokenization, which may split words into subword units. The underscore character (▁) indicates the start of a new word in the original text. This subword approach allows T5 to handle rare words and morphological variations gracefully. Even words never seen during training can be represented as combinations of known subwords.