1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4from torch.utils.data import Dataset, DataLoader
5import tiktoken  # OpenAI's BPE tokenizer
6
7device = 'cuda' if torch.cuda.is_available() else 'cpu'
8
9config = {
10    'vocab_size': 50257,    # GPT-2 vocabulary size
11    'd_model': 384,         # Embedding dimension
12    'n_layers': 6,          # Number of transformer blocks
13    'n_heads': 6,           # Number of attention heads
14    'max_seq_len': 256,     # Maximum sequence length
15    'dropout': 0.1,

16    'batch_size': 32,
17    'learning_rate': 3e-4,
18    'num_epochs': 10
19}

1tokenizer = tiktoken.get_encoding("gpt2")
2
3text = "Hello, how are you doing today?"
4tokens = tokenizer.encode(text)
5# [15496, 11, 703, 389, 345, 1804, 1909, 30]
6
7# Decode back to text
8decoded = tokenizer.decode(tokens)  # "Hello, how are you doing today?"
9
10# Inspect individual tokens
11for tid in tokens:
12    print(f"  {tid}: '{tokenizer.decode([tid])}'")

1class TextDataset(Dataset):
2    def __init__(self, text_file, tokenizer, max_seq_len):
3        with open(text_file, 'r', encoding='utf-8') as f:
4            text = f.read()
5        self.tokens = tokenizer.encode(text)
6        self.max_seq_len = max_seq_len
7
8    def __len__(self):
9        return len(self.tokens) - self.max_seq_len
10
11    def __getitem__(self, idx):
12        chunk = self.tokens[idx : idx + self.max_seq_len + 1]
13        x = torch.tensor(chunk[:-1], dtype=torch.long)  # Input

14        y = torch.tensor(chunk[1:],  dtype=torch.long)  # Target
15        return x, y
16
17dataset = TextDataset('shakespeare.txt', tokenizer, config['max_seq_len'])
18dataloader = DataLoader(dataset, batch_size=config['batch_size'], shuffle=True)

1class Embeddings(nn.Module):
2    def __init__(self, vocab_size, d_model, max_seq_len, dropout):
3        super().__init__()
4        self.token_embed = nn.Embedding(vocab_size, d_model)
5        self.pos_embed = nn.Embedding(max_seq_len, d_model)
6        self.dropout = nn.Dropout(dropout)
7
8    def forward(self, x):
9        seq_len = x.size(1)
10        tok_emb = self.token_embed(x)                        # (B, T, D)
11        pos_emb = self.pos_embed(torch.arange(seq_len, device=x.device))  # (T, D)
12        return self.dropout(tok_emb + pos_emb)               # Broadcasting adds positions

1class MultiHeadAttention(nn.Module):
2    def __init__(self, d_model, n_heads, dropout):
3        super().__init__()
4        assert d_model % n_heads == 0
5        self.n_heads = n_heads
6        self.head_dim = d_model // n_heads
7
8        self.q_linear = nn.Linear(d_model, d_model)
9        self.k_linear = nn.Linear(d_model, d_model)
10        self.v_linear = nn.Linear(d_model, d_model)
11        self.out_linear = nn.Linear(d_model, d_model)
12        self.dropout = nn.Dropout(dropout)

