An Introduction to Transformer

import torch 
from torch import nn

batch_size = 32
seq_len = 6     # squence length

n_embd = 10     # embedding dim (d_e metioned above)

x = torch.randn(batch_size, seq_len, n_embd)

class SelfAttention(nn.Module):
    '''Implement a single head self-attention.'''

    def __init__(self, head_size):
        super().__init__()
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)


    def forward(self, x):
        B, T, d = x.shape   # x shape: [batch_size, max_len, n_embd]

        q = self.query(x)                       # [B, T, head_size]

        k = self.query(x)                       # [B, T, head_size]

        v = self.value(x)                       # [B, T, head_size]

        wei = q @ k.transpose(-2, -1) * d**0.5  # [B, T, T]

        wei = F.softmax(wei, dim=-1) 

        out = wei @ v                           # [B, T, head_size]

        return out


head_size = 8   # dimension of query, key and value (set to be the same by default)        

slf_attn = SelfAttention(head_size)        
print(slf_attn(x).shape)       # [batch_size, max_len, head_size]

import torch 
from torch import nn

batch_size = 32
seq_len = 6
n_embd = 10
x = torch.randn(batch_size, seq_len, n_embd)



class MultiHeadAttention(nn.Module):
    '''Implement multiple heads of self-attetition in parallel.'''
    def __init__(self, num_heads, single_head_size):
        super().__init__()
        self.heads = nn.ModuleList([SelfAttention(single_head_size) for _ in range(num_heads)])


    def forward(self, x):

        return torch.cat([h(x) for h in self.heads], dim=-1)


head_size = 8  
num_heads = 4
single_head_size = head_size // num_heads
multihead_attn = MultiHeadAttention(num_heads, single_head_size)        
print(multihead_attn(x).shape)       # [batch_size, max_len, single_head_size*num_heads]

import torch 
from torch import nn

batch_size = 32
seq_len = 6
n_embd = 10
x = torch.randn(batch_size, seq_len, n_embd)


class LayerNorm(nn.Module):
    "Implements LayerNorm in Transformer. For convenience, just use nn.LayerNorm()."
    def __init__(self, in_dim, eps=1e-6):
        super().__init__()

        self.a = nn.Parameter(torch.ones(in_dim))   # the 'n_embd' dimension

        self.b = nn.Parameter(torch.ones(in_dim))
        self.eps = eps

    def forward(self, x):
        # normalize over the 'n_embd' dimension.

        means = x.mean(dim=-1, keep_dim=True)
        stds = x.std(dim=-1, keep_dim=True)

        out = self.a * (x - means) / (stds + self.eps) + self.b

        return out


layer_norm = nn.LayerNorm(n_embd)
out = layer_norm(x)
print(out.shape)  

import torch 
from torch import nn

class PositionWiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, in_dim, hidden_dim):
        super().__init__()

        self.net = nn.Sequential(nn.Linear(in_dim, hidden_dim),  
                                 nn.ReLU(),
                                 nn.Linear(hidden_dim, in_dim),
                                 )


    def forward(self, x):

        return self.net(x)

import numpy as np
import torch 
from torch import nn

class PositionalEncoding(nn.Module):
    def __init__(self, n_embd, max_len=500):
        super(PositionalEncoding, self).__init__()  

        # Compute the positional encodings once in log space.

        pe = torch.zeros(max_len, n_embd)
        position = torch.arange(0, max_len).unsqueeze(dim=1)
        div_term = torch.exp(torch.arange(0, n_embd, 2) * -(np.log(10000) / d_model))

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(dim=0) # (1, max_len, d_model)

        self.register_buffer('pe', pe)

    def forward(self, embedding):
        "Output token embedding + positional encoding."
        # embedding shape: [batch_size, len, n_embd]

        x = embedding + self.pe[:, :embedding.size(1), :].requires_grad_(False)

        return x

import torch 
from torch import nn

