PanguAlpha evaluation net Args: backbone: backbone network of PanguAlpha generate: enable generate mode Inputs: input_ids: the tokenized inpus init_reset: whether reset saved states Returns: outputs: Tensor, corresponding output for different
| 595 | |
| 596 | |
| 597 | class LogitsNet(nn.Cell): |
| 598 | """ |
| 599 | PanguAlpha evaluation net |
| 600 | Args: |
| 601 | backbone: backbone network of PanguAlpha |
| 602 | generate: enable generate mode |
| 603 | Inputs: |
| 604 | input_ids: the tokenized inpus |
| 605 | init_reset: whether reset saved states |
| 606 | Returns: |
| 607 | outputs: Tensor, corresponding output for different tasks |
| 608 | """ |
| 609 | |
| 610 | def __init__(self, backbone, generate=False, pad_token=6, seq_length=2048): |
| 611 | super(LogitsNet, self).__init__(auto_prefix=False) |
| 612 | self.backbone = backbone |
| 613 | self.pad_token = pad_token |
| 614 | self.argmax = P.Argmax() |
| 615 | self.generate = generate |
| 616 | self.topk = P.TopK(sorted=True).shard(((1, 1),)) |
| 617 | self.gather = P.Gather().shard(((1, 1), (1,))) |
| 618 | self.log_softmax = P.LogSoftmax().shard(((1, 1, 1),)) |
| 619 | self.get_attention_mask = AttentionMask(seq_length) |
| 620 | self.expand = P.ExpandDims().shard(((1, 1, 1),)) |
| 621 | self.all_ones_attention_mask = Tensor(np.ones((1, 1, seq_length)), mstype.float32) |
| 622 | self.not_equal = P.NotEqual().shard(((1, 1), ())) |
| 623 | |
| 624 | def construct(self, input_ids, init_reset=True, batch_valid_length=None, attention_mask=None): |
| 625 | """evaluation net""" |
| 626 | # input_mask = F.cast(F.not_equal(input_ids, self.pad_token), mstype.float32) |
| 627 | input_mask = F.cast(self.not_equal(input_ids, self.pad_token), mstype.float32) |
| 628 | bs, seq_length = F.shape(input_ids) |
| 629 | if attention_mask is None: |
| 630 | if self.is_first_iteration is False: |
| 631 | attention_mask = P.Tile()(self.all_ones_attention_mask, (bs, 1, 1)) |
| 632 | else: |
| 633 | attention_mask = self.get_attention_mask(input_mask) |
| 634 | input_position = F.tuple_to_array(F.make_range(seq_length)) |
| 635 | input_position = P.Tile()(input_position, (bs, 1)) |
| 636 | logits = self.backbone(input_ids, input_position, attention_mask, |
| 637 | init_reset, batch_valid_length) |
| 638 | |
| 639 | return logits |
| 640 | |
| 641 | |
| 642 | class PanGUAlphaWithFinetuneLoss(Cell): |