Causal Language Modeling › A Working Example: Training a Tiny CLM

import torch

# Sample text for training
text = """To be or not to be that is the question
Whether tis nobler in the mind to suffer
The slings and arrows of outrageous fortune"""

# Character-level tokenization
chars = sorted(set(text))
char_to_idx = {c: i for i, c in enumerate(chars)}
idx_to_char = {i: c for c, i in char_to_idx.items()}
vocab_size = len(chars)

# Encode the text
encoded = torch.tensor([char_to_idx[c] for c in text])

class TinyCausalLM(nn.Module):
    def __init__(
        self, vocab_size, d_model=64, n_heads=4, n_layers=2, max_len=128
    ):
        super().__init__()
        self.d_model = d_model

        # Token and position embeddings
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_len, d_model)

        # Transformer layers
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=d_model * 4,
            dropout=0.1,
            batch_first=True,
        )
        self.transformer = nn.TransformerEncoder(
            encoder_layer, num_layers=n_layers
        )

        # Output projection
        self.output = nn.Linear(d_model, vocab_size)

        # Register causal mask buffer
        self.register_buffer("causal_mask", None)

    def forward(self, x):
        batch_size, seq_len = x.shape

        # Create position indices
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0)

        # Combine embeddings
        h = self.token_emb(x) + self.pos_emb(positions)

        # Create causal mask if needed
        if self.causal_mask is None or self.causal_mask.size(0) < seq_len:
            mask = torch.triu(
                torch.ones(seq_len, seq_len, device=x.device), diagonal=1
            )
            self.causal_mask = mask.masked_fill(mask == 1, float("-inf"))

        # Apply transformer with causal masking
        h = self.transformer(
            h, mask=self.causal_mask[:seq_len, :seq_len], is_causal=True
        )

        # Project to vocabulary
        logits = self.output(h)

        return logits

def get_batch(data, batch_size=32, seq_len=32):
    """Get a random batch of sequences."""
    # Random starting positions
    starts = torch.randint(0, len(data) - seq_len - 1, (batch_size,))

    # Extract sequences
    x = torch.stack([data[s : s + seq_len] for s in starts])
    y = torch.stack([data[s + 1 : s + seq_len + 1] for s in starts])

    return x, y


# Training loop
model = TinyCausalLM(vocab_size)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

losses = []
for step in range(500):
    # Get batch
    x, y = get_batch(encoded, batch_size=16, seq_len=32)

    # Forward pass
    logits = model(x)
    loss = F.cross_entropy(logits.view(-1, vocab_size), y.view(-1))

    # Backward pass
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    losses.append(loss.item())

def generate(model, start_text, max_new_tokens=100, temperature=0.8):
    """Generate text autoregressively."""
    model.eval()

    # Encode the prompt
    tokens = [char_to_idx[c] for c in start_text]
    tokens = torch.tensor(tokens).unsqueeze(0)

    with torch.no_grad():
        for _ in range(max_new_tokens):
            # Get predictions for last position
            logits = model(tokens)
            next_logits = logits[0, -1, :] / temperature

            # Sample from distribution
            probs = F.softmax(next_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)

            # Append to sequence
            tokens = torch.cat([tokens, next_token.unsqueeze(0)], dim=1)

    # Decode
    generated = "".join([idx_to_char[t.item()] for t in tokens[0]])
    return generated