| 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""" |