13importmath1415importtorch16importtorch.nnasnn1718fromlabml\_nn.utilsimportclone\_module\_list19from.feed\_forwardimportFeedForward20from.mhaimportMultiHeadAttention21from.positional\_encodingimportget\_positional\_encoding

24classEmbeddingsWithPositionalEncoding(nn.Module):

31def\_\_init\_\_(self,d\_model:int,n\_vocab:int,max\_len:int=5000):32super().\_\_init\_\_()33self.linear=nn.Embedding(n\_vocab,d\_model)34self.d\_model=d\_model35self.register\_buffer('positional\_encodings',get\_positional\_encoding(d\_model,max\_len))

37defforward(self,x:torch.Tensor):38pe=self.positional\_encodings[:x.shape[0]].requires\_grad\_(False)39returnself.linear(x)\*math.sqrt(self.d\_model)+pe

42classEmbeddingsWithLearnedPositionalEncoding(nn.Module):

49def\_\_init\_\_(self,d\_model:int,n\_vocab:int,max\_len:int=5000):50super().\_\_init\_\_()51self.linear=nn.Embedding(n\_vocab,d\_model)52self.d\_model=d\_model53self.positional\_encodings=nn.Parameter(torch.zeros(max\_len,1,d\_model),requires\_grad=True)

55defforward(self,x:torch.Tensor):56pe=self.positional\_encodings[:x.shape[0]]57returnself.linear(x)\*math.sqrt(self.d\_model)+pe

60classTransformerLayer(nn.Module):

78def\_\_init\_\_(self,\*,79d\_model:int,80self\_attn:MultiHeadAttention,81src\_attn:MultiHeadAttention=None,82feed\_forward:FeedForward,83dropout\_prob:float):

91super().\_\_init\_\_()92self.size=d\_model93self.self\_attn=self\_attn94self.src\_attn=src\_attn95self.feed\_forward=feed\_forward96self.dropout=nn.Dropout(dropout\_prob)97self.norm\_self\_attn=nn.LayerNorm([d\_model])98ifself.src\_attnisnotNone:99self.norm\_src\_attn=nn.LayerNorm([d\_model])100self.norm\_ff=nn.LayerNorm([d\_model])

102self.is\_save\_ff\_input=False

104defforward(self,\*,105x:torch.Tensor,106mask:torch.Tensor,107src:torch.Tensor=None,108src\_mask:torch.Tensor=None):

110z=self.norm\_self\_attn(x)

112self\_attn=self.self\_attn(query=z,key=z,value=z,mask=mask)

114x=x+self.dropout(self\_attn)

119ifsrcisnotNone:

121z=self.norm\_src\_attn(x)

123attn\_src=self.src\_attn(query=z,key=src,value=src,mask=src\_mask)

125x=x+self.dropout(attn\_src)

128z=self.norm\_ff(x)

130ifself.is\_save\_ff\_input:131self.ff\_input=z.clone()

133ff=self.feed\_forward(z)

135x=x+self.dropout(ff)136137returnx

140classEncoder(nn.Module):

147def\_\_init\_\_(self,layer:TransformerLayer,n\_layers:int):148super().\_\_init\_\_()

150self.layers=clone\_module\_list(layer,n\_layers)

152self.norm=nn.LayerNorm([layer.size])

154defforward(self,x:torch.Tensor,mask:torch.Tensor):

156forlayerinself.layers:157x=layer(x=x,mask=mask)

159returnself.norm(x)

162classDecoder(nn.Module):

169def\_\_init\_\_(self,layer:TransformerLayer,n\_layers:int):170super().\_\_init\_\_()

172self.layers=clone\_module\_list(layer,n\_layers)

174self.norm=nn.LayerNorm([layer.size])

176defforward(self,x:torch.Tensor,memory:torch.Tensor,src\_mask:torch.Tensor,tgt\_mask:torch.Tensor):

178forlayerinself.layers:179x=layer(x=x,mask=tgt\_mask,src=memory,src\_mask=src\_mask)

181returnself.norm(x)

184classGenerator(nn.Module):

194def\_\_init\_\_(self,n\_vocab:int,d\_model:int):195super().\_\_init\_\_()196self.projection=nn.Linear(d\_model,n\_vocab)

198defforward(self,x):199returnself.projection(x)

202classEncoderDecoder(nn.Module):

209def\_\_init\_\_(self,encoder:Encoder,decoder:Decoder,src\_embed:nn.Module,tgt\_embed:nn.Module,generator:nn.Module):210super().\_\_init\_\_()211self.encoder=encoder212self.decoder=decoder213self.src\_embed=src\_embed214self.tgt\_embed=tgt\_embed215self.generator=generator

219forpinself.parameters():220ifp.dim()\>1:221nn.init.xavier\_uniform\_(p)

223defforward(self,src:torch.Tensor,tgt:torch.Tensor,src\_mask:torch.Tensor,tgt\_mask:torch.Tensor):

225enc=self.encode(src,src\_mask)

227returnself.decode(enc,src\_mask,tgt,tgt\_mask)

229defencode(self,src:torch.Tensor,src\_mask:torch.Tensor):230returnself.encoder(self.src\_embed(src),src\_mask)

232defdecode(self,memory:torch.Tensor,src\_mask:torch.Tensor,tgt:torch.Tensor,tgt\_mask:torch.Tensor):233returnself.decoder(self.tgt\_embed(tgt),memory,src\_mask,tgt\_mask)