The PanguAlpha network consisting of two parts the backbone and the head Args: config(PanguAlphaConfig): the config of network Inputs: input_ids: the tokenized inputs input_mask: the mask indicating whether each position is a valid input past: the previou
| 469 | |
| 470 | |
| 471 | class PanguAlphaModel(nn.Cell): |
| 472 | """ |
| 473 | The PanguAlpha network consisting of two parts the backbone and the head |
| 474 | Args: |
| 475 | config(PanguAlphaConfig): the config of network |
| 476 | Inputs: |
| 477 | input_ids: the tokenized inputs |
| 478 | input_mask: the mask indicating whether each position is a valid input |
| 479 | past: the previous feature map |
| 480 | Returns: |
| 481 | logits: Tensor: the logits of the corresponding inputs with shape (batch_size, seq_length, vocab_size) |
| 482 | """ |
| 483 | |
| 484 | def __init__(self, config): |
| 485 | super(PanguAlphaModel, self).__init__() |
| 486 | # Network head to get logits over vocabulary |
| 487 | copied_parallel_config = copy.deepcopy(config.parallel_config) |
| 488 | if copied_parallel_config.pipeline_stage > 1: |
| 489 | copied_parallel_config.vocab_emb_dp = False |
| 490 | self.head = PanGuHead( |
| 491 | hidden_size=config.hidden_size, |
| 492 | parallel_config=copied_parallel_config, |
| 493 | ) |
| 494 | self.head.pipeline_stage = config.parallel_config.pipeline_stage - 1 |
| 495 | self.backbone = PanguAlpha_Model(config) |
| 496 | self.backbone.embedding.word_embedding.embedding_table.add_pipeline_stage(self.head.pipeline_stage) |
| 497 | self.print = P.Print() |
| 498 | |
| 499 | def construct(self, input_ids, input_position, attention_mask, |
| 500 | init_reset=True, batch_valid_length=None): |
| 501 | output_states, word_table = self.backbone(input_ids, input_position, attention_mask, |
| 502 | init_reset, batch_valid_length) |
| 503 | self.print("PanguAlphaModel: output_states", output_states) |
| 504 | self.print("PanguAlphaModel: word_table", word_table) |
| 505 | logits = self.head(output_states, word_table) |
| 506 | self.print("PanguAlphaModel: logits", logits) |
| 507 | |
| 508 | return logits |
| 509 | |
| 510 | |
| 511 | class PanGUAlphaWithLoss(Cell): |
no outgoing calls
no test coverage detected