13
14    def forward(self, x, mask=None):
15        B, T, D = x.shape
16        Q = self.q_linear(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
17        K = self.k_linear(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
18        V = self.v_linear(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
19        # Q, K, V shape: (B, n_heads, T, head_dim)

1        # Scaled dot-product attention
2        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
3
4        if mask is not None:
5            scores = scores.masked_fill(mask == 0, float('-inf'))
6
7        attn_weights = F.softmax(scores, dim=-1)
8        attn_weights = self.dropout(attn_weights)
9
10        # Apply attention to values and concatenate heads
11        out = torch.matmul(attn_weights, V)                  # (B, n_heads, T, head_dim)
12        out = out.transpose(1, 2).contiguous().view(B, T, D) # (B, T, D)
13        return self.out_linear(out)

1class FeedForward(nn.Module):
2    def __init__(self, d_model, dropout):
3        super().__init__()
4        self.net = nn.Sequential(
5            nn.Linear(d_model, 4 * d_model),
6            nn.GELU(),                        # GPT uses GELU, not ReLU
7            nn.Linear(4 * d_model, d_model),
8            nn.Dropout(dropout)
9        )
10    def forward(self, x):
11        return self.net(x)

12
13class TransformerBlock(nn.Module):
14    def __init__(self, d_model, n_heads, dropout):
15        super().__init__()
16        self.ln1 = nn.LayerNorm(d_model)
17        self.attention = MultiHeadAttention(d_model, n_heads, dropout)
18        self.ln2 = nn.LayerNorm(d_model)
19        self.ffn = FeedForward(d_model, dropout)
20
21    def forward(self, x, mask):
22        x = x + self.attention(self.ln1(x), mask)   # Pre-norm + residual

23        x = x + self.ffn(self.ln2(x))               # Pre-norm + residual
24        return x

1class GPT(nn.Module):
2    def __init__(self, vocab_size, d_model, n_layers, n_heads, max_seq_len, dropout):
3        super().__init__()
4        self.max_seq_len = max_seq_len
5        self.embeddings = Embeddings(vocab_size, d_model, max_seq_len, dropout)
6        self.blocks = nn.ModuleList([
7            TransformerBlock(d_model, n_heads, dropout) for _ in range(n_layers)
8        ])
9        self.ln_f = nn.LayerNorm(d_model)
10        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
11        self.apply(self._init_weights)
12

13    def _init_weights(self, module):
14        if isinstance(module, nn.Linear):
15            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
16            if module.bias is not None:
17                torch.nn.init.zeros_(module.bias)
18        elif isinstance(module, nn.Embedding):
19            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

1    def forward(self, x, targets=None):
2        seq_len = x.size(1)
3        mask = create_causal_mask(seq_len, x.device)
4
5        x = self.embeddings(x)
6        for block in self.blocks:
7            x = block(x, mask)
8        x = self.ln_f(x)
9        logits = self.lm_head(x)          # (batch, seq_len, vocab_size)
10
11        loss = None
12        if targets is not None:
13            loss = F.cross_entropy(
14                logits.view(-1, logits.size(-1)),
15                targets.view(-1)
16            )

17        return logits, loss

1model = GPT(
2    config['vocab_size'], config['d_model'], config['n_layers'],
3    config['n_heads'], config['max_seq_len'], config['dropout']
4).to(device)
5
6optimizer = torch.optim.AdamW(
7    model.parameters(), lr=config['learning_rate'],
8    betas=(0.9, 0.95), weight_decay=0.1
9)
10
11model.train()
12for epoch in range(config['num_epochs']):
13    total_loss = 0
14    for x, y in dataloader:

15        x, y = x.to(device), y.to(device)
16        logits, loss = model(x, targets=y)
17        optimizer.zero_grad()
18        loss.backward()
19        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
20        optimizer.step()
21        total_loss += loss.item()
22    print(f"Epoch {epoch+1}, Loss: {total_loss / len(dataloader):.4f}")

1class GPT(nn.Module):
2    def __init__(self, vocab_size, d_model, n_layers, n_heads, max_seq_len, dropout):
3        super().__init__()
4        self.embeddings = Embeddings(vocab_size, d_model, max_seq_len, dropout)
5        # ... blocks, ln_f as before ...
6        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
7
8        # Tie weights: lm_head uses same matrix as token embeddings
9        self.lm_head.weight = self.embeddings.token_embed.weight
10        # Saves ~20% parameters for large vocabularies!

1accumulation_steps = 8  # Simulate 8x larger batch
2optimizer.zero_grad()
3
4for i, (x, y) in enumerate(dataloader):
5    x, y = x.to(device), y.to(device)
6    logits, loss = model(x, targets=y)
7    loss = loss / accumulation_steps  # Normalize
8    loss.backward()
9
10    if (i + 1) % accumulation_steps == 0:
11        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
12        optimizer.step()
13        optimizer.zero_grad()

1import math
2
3def get_lr(step, warmup_steps=1000, max_steps=50000, max_lr=3e-4, min_lr=3e-5):
4    if step < warmup_steps:
5        return max_lr * step / warmup_steps          # Linear warmup
6    decay_ratio = (step - warmup_steps) / (max_steps - warmup_steps)
7    return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * decay_ratio))
8
9# Usage in training loop
10for step in range(max_steps):
11    lr = get_lr(step)
12    for param_group in optimizer.param_groups:
13        param_group['lr'] = lr
14    # ... training step ...

1scaler = torch.amp.GradScaler()
2
3for x, y in dataloader:
4    x, y = x.to(device), y.to(device)
5
6    with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
7        logits, loss = model(x, targets=y)
8
9    scaler.scale(loss).backward()
10    scaler.unscale_(optimizer)
11    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
12    scaler.step(optimizer)
13    scaler.update()
14    optimizer.zero_grad()

1@torch.no_grad()
2def generate_greedy(model, tokenizer, prompt, max_new_tokens=50):
3    model.eval()
4    tokens = tokenizer.encode(prompt)
5    x = torch.tensor([tokens], dtype=torch.long, device=device)
6
7    for _ in range(max_new_tokens):
8        x_crop = x[:, -model.max_seq_len:]
9        logits, _ = model(x_crop)
10        logits = logits[:, -1, :]                # Last position only
11        next_token = torch.argmax(logits, dim=-1, keepdim=True)
12        x = torch.cat([x, next_token], dim=1)
13
14    return tokenizer.decode(x[0].tolist())

1def sample_next_token(logits, temperature=1.0, top_k=None, top_p=None):
2    logits = logits / temperature
3
4    if top_k is not None:
5        top_k = min(top_k, logits.size(-1))
6        threshold = torch.topk(logits, top_k)[0][..., -1, None]
7        logits[logits < threshold] = float('-inf')
8
9    if top_p is not None:
10        sorted_logits, sorted_idx = torch.sort(logits, descending=True)
11        cumprobs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
12        remove = cumprobs > top_p
13        remove[..., 1:] = remove[..., :-1].clone()
14        remove[..., 0] = False
15        logits[sorted_idx[remove]] = float('-inf')
16

17    probs = F.softmax(logits, dim=-1)
18    return torch.multinomial(probs, num_samples=1)

1@torch.no_grad()
2def generate(model, tokenizer, prompt, max_new_tokens=100,
3             temperature=0.8, top_k=40, top_p=0.9):
4    model.eval()
5    tokens = tokenizer.encode(prompt)
6    x = torch.tensor([tokens], dtype=torch.long, device=device)
7
8    for _ in range(max_new_tokens):
9        x_crop = x[:, -model.max_seq_len:]
10        logits, _ = model(x_crop)
11        next_token = sample_next_token(
12            logits[0, -1, :], temperature=temperature, top_k=top_k, top_p=top_p
13        )
14        x = torch.cat([x, next_token.unsqueeze(0)], dim=1)

15
16    return tokenizer.decode(x[0].tolist())
17
18# Generate text
19print(generate(model, tokenizer, "Once upon a time"))