class Embeddings(nn.Module):
    def __init__(self, vocab_size, n_embd):l
        super(Embeddings, self).__init__()
        self.embedding = nn.Embedding(vocab_size, n_embd)

    def forward(self, x):

        return self.embedding(x)

torch.manual_seed(123)
B, T, C = 4, 8, 32  # batch_size, seq_len, n_embd

x = torch.randn(B, T, C)

head_size = 16
key = nn.Linear(C, head_size, bias=False)
value = nn.Linear(C, head_size, bias=False)
query = nn.Linear(C, head_size, bias=False)

k = key(x)
q = query(x)
wei = q @ k.transpose(-2, -1)  # [B, T, C] @ [B, C, T] ---> [B, T, T]

# implement sequence mask

tril = torch.tril(torch.ones(T, T))
wei = wei.masked_fill(tril == 0, float('-inf'))  
wei = F.softmax(wei, dim=-1)

v = value(x)
out = wei @ v

print(out.shape)   # [4, 8, 16]

 import torch
 from torch import nn

 class LabelSmoothing(nn.Module):
    def __init__(self, vocab_size, smoothing=0.1):
        super().__init__()

        self.vocab_size = vocab_size
        self.smoothing = smoothing

    def forward(self, target):
        # target shape: [batch_size, seq_len, ] or [N, ]

        onehot = F.one_hot(target, self.vocab_size)

        return (1 - self.smoothing) * onehot + (self.smoothing / (self.vocab_size - 1)) * (1 - onehot)


vocab_size = 5
label_smooth = LabelSmoothing(vocab_size)
# target = torch.randint(0, 4, (3, 6))  # [batch_size, seq_len]

 target = torch.LongTensor([2, 1, 0, 3, 3])

print(label_smooth(target))        

# get all the unique characters

chars = sorted(list(set(data)))     # chars[0]: '\n',  chars[1]: ' '

vocab_size = len(chars)
print("all the unique characters:", ''.join(chars))
print(f"vocab size: {vocab_size}")


# all the unique characters: 

# !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz

# vocab size: 65



# create a mapping from characters to integers

stoi = {ch:i for i, ch in enumerate(chars)}
itos = {i:ch for i, ch in enumerate(chars)}

def encode(string):
    '''encoder: take a string, output a list of integers'''
    return [stoi[s] for s in string]

def decode(ids):
    '''decoder: take a list of integers, output a string'''
    return ''.join(itos[i] for i in ids)


print(encode("hi there"))
print(decode(encode("hi there")))


# [46, 47, 1, 58, 46, 43, 56, 43]

# hi there

# create the train, validation split

n = int(0.9 * len(data))   # number of tokens

train_data = torch.as_tensor(encode(data[:n]), dtype=torch.long)
val_data = torch.as_tensor(encode(data[n:]), dtype=torch.long)
print(f"train has {len(train_data):,} tokens")
print(f"val has {len(val_data):,} tokens")


# train has 1,003,854 tokens

# val has 111,540 tokens

context_len = 8
x = train_data[:context_len]
y = train_data[1:context_len+1]
print(f"x: {x}")
print(f"y: {y}")
for t in range(context_len):
    context = x[:t+1]
    target = y[t]
    print(f"when input is {context}, the target is {target}.")



# x: tensor([18, 47, 56, 57, 58,  1, 15, 47])

# y: tensor([47, 56, 57, 58,  1, 15, 47, 58])

# when input is tensor([18]), the target is 47.

# when input is tensor([18, 47]), the target is 56.

# when input is tensor([18, 47, 56]), the target is 57.

# when input is tensor([18, 47, 56, 57]), the target is 58.

# when input is tensor([18, 47, 56, 57, 58]), the target is 1.

# when input is tensor([18, 47, 56, 57, 58,  1]), the target is 15.

# when input is tensor([18, 47, 56, 57, 58,  1, 15]), the target is 47.

