Data Preparation and Tokenization

Welcome back to Bringing Transformers to Life: Training & Inference! Excellent progress so far. In our first lesson, you successfully assembled a complete Transformer model by integrating all the architectural components we've built throughout this learning journey. You now have a powerful, production-ready model capable of handling sequence-to-sequence tasks.

However, even the most sophisticated model is only as good as the data it receives. In this second lesson, we shift our focus to a foundational yet critical aspect: data preparation and tokenization. While our Transformer model is ready to process sequences, we need to bridge the gap between raw text and the numerical tensors that neural networks understand. This lesson will guide you through building a complete data preprocessing pipeline, from raw sentences to properly batched, padded tensors ready for training. We'll explore vocabulary building, special token handling, and efficient data loading strategies that ensure your Transformer receives well-structured input for optimal learning.

Before diving into implementation, let's establish the complete data pipeline for sequence-to-sequence learning. Unlike simple classification tasks, seq2seq models require careful coordination between source and target sequences, each potentially having different vocabularies and lengths.

The pipeline consists of several interconnected stages: tokenization converts raw text into discrete tokens (words or characters), vocabulary building creates mappings between tokens and numerical indices, special token integration handles sequence boundaries and unknown words, and dynamic batching efficiently groups variable-length sequences. Each stage must handle the unique challenges of seq2seq tasks, such as different source and target languages, varying sentence lengths within batches, and the need for teacher forcing during training. Understanding this pipeline is crucial because any inefficiency or error propagates through training, potentially degrading model performance or causing training instability.

Our vocabulary system forms the backbone of text-to-tensor conversion. Let's examine how we create a robust vocabulary that handles both known and unknown tokens:

class Vocabulary:
    def __init__(self):
        self.token2idx = {'<PAD>': 0, '<SOS>': 1, '<EOS>': 2, '<UNK>': 3}
        self.idx2token = {0: '<PAD>', 1: '<SOS>', 2: '<EOS>', 3: '<UNK>'}
        self.size = 4
    
    def add_token(self, token):
        if token not in self.token2idx:
            self.token2idx[token] = self.size
            self.idx2token[self.size] = token
            self.size += 1
    
    def build_vocab(self, sentences, min_freq=1):
        """Build vocabulary from list of sentences"""
        counter = Counter()
        for sentence in sentences:
            counter.update(sentence.split())
        
        for token, freq in counter.items():
            if freq >= min_freq:
                self.add_token(token)

This vocabulary design incorporates four essential special tokens that enable robust sequence handling. The <PAD> token (index 0) allows batching sequences of different lengths, <SOS> marks sequence starts for decoder initialization, <EOS> signals sequence completion, and <UNK> gracefully handles out-of-vocabulary words. The build_vocab method uses frequency-based filtering through min_freq, preventing rare tokens that might cause overfitting while maintaining vocabulary manageability. Notice how we maintain bidirectional mappings: token2idx for encoding text into numbers and idx2token for decoding numbers back to text.

Let's complete our vocabulary with the encoding and decoding methods:

    def encode(self, sentence):
        """Convert sentence to token indices"""
        tokens = sentence.split()
        return [self.token2idx.get(token, self.token2idx['<UNK>']) 
                for token in tokens]
            
    def decode(self, indices):
        """Convert token indices back to sentence"""
        tokens = [self.idx2token[idx] for idx in indices if idx != self.token2idx['<PAD>']]
        return ' '.join(tokens).replace('<SOS>', '').replace('<EOS>', '').strip()

The encode method demonstrates graceful unknown token handling using the get method with <UNK> as the default. This ensures that even completely unseen words don't break the pipeline, instead being mapped to a special unknown token that the model can learn to handle appropriately. The decode method performs the inverse transformation while intelligently filtering out padding tokens and cleaning up special tokens that shouldn't appear in the final output. This asymmetric processing (adding special tokens during encoding, removing them during decoding) is fundamental to seq2seq training and inference workflows.

To demonstrate our pipeline, we'll create synthetic translation data and a PyTorch dataset class. We consider a very simple translation task: given a string consisting of multiple words, return another string in which the word order is reversed; for example, given the string "how are you", the model should return "you are how". While this task may seem trivial, it allows us to keep the training process complexity within reasonable time boundaries, while also building a scaffolding that can easily scale to real-world tasks (such as language translation) by simply changing the data source. In other words, this can be seen as the "Hello World" for a Transformer!

As first step, let's generate the data:

def create_synthetic_data(num_samples=1000, seed=42):
    """Create synthetic translation data (English to reversed English)"""
    random.seed(seed)
    templates = [
        "hello world", "good morning", "how are you", "nice to meet you",
        "thank you very much", "see you later", "have a nice day"
    ]
    
    src_sentences = []
    tgt_sentences = []
    
    for _ in range(num_samples):
        template = random.choice(templates)
        src_sentences.append(template)
        tgt_sentences.append(" ".join(template.split()[::-1]))
    
    return src_sentences, tgt_sentences

This synthetic data generator creates a controlled learning environment perfect for understanding data pipeline mechanics. By reversing word order ("hello world" becomes "world hello"), we create a deterministic translation task that's complex enough to require the full seq2seq architecture yet simple enough to debug and verify.

Now let's implement the dataset class:

