Input:  "The cat"
Output: " sat"
Next:   " on"
→ "The cat sat on the mat"

def create_causal_mask(seq_len):
    return torch.triu(torch.ones(seq_len, seq_len), diagonal=1) == 0

Mask:
[[1, 0, 0, 0],
 [1, 1, 0, 0],
 [1, 1, 1, 0],
 [1, 1, 1, 1]]

import torch
import torch.nn as nn
import torch.nn.functional as F

class MiniGPT(nn.Module):
    def __init__(self, vocab_size=1000, n_embd=64, n_head=4, n_layer=4, max_seq=128, dropout=0.1):
        super().__init__()
        self.max_seq = max_seq
        self.n_embd = n_embd

        # Token + Position
        self.token_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(max_seq, n_embd)

        # Decoder blocks
        self.blocks = nn.ModuleList([
            TransformerBlock(n_embd, n_head, n_embd*4, dropout)
            for _ in range(n_layer)
        ])

        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)

        self.dropout = nn.Dropout(dropout)

        # Init
        self.apply(self._init_weights)

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

class TransformerBlock(nn.Module):
    def __init__(self, n_embd, n_head, n_ff, dropout):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalMultiHeadAttention(n_embd, n_head, dropout)
        self.ln2 = nn.LayerNorm(n_embd)
        self.ff = nn.Sequential(
            nn.Linear(n_embd, n_ff),
            nn.GELU(),
            nn.Linear(n_ff, n_embd),
            nn.Dropout(dropout)
        )

    def forward(self, x, cache=None):
        attn_out, new_cache = self.attn(self.ln1(x), cache)
        x = x + attn_out
        x = x + self.ff(self.ln2(x))
        return x, new_cache

class CausalMultiHeadAttention(nn.Module):
    def __init__(self, n_embd, n_head, dropout):
        super().__init__()
        self.n_head = n_head
        self.d_k = n_embd // n_head
        self.Wq = nn.Linear(n_embd, n_embd)
        self.Wk = nn.Linear(n_embd, n_embd)
        self.Wv = nn.Linear(n_embd, n_embd)
        self.Wo = nn.Linear(n_embd, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, cache=None):
        B, T, C = x.shape
        q = self.Wq(x).view(B, T, self.n_head, self.d_k).transpose(1, 2)
        k = self.Wk(x).view(B, T, self.n_head, self.d_k).transpose(1, 2)
        v = self.Wv(x).view(B, T, self.n_head, self.d_k).transpose(1, 2)

        # KV Cache
        if cache is not None:
            k_cache, v_cache = cache
            k = torch.cat([k_cache, k], dim=2)
            v = torch.cat([v_cache, v], dim=2)

        # Scaled dot-product
        att = (q @ k.transpose(-2, -1)) * (1.0 / (self.d_k ** 0.5))
        mask = torch.tril(torch.ones(T + (k.size(2) - T) if cache else T, 
                                      k.size(2), device=x.device))
        att = att.masked_fill(mask == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.dropout(att)
        y = att @ v

        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.Wo(y)

        new_cache = (k, v) if T == 1 else None  # Only cache during generation
        return y, new_cache

    def forward(self, idx, targets=None, cache=None):
        B, T = idx.shape
        assert T <= self.max_seq

        # Embeddings
        tok_emb = self.token_emb(idx)
        pos_emb = self.pos_emb(torch.arange(T, device=idx.device))
        x = self.dropout(tok_emb + pos_emb)

        # Forward through blocks
        new_caches = []
        for i, block in enumerate(self.blocks):
            cache_i = cache[i] if cache else None
            x, new_cache = block(x, cache_i)
            new_caches.append(new_cache)

        x = self.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))

        return logits, loss, new_caches

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        cache = [None] * len(self.blocks)
        for _ in range(max_new_tokens):
            logits, _, cache = self(idx, cache=cache)
            logits = logits[:, -1, :] / temperature

            if top_k:
                v, _ = torch.topk(logits, top_k)
                logits = logits.masked_fill(logits < v[:, [-1]], float('-inf'))

            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, idx_next], dim=1)

            if idx_next.item() == 0:  # EOS
                break
        return idx

# === FULL MINI-GPT (64-dim) ===
class MiniGPT(nn.Module):
    def __init__(self, vocab_size=50257, n_embd=64, n_head=4, n_layer=4, max_seq=128):
        super().__init__()
        self.max_seq = max_seq
        self.token_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(max_seq, n_embd)
        self.blocks = nn.ModuleList([TransformerBlock(n_embd, n_head, n_embd*4) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, std=0.02)

    def forward(self, idx, targets=None, cache=None):
        B, T = idx.shape
        x = self.token_emb(idx) + self.pos_emb(torch.arange(T, device=idx.device))
        new_cache = []
        for i, block in enumerate(self.blocks):
            c = cache[i] if cache else None
            x, nc = block(x, c)
            new_cache.append(nc)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) if targets is not None else None
        return logits, loss, new_cache

    @torch.no_grad()
    def generate(self, idx, max_new_tokens=50):
        cache = [None] * len(self.blocks)
        for _ in range(max_new_tokens):
            logits, _, cache = self(idx, cache=cache)
            next_token = torch.multinomial(F.softmax(logits[:, -1, :], dim=-1), 1)
            idx = torch.cat([idx, next_token], dim=1)
        return idx

# Download tiny shakespeare
!wget https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt -O tiny.txt

text = open('tiny.txt').read()
chars = sorted(list(set(text)))
vocab_size = len(chars)
stoi = {ch:i for i,ch in enumerate(chars)}
itos = {i:ch for i,ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

data = torch.tensor(encode(text), dtype=torch.long)
train_data = data[:int(0.9*len(data))]
val_data = data[int(0.9*len(data)):]

# Model
model = MiniGPT(vocab_size=vocab_size, n_embd=64, n_head=4, n_layer=4)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

# Train
for step in range(1000):
    xb = train_data[torch.randint(len(train_data)-32, (32,))]
    xb = xb.unfold(0, 32, 1).t().contiguous()[:, :-1]
    yb = xb[:, 1:]
    xb = xb[:, :-1]

    logits, loss, _ = model(xb, yb)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if step % 100 == 0:
        print(f"Step {step}, Loss: {loss.item():.4f}")

context = torch.tensor(encode("ROMEO:"), dtype=torch.long).unsqueeze(0)
generated = model.generate(context, max_new_tokens=200)
print(decode(generated[0].tolist()))

import time

model.eval()
context = torch.tensor(encode("To be or not to be"), dtype=torch.long).unsqueeze(0)

# Without cache
start = time.time()
for _ in range(50):
    model(context)
no_cache = time.time() - start

# With cache
cache = [None] * 4
start = time.time()
for _ in range(50):
    _, _, cache = model(context, cache=cache)
with_cache = time.time() - start

print(f"No cache: {no_cache:.3f}s, With cache: {with_cache:.3f}s, Speedup: {no_cache/with_cache:.1f}x")