Spaces:
Running on Zero
Running on Zero
| import torch | |
| import torch.nn as nn | |
| from src.transformer.attention import Attention | |
| from src.transformer.feed_forward import FeedForward | |
| from src.transformer.mutli_head_attention import MultiHeadAttention | |
| class TransformerBlock(nn.Module): | |
| def __init__(self, d_model: int, d_ff: int, context_length: int, dropout=0.1, number_of_heads=1): | |
| super().__init__() | |
| self.number_of_heads = number_of_heads | |
| print(f"\tTransformer Block: number of heads: {number_of_heads}") | |
| if self.number_of_heads == 1: | |
| self.attention = Attention(d_model, context_length) | |
| else: | |
| self.attention = MultiHeadAttention(d_model, context_length, number_of_heads) | |
| self.feed_forward = FeedForward(d_model, d_ff) | |
| self.layer_norm_1 = nn.LayerNorm(d_model) | |
| self.layer_norm_2 = nn.LayerNorm(d_model) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x, padding_mask: torch.Tensor | None = None,) -> torch.Tensor: | |
| attention_output = self.attention(x, padding_mask=padding_mask) | |
| x = self.layer_norm_1( | |
| x + self.dropout(attention_output) | |
| ) | |
| ffn_output = self.feed_forward(x) | |
| x = self.layer_norm_2( | |
| x + self.dropout(ffn_output) #x is the residual/skip connection... makes sure the original x still influences the learning | |
| ) | |
| return x |