Transformer language model. Arguments: transformer_hparams: transformer hyperparameters attention_mask_func: a function that takes `unmaksed-attention-scores` with size [b, np, s, s] and an `attention-mask` and will apply the masking. The function should
| 479 | |
| 480 | |
| 481 | class TransformerLanguageModel(MegatronModule): |
| 482 | """Transformer language model. |
| 483 | |
| 484 | Arguments: |
| 485 | transformer_hparams: transformer hyperparameters |
| 486 | attention_mask_func: a function that takes `unmaksed-attention-scores` |
| 487 | with size [b, np, s, s] and an `attention-mask` and will apply |
| 488 | the masking. The function should return a masked score of the |
| 489 | same size [b, np, s, s]. |
| 490 | masked-attention-scores = attention_mask_func( |
| 491 | unmaksed-attention-scores, attention-mask) |
| 492 | vocab_size: vocabulary size |
| 493 | max_sequence_length: maximum size of sequence. This |
| 494 | is used for positional embedding |
| 495 | embedding_dropout_prob: dropout probability for embeddings |
| 496 | num_tokentypes: size of the token-type embeddings. 0 value |
| 497 | will ignore this embedding |
| 498 | """ |
| 499 | |
| 500 | def __init__(self, |
| 501 | init_method, |
| 502 | output_layer_init_method, |
| 503 | num_tokentypes=0, |
| 504 | add_pooler=False): |
| 505 | super(TransformerLanguageModel, self).__init__() |
| 506 | args = get_args() |
| 507 | |
| 508 | self.hidden_size = args.hidden_size |
| 509 | self.num_tokentypes = num_tokentypes |
| 510 | self.init_method = init_method |
| 511 | self.add_pooler = add_pooler |
| 512 | |
| 513 | # Embeddings |
| 514 | self.embedding = Embedding(self.hidden_size, |
| 515 | args.padded_vocab_size, |
| 516 | args.max_position_embeddings, |
| 517 | args.hidden_dropout, |
| 518 | self.init_method, |
| 519 | self.num_tokentypes) |
| 520 | self._embedding_key = 'embedding' |
| 521 | |
| 522 | # Query embeddings |
| 523 | self.topQueryEmbedding = QueryEmbedding(self.hidden_size, |
| 524 | args.padded_vocab_size, |
| 525 | args.max_position_embeddings, |
| 526 | args.hidden_dropout, |
| 527 | self.init_method, |
| 528 | self.num_tokentypes) |
| 529 | self._topQueryEmbedding_key = 'topQueryEmbedding' |
| 530 | |
| 531 | # Transformer |
| 532 | self.transformer = ParallelTransformer( |
| 533 | self.init_method, |
| 534 | output_layer_init_method) |
| 535 | self._transformer_key = 'transformer' |
| 536 | |
| 537 | def set_input_tensor(self, input_tensor): |
| 538 | """See megatron.model.transformer.set_input_tensor()""" |
no outgoing calls
no test coverage detected