| 366 | class T5Model(nn.Module): |
| 367 | |
| 368 | def __init__(self, |
| 369 | vocab_size, |
| 370 | dim, |
| 371 | dim_attn, |
| 372 | dim_ffn, |
| 373 | num_heads, |
| 374 | encoder_layers, |
| 375 | decoder_layers, |
| 376 | num_buckets, |
| 377 | shared_pos=True, |
| 378 | dropout=0.1): |
| 379 | super(T5Model, self).__init__() |
| 380 | self.vocab_size = vocab_size |
| 381 | self.dim = dim |
| 382 | self.dim_attn = dim_attn |
| 383 | self.dim_ffn = dim_ffn |
| 384 | self.num_heads = num_heads |
| 385 | self.encoder_layers = encoder_layers |
| 386 | self.decoder_layers = decoder_layers |
| 387 | self.num_buckets = num_buckets |
| 388 | |
| 389 | # layers |
| 390 | self.token_embedding = nn.Embedding(vocab_size, dim) |
| 391 | self.encoder = T5Encoder(self.token_embedding, dim, dim_attn, dim_ffn, |
| 392 | num_heads, encoder_layers, num_buckets, |
| 393 | shared_pos, dropout) |
| 394 | self.decoder = T5Decoder(self.token_embedding, dim, dim_attn, dim_ffn, |
| 395 | num_heads, decoder_layers, num_buckets, |
| 396 | shared_pos, dropout) |
| 397 | self.head = nn.Linear(dim, vocab_size, bias=False) |
| 398 | |
| 399 | # initialize weights |
| 400 | self.apply(init_weights) |
| 401 | |
| 402 | def forward(self, encoder_ids, encoder_mask, decoder_ids, decoder_mask): |
| 403 | x = self.encoder(encoder_ids, encoder_mask) |