Spaces:
Running on Zero
Running on Zero
| import torch | |
| import torch.nn as nn | |
| from src.transformer.transformer_block import TransformerBlock | |
| class Transformer(nn.Module): | |
| def __init__( | |
| self, | |
| d_model: int, | |
| d_ff: int, | |
| context_length: int, | |
| num_layers: int, | |
| number_of_heads: int, | |
| dropout = 0.1 | |
| ): | |
| super().__init__() | |
| self.blocks = nn.ModuleList([ | |
| TransformerBlock( | |
| d_model=d_model, | |
| d_ff=d_ff, | |
| context_length=context_length, | |
| dropout=dropout, | |
| number_of_heads=number_of_heads | |
| ) | |
| for _ in range(num_layers) | |
| ]) | |
| def forward(self, | |
| x: torch.Tensor, | |
| padding_mask: torch.Tensor | None = None,) -> torch.Tensor: | |
| for block in self.blocks: | |
| x = block(x,padding_mask=padding_mask,) | |
| return x |