File size: 932 Bytes
38e2dac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | """
Source: https://github.com/karpathy/nanoGPT/blob/master/model.py
"""
from torch import nn
from models.attention import CausalSelfAttention
from models.decoder_ffn import FFNN
class DecoderBlock(nn.Module):
def __init__(
self,
d,
H,
T,
bias=False,
dropout=0.2,
):
"""
Arguments:
d: size of embedding dimension
H: number of attention heads
T: maximum length of input sequences (in tokens)
bias: whether or not to use bias in linear layers
dropout: probability of dropout
"""
super().__init__()
self.ln_1 = nn.LayerNorm(d)
self.attn = CausalSelfAttention(d, H, T, bias, dropout)
self.ln_2 = nn.LayerNorm(d)
self.ffnn = FFNN(d, bias, dropout)
def forward(self, x):
x = x + self.attn(self.ln_1(x))
x = x + self.ffnn(self.ln_2(x))
return x
|