class TranslationDataset(Dataset):
    def __init__(self, src_sentences, tgt_sentences, src_vocab, tgt_vocab, max_len=50):
        self.src_sentences = src_sentences
        self.tgt_sentences = tgt_sentences
        self.src_vocab = src_vocab
        self.tgt_vocab = tgt_vocab
        self.max_len = max_len
    
    def __getitem__(self, idx):
        src_sentence = self.src_sentences[idx]
        tgt_sentence = self.tgt_sentences[idx]
        
        # Encode sentences
        src_indices = self.src_vocab.encode(src_sentence)
        tgt_indices = self.tgt_vocab.encode('<SOS> ' + tgt_sentence) 
        tgt_output = self.tgt_vocab.encode(tgt_sentence + ' <EOS>')
        
        return {
            'src': torch.tensor(src_indices, dtype=torch.long),
            'tgt': torch.tensor(tgt_indices, dtype=torch.long),
            'tgt_output': torch.tensor(tgt_output, dtype=torch.long)
        }

Notice how the target sequence handling prepares for a training technique called teacher forcing: tgt_indices starts with <SOS> for the decoder input, while tgt_output ends with <EOS> for loss computation. This offset pattern is essential for how the model will be trained, allowing the decoder to receive the correct previous token as input while learning to predict the next token. While you don’t need to fully understand teacher forcing yet, know that this setup is standard for sequence-to-sequence training and will be explained in detail in the next lesson. The dataset returns three tensors per sample, creating the essential triplet structure needed for efficient seq2seq training.

Efficient batching requires dynamic padding since sequences within a batch typically have different lengths. Let's implement a custom collate function:

def collate_fn(batch):
    """Custom collate function for dynamic padding"""
    src_batch = [item['src'] for item in batch]
    tgt_batch = [item['tgt'] for item in batch]
    tgt_output_batch = [item['tgt_output'] for item in batch]
    
    # Pad sequences to max length in batch
    src_batch = torch.nn.utils.rnn.pad_sequence(src_batch, batch_first=True, padding_value=0)
    tgt_batch = torch.nn.utils.rnn.pad_sequence(tgt_batch, batch_first=True, padding_value=0)
    tgt_output_batch = torch.nn.utils.rnn.pad_sequence(tgt_output_batch, batch_first=True, padding_value=0)
    
    return {
        'src': src_batch,
        'tgt': tgt_batch,
        'tgt_output': tgt_output_batch
    }

This custom collate function implements adaptive padding that only pads to the maximum length within each batch, not across the entire dataset. This strategy significantly reduces computational waste compared to global padding while maintaining tensor uniformity required for efficient GPU processing. The pad_sequence function automatically handles the padding logic, using our designated padding value (0, corresponding to the <PAD> token). The batch_first=True parameter ensures compatibility with our Transformer implementation's expected input format.

Let's execute our complete pipeline and examine the outputs to verify correct functionality:

def main():
    """Test data preparation pipeline"""    
    # Create synthetic data
    src_sentences, tgt_sentences = create_synthetic_data(num_samples=500)
    
    print(f"Created {len(src_sentences)} sentence pairs")
    print("Sample data:")
    for i in range(3):
        print(f"  Source: {src_sentences[i]}")
        print(f"  Target: {tgt_sentences[i]}")
    
    # Build vocabularies
    src_vocab = Vocabulary()
    tgt_vocab = Vocabulary()
    
    src_vocab.build_vocab(src_sentences)
    tgt_vocab.build_vocab(tgt_sentences)
    
    print(f"Source vocabulary size: {src_vocab.size}")
    print(f"Target vocabulary size: {tgt_vocab.size}")

The initial output demonstrates our synthetic data generation and vocabulary building:

Created 500 sentence pairs
Sample data:
  Source: see you later
  Target: later you see
  Source: have a nice day
  Target: day nice a have
  Source: thank you very much
  Target: much very you thank
Source vocabulary size: 22
Target vocabulary size: 22

Both vocabularies reach identical sizes (22 tokens) because our synthetic task uses the same words in both source and target, just in different orders. Let's complete the validation by testing batch processing:

    # Create dataset and dataloader
    dataset = TranslationDataset(src_sentences, tgt_sentences, src_vocab, tgt_vocab)
    dataloader = DataLoader(dataset, batch_size=4, shuffle=True, collate_fn=collate_fn)
    
    # Test one batch
    batch = next(iter(dataloader))
    
    print(f"Batch source shape: {batch['src'].shape}")
    print(f"Batch target shape: {batch['tgt'].shape}")
    print(f"Batch target output shape: {batch['tgt_output'].shape}")
    
    # Decode sample
    sample_src = batch['src'][0]
    sample_tgt = batch['tgt_output'][0]
    
    decoded_src = src_vocab.decode(sample_src.tolist())
    decoded_tgt = tgt_vocab.decode(sample_tgt.tolist())
    
    print(f"Decoded source: '{decoded_src}'")
    print(f"Decoded target: '{decoded_tgt}'")

The final output confirms successful tensor creation and round-trip encoding/decoding:

Batch source shape: torch.Size([4, 4])
Batch target shape: torch.Size([4, 5])
Batch target output shape: torch.Size([4, 5])
Decoded source: 'nice to meet you'
Decoded target: 'you meet to nice'

Notice how the target sequences are longer (5 tokens) than the source sequences (4 tokens) due to the added <SOS> and <EOS> tokens. The successful decoding demonstrates that our vocabulary system correctly handles the complete encode-decode cycle, preserving semantic content while managing special tokens appropriately.

You've successfully built a comprehensive data preparation pipeline that transforms raw text into training-ready tensors for sequence-to-sequence learning. This pipeline elegantly handles vocabulary construction, special token management, dynamic padding, and efficient batching — all critical components for successful Transformer training. Your implementation demonstrates a sophisticated understanding of the challenges inherent in seq2seq data processing, from handling variable-length sequences to implementing teacher forcing preparation.

The robust foundation you've established will serve as the cornerstone for all subsequent training and inference work. In the upcoming practice exercises, you'll apply this pipeline to more complex scenarios and explore how different tokenization strategies impact model performance, deepening your expertise in the crucial intersection between raw data and neural network training.