Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| encoder_block_size = 33 | |
| decoder_block_size = 30 | |
| class Head(nn.Module): | |
| """ one self-attention head """ | |
| def __init__(self, n_embd, d_k, dropout, mask=0): # d_k is dimention of key , nomaly d_k = n_embd / 4 | |
| super().__init__() | |
| self.mask = mask | |
| self.key = nn.Linear(n_embd, d_k, bias=False, device=device) | |
| self.query = nn.Linear(n_embd, d_k, bias=False, device=device) | |
| self.value = nn.Linear(n_embd, d_k, bias=False, device=device) | |
| if mask: | |
| self.register_buffer('tril', torch.tril(torch.ones(encoder_block_size, encoder_block_size, device=device))) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x, encoder_output = None): | |
| B,T,C = x.shape | |
| if encoder_output is not None: | |
| k = self.key(encoder_output) | |
| Be, Te, Ce = encoder_output.shape | |
| else: | |
| k = self.key(x) # (B,T,d_k) | |
| q = self.query(x) # (B,T,d_k) | |
| # compute attention scores | |
| wei = q @ k.transpose(-2, -1) * C**-0.5 # (B,T,T) | |
| if self.mask: | |
| if encoder_output is not None: | |
| wei = wei.masked_fill(self.tril[:T, :Te] == 0, float('-inf')) # (B,T,T) | |
| else: | |
| wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B,T,T) | |
| wei = F.softmax(wei, dim=-1) | |
| wei = self.dropout(wei) | |
| # perform weighted aggregation of values | |
| if encoder_output is not None: | |
| v = self.value(encoder_output) | |
| else: | |
| v = self.value(x) | |
| out = wei @ v # (B,T,C) | |
| return out | |
| class MultiHeadAttention(nn.Module): | |
| """ multiple self attention heads in parallel """ | |
| def __init__(self, n_embd, num_head, d_k, dropout, mask=0): | |
| super().__init__() | |
| self.heads = nn.ModuleList([Head(n_embd, d_k, dropout, mask) for _ in range(num_head)]) | |
| self.proj = nn.Linear(n_embd, n_embd) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x, encoder_output=None): | |
| out = torch.cat([h(x, encoder_output) for h in self.heads], dim=-1) | |
| out = self.dropout(self.proj(out)) | |
| return out | |
| class FeedForward(nn.Module): | |
| """ multiple self attention heads in parallel """ | |
| def __init__(self, n_embd, dropout): | |
| super().__init__() | |
| self.net = nn.Sequential( | |
| nn.Linear(n_embd, 4 * n_embd), | |
| nn.ReLU(), | |
| nn.Linear(4 * n_embd, n_embd), | |
| nn.Dropout(dropout) | |
| ) | |
| def forward(self, x): | |
| return self.net(x) | |
| class encoderBlock(nn.Module): | |
| """ Tranformer encoder block : communication followed by computation """ | |
| def __init__(self, n_embd, n_head, dropout): | |
| super().__init__() | |
| d_k = n_embd // n_head | |
| self.sa = MultiHeadAttention(n_embd, n_head, d_k, dropout) | |
| self.ffwd = FeedForward(n_embd, dropout) | |
| self.ln1 = nn.LayerNorm(n_embd) | |
| self.ln2 = nn.LayerNorm(n_embd) | |
| def forward(self, x, encoder_output=None): | |
| x = x + self.sa(self.ln1(x), encoder_output) | |
| x = x + self.ffwd(self.ln2(x)) | |
| return x | |
| class Encoder(nn.Module): | |
| def __init__(self, n_embd, n_head, n_layers, dropout): | |
| super().__init__() | |
| self.token_embedding_table = nn.Embedding(input_vocab_size, n_embd) # n_embd: input embedding dimension | |
| self.position_embedding_table = nn.Embedding(encoder_block_size, n_embd) | |
| self.blocks = nn.Sequential(*[encoderBlock(n_embd, n_head, dropout) for _ in range(n_layers)]) | |
| self.ln_f = nn.LayerNorm(n_embd) # final layer norm | |
| def forward(self, idx): | |
| B, T = idx.shape | |
| tok_emb = self.token_embedding_table(idx) # (B,T,n_embd) | |
| pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,n_embd) | |
| x = tok_emb + pos_emb # (B,T,n_embd) | |
| x = self.blocks(x) # apply one attention layer (B,T,C) | |
| x = self.ln_f(x) # (B,T,C) | |
| return x | |
| class decoderBlock(nn.Module): | |
| """ Tranformer decoder block : self communication then cross communication followed by computation """ | |
| def __init__(self, n_embd, n_head, dropout): | |
| super().__init__() | |
| d_k = n_embd // n_head | |
| self.sa = MultiHeadAttention(n_embd, n_head, d_k, dropout, mask = 1) | |
| self.ca = MultiHeadAttention(n_embd, n_head, d_k, dropout, mask = 1) | |
| self.ffwd = FeedForward(n_embd, dropout) | |
| self.ln1 = nn.LayerNorm(n_embd, device=device) | |
| self.ln2 = nn.LayerNorm(n_embd, device=device) | |
| self.ln3 = nn.LayerNorm(n_embd, device=device) | |
| def forward(self, x_encoder_output): | |
| x = x_encoder_output[0] | |
| encoder_output = x_encoder_output[1] | |
| x = x + self.sa(self.ln1(x)) | |
| x = x + self.ca(self.ln2(x), encoder_output) | |
| x = x + self.ffwd(self.ln3(x)) | |
| return (x,encoder_output) | |
| class Decoder(nn.Module): | |
| def __init__(self, n_embd, n_head, n_layers, dropout): | |
| super().__init__() | |
| self.token_embedding_table = nn.Embedding(output_vocab_size, n_embd) # n_embd: input embedding dimension | |
| self.position_embedding_table = nn.Embedding(decoder_block_size, n_embd) | |
| self.blocks = nn.Sequential(*[decoderBlock(n_embd, n_head=n_head, dropout=dropout) for _ in range(n_layers)]) | |
| self.ln_f = nn.LayerNorm(n_embd) # final layer norm | |
| self.lm_head = nn.Linear(n_embd, output_vocab_size) | |
| def forward(self, idx, encoder_output, targets=None): | |
| B, T = idx.shape | |
| tok_emb = self.token_embedding_table(idx) # (B,T,n_embd) | |
| pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,n_embd) | |
| x = tok_emb + pos_emb # (B,T,n_embd) | |
| x =self.blocks((x, encoder_output)) | |
| x = self.ln_f(x[0]) # (B,T,C) | |
| logits = self.lm_head(x) # (B,T,output_vocab_size) | |
| if targets is None: | |
| loss = None | |
| else: | |
| B, T, C = logits.shape | |
| temp_logits = logits.view(B*T, C) | |
| targets = targets.reshape(B*T) | |
| loss = F.cross_entropy(temp_logits, targets.long()) | |
| # print(logits) | |
| # out = torch.argmax(logits) | |
| return logits, loss | |