# when input is tensor([18, 47, 56, 57, 58,  1, 15, 47]), the target is 58.

# create data loader

context_len = 8 
batch_size = 4


def get_batch(data, context_len, atch_size):
    '''generate a batch of data of inputs x and targets y'''
    idx = torch.randint(len(data) - context_len, (batch_size, ))

    x = torch.stack([data[i:i+context_len] for i in idx])
    y = torch.stack([data[i+1:i+context_len+1] for i in idx])

    return x, y



xb, yb = get_batch(train_data, context_len, batch_size)
print(xb.shape)  # [4, 8]

print(yb.shape)  # [4, 8]

class MultiHeadSelfAttention(nn.Module):
    def __init__(self, context_len, d_model, h, dropout=0.1):
        super().__init__()

        assert d_model % h == 0
        self.h = h
        self.d = d_model // h

        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.proj = nn.Linear(d_model, d_model)

        self.attn_dropout = nn.Dropout(p=dropout)
        self.dropout = nn.Dropout(p=dropout)

        self.attn = None
        self.register_buffer('mask', torch.tril(torch.ones(1, 1, context_len, context_len)).type(torch.int8))

    def forward(self, x):
        # x shape: [batch_size, len, d_model]

        # mask shape: [1, 1, len, len]

        b, t, d_model = x.shape  

        # pre-attention linear projection:    

        q, k, v = self.qkv(x).split(d_model, dim=-1)

        # separate different heads: [b, len, h*d] ==> [b, len, h, d]

        q = q.view(b, t, self.h, self.d)
        k = k.view(b, t, self.h, self.d)
        v = v.view(b, t, self.h, self.d)

        # transpose for attention dot product: [b, len, h, d] ==> [b, h, len, d]

        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        att = q @ k.transpose(-2, -1) / math.sqrt(self.d)
        att = att.masked_fill(self.mask[:, :, :t, :t] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        self.attn = att
        att = self.attn_dropout(att)   # Not mentioned in the paper but practically useful

        out = att @ v            # [b, h, len, d]

        # transpose to move the head dimension back: [b, h, len, d] ==> [b, len, h, d]

        # combine the last two dimensions to concatenate all the heads together: [b, len, h*d]

        out = out.transpose(1, 2).contiguous().view(b, t, d_model)
        out = self.proj(out)  # [b, h, len, d]

        del q, k ,v

        return self.dropout(out)


class PositionWiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, d_model, dropout=0.1):
        super().__init__()

        self.w1 = nn.Linear(d_model, 4 * d_model)  
        self.w2 = nn.Linear(4 * d_model, d_model)  
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x):
        # x shape: [b, len, d_model]

        x = self.w2(F.relu(self.w1(x)))         

        return self.dropout(x)



class DecoderBlock(nn.Module):
    def __init__(self, context_len, d_model, h, dropout):
        super(DecoderBlock, self).__init__()
        self.dec_attn = MultiHeadSelfAttention(context_len, d_model, h, dropout)
        self.feed_forward = PositionWiseFeedForward(d_model, dropout)
        self.layer_norm1 = nn.LayerNorm(d_model)
        self.layer_norm2 = nn.LayerNorm(d_model)

    def forward(self, x):
        # pre-LayerNorm

        x = x + self.dec_attn(self.layer_norm1(x))
        x = x + self.feed_forward(self.layer_norm2(x))

        return x


class Decoder(nn.Module):
    def __init__(self, context_len, d_model, h, N, dropout):
        super(Decoder, self).__init__()

        self.decoder = nn.Sequential(*[DecoderBlock(context_len, d_model, h, dropout) for _ in range(N)])  

        self.layer_norm = nn.LayerNorm(d_model)


    def forward(self, seq):

        dec_output = self.decoder(seq)

        return self.layer_norm(dec_output)




