Added files

transformer_from_scratch/model.py

@@ -62,4 +62,112 @@ class LayerNormalization(nn.Module):
 class FeedForwardBlock(nn.Module):
     def __init__(self, d_model: int, d_ff: int, dropout: float):
         super().__init__()
-        self.linear1 = nn.Linear(d_model, d_ff)
+        self.linear1 = nn.Linear(d_model, d_ff)  # W1 and b1, bias = True
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(d_ff, d_model)  # W2 and b2, bias = True
+
+    def forward(self, x):
+        # (Batch, Seq_len, d_model) --> (Batch, Seq_len, d_ff) --> (Batch, Seq-len, d_model)
+        return self.linear2(self.dropout(torch.relu(self.linear1(x))))
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, d_model: int, h: int, dropout: float):
+        super().__init__()
+        self.d_model = d_model
+        self.h = h
+
+        assert d_model % h == 0, "d_model must be divisible by h"
+
+        self.d_k = d_model // h
+        self.w_q = nn.Linear(d_model, d_model)  # Wq
+        self.w_k = nn.Linear(d_model, d_model)  # Wk
+        self.w_v = nn.Linear(d_model, d_model)  # Wv
+
+        self.wo = nn.Linear(d_model, d_model)  # Wo
+        self.dropout = nn.Dropout(dropout)
+
+    @staticmethod
+    def attention(query, key, value, mask, dropout: nn.Dropout):
+        d_k = query.size(-1)
+
+        # (Batch, h, seq_len, d_k) --> (Batch, h, seq_len, seq_len)
+        attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)
+        if mask is not None:
+            attention_scores.masked_fill_(mask == 0, -1e9)
+        attention_scores = attention_scores.softmax(
+            dim=-1
+        )  # (Batch, h, seq_Len, seq_len)
+        if dropout is not None:
+            attention_scores = dropout(attention_scores)
+
+        return (attention_scores @ value), attention_scores
+
+    def forward(self, q, k, v, mask):
+        query = self.w_q(q)  # (Batch, Seq_Len, d_model) --> (Batch, Seq_Len, d_model)
+        key = self.w_k(k)  # (Batch, Seq_Len, d_model) --> (Batch, Seq_Len, d_model)
+        value = self.w_v(v)  # (Batch, Seq_Len, d_model) --> (Batch, Seq_Len, d_model)
+
+        # (Batch, Seq_Len, d_model) --> (Batch, Seq_len, h, d_k) --> (Batch, h, Seq_len, d_k)
+        query = query.view(query.shape[0], query.shape[1], self.h, self.d_k).transpose(
+            1, 2
+        )
+        key = key.view(key.shape[0], key.shape[1], self.h, self.d_k).transpose(1, 2)
+        value = value.view(value.shape[0], value.shape[1], self.h, self.d_k).transpose(
+            1, 2
+        )
+
+        x, attention_scores = MultiHeadAttention.attention(
+            query, key, value, mask, self.dropout
+        )
+
+        # (Batch, h, seq_len, d_k) --> (Batch, Seq_len, h, d_k) --> (Batch, Seq_len, d_model), contiguous - in place
+        x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.h * self.d_k)
+
+        # (Batch, Seq_len, d_model) --> (Batch, Seq_len, d_model)
+        return self.wo(x)
+
+
+class ResidualConnection(nn.Module):
+    def __init__(self, dropout: float):
+        super().__init__()
+        self.dropout = nn.Dropout(dropout)
+        self.norm = LayerNormalization()
+
+    def forward(self, x, sublayer):  # sublayer - the previous layer
+        return x + self.dropout(sublayer(self.norm(x)))
+
+
+class EncoderBlock(nn.Module):
+    def __init__(
+        self,
+        self_attention_block: MultiHeadAttention,
+        feed_forward_block: FeedForwardBlock,
+        dropout: float,
+    ):
+        super().__init__()
+        self.self_attention_block = self_attention_block
+        self.feed_forward_block = feed_forward_block
+        self.residual_connections = nn.ModuleList(
+            [ResidualConnection(dropout) for _ in range(2)]
+        )
+
+    def forward(self, x, src_mask):
+        x = self.residual_connections[0](
+            x, lambda x: self.self_attention_block(x, x, x, src_mask)
+        )
+        x = self.residual_connections[1](x, self.feed_forward_block)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(self, layers: nn.ModuleList):
+        super().__init__()
+        self.layers = layers
+        self.norm = LayerNormalization()
+
+    def forward(self, x, mask):
+        for layer in self.layers:
+            x = layer(x, mask)
+
+        return self.norm(x)