class GPT(nn.Module):
    def __init__(self, vocab_size, context_len, d_model, h, N, dropout):
        super().__init__()

        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.positional_embedding = nn.Embedding(context_len, d_model)   # Instead of using encoding as the original paper, embedding is used here
        self.decoder = Decoder(context_len, d_model, h, N, dropout)
        self.linear = nn.Linear(d_model, vocab_size)

        self.context_len = context_len

        # # initialize the weights

        # for p in self.parameters():

        #     if p.dim() > 1:

        #         nn.init.xavier_uniform_(p)

    def forward(self, seq):
        # seq shape: [batch_size, seq_len]

        b, t = seq.shape
        token_emb = self.token_embedding(seq)  # [b, t, d_model]

        pos_emb =self.positional_embedding(torch.arange(t)) # [t, d_model]

        x = token_emb + pos_emb   # [b, t, d_model]

        x = self.decoder(x)  # [b, t, d_model]

        logits = self.linear(x)   # [b, t, vocab_size]

        return logits



    @torch.no_grad()
    def generate(self, seq, max_len):
        # seq is [B, t] array of indices in the current context

        for _ in range(max_len):
            # crop seq to the last 'context_len' tokens, or the positional embedding will fail

            seq_crop = seq[:, -self.context_len:]
            logits = self(seq_crop)
            # focus on only the last time step

            logits = logits[:, -1, :]   # [B, vocab_size]

            probs = F.softmax(logits, dim=-1)
            # sample from the distribution

            idx_next = torch.multinomial(probs, num_samples=1)  # [B, 1]

            # append sampled index to the running sequence

            seq = torch.cat([seq, idx_next], dim=-1)  # [B, t+1]

        return seq

def compute_loss(model, x, y):
    # x, y shape:[b, t]
    logits = model(x)
    b, t, vocab_size = logits.shape
    logits = logits.contiguous().view(-1, vocab_size)
    y = y.contiguous().view(-1)
    loss = F.cross_entropy(logits, y)

    return loss


@torch.no_grad()
def eval_loss(model, eval_iters):
    train_loss = val_loss = 0

    model.eval()
    for _ in range(eval_iters):
        x, y = get_batch(train_data, context_len, batch_size)
        train_loss += compute_loss(model, x, y).item()

        x, y = get_batch(val_data, context_len, batch_size)
        val_loss += compute_loss(model, x, y).item()        
    model.train()

    return train_loss / eval_iters, val_loss / eval_iters


vocab_size = len(chars)
context_len = 8  # context of up to 8 previous characters

d_model = 32     # token embedding dim

h = 4            # number of heads in self-attention   

N = 2            # number of blocks in decoder

dropout = 0.1
batch_size = 32
learning_rate = 1e-3
max_iters = 5000
eval_interval = 500  # keep frequent because we'll overfit

eval_iters = 200  

seed = 3
torch.manual_seed(seed)

gpt = GPT(vocab_size, context_len, d_model, h, N, dropout)
optimizer = optim.Adam(gpt.parameters(), lr=learning_rate)


train_losses = []
val_losses = []
start = time.time()
for step in range(max_iters):   
    # every once in a while evaluate the train and val loss

    if step % eval_interval == 0:
        train_loss, val_loss = eval_loss(gpt, eval_iters)
        # print('='*30)

        print(f"step: {step}, train loss: {train_loss:.4f}, eval loss: {val_loss:.4f}, time: {timedelta(seconds=time.time()-start)}")
        # print('='*30)

        train_losses.append(train_loss)
        val_losses.append(val_loss)

    xb, yb = get_batch(train_data, context_len, batch_size)

    loss = compute_loss(gpt, xb, yb)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()


# plt.plot(train_losses, label='train_loss')

# plt.plot(val_losses, label='val_loss')

# plt.legend()


gpt.eval()
start = torch.zeros((1, 1), dtype=torch.long)
print(decode(gpt.generate(start, max_len=500)[0].tolist()))