PanguAlpha evaluation net Args: backbone: backbone network of PanguAlpha generate: enable generate mode Inputs: input_ids: the tokenized inpus current_index: the index of current token init_reset: whether reset saved states Returns: ou
| 560 | |
| 561 | |
| 562 | class EvalNet(nn.Cell): |
| 563 | """ |
| 564 | PanguAlpha evaluation net |
| 565 | Args: |
| 566 | backbone: backbone network of PanguAlpha |
| 567 | generate: enable generate mode |
| 568 | Inputs: |
| 569 | input_ids: the tokenized inpus |
| 570 | current_index: the index of current token |
| 571 | init_reset: whether reset saved states |
| 572 | Returns: |
| 573 | outputs: Tensor, corresponding output for different tasks |
| 574 | """ |
| 575 | |
| 576 | def __init__(self, backbone, generate=False, pad_token=6, seq_length=2048): |
| 577 | super(EvalNet, self).__init__(auto_prefix=False) |
| 578 | self.backbone = backbone |
| 579 | self.pad_token = pad_token |
| 580 | self.argmax = P.Argmax() |
| 581 | self.generate = generate |
| 582 | self.topk = P.TopK(sorted=True).shard(((1, 1),)) |
| 583 | self.gather = P.Gather().shard(((1, 1), (1,))) |
| 584 | self.log_softmax = P.LogSoftmax().shard(((1, 1, 1),)) |
| 585 | self.get_attention_mask = AttentionMask(seq_length) |
| 586 | self.expand = P.ExpandDims().shard(((1, 1, 1),)) |
| 587 | self.all_ones_attention_mask = Tensor(np.ones((1, 1, seq_length)), mstype.float32) |
| 588 | self.print = P.Print() |
| 589 | |
| 590 | def construct(self, input_ids, current_index, init_reset=True, batch_valid_length=None): |
| 591 | """evaluation net""" |
| 592 | input_mask = F.cast(F.not_equal(input_ids, self.pad_token), mstype.float32) |
| 593 | bs, seq_length = F.shape(input_ids) |
| 594 | if self.is_first_iteration is False: |
| 595 | attention_mask = P.Tile()(self.all_ones_attention_mask, (bs, 1, 1)) |
| 596 | else: |
| 597 | attention_mask = self.get_attention_mask(input_mask) |
| 598 | self.print("EvalNet: attention_mask", attention_mask) |
| 599 | input_position = F.tuple_to_array(F.make_range(seq_length)) |
| 600 | self.print("EvalNet: input_position_0", input_position) |
| 601 | input_position = P.Tile()(input_position, (bs, 1)) |
| 602 | self.print("EvalNet: input_position_1", input_position) |
| 603 | logits = self.backbone(input_ids, input_position, attention_mask, |
| 604 | init_reset, batch_valid_length) |
| 605 | self.print("EvalNet: logits", logits) |
| 606 | index = current_index.view(-1, ) |
| 607 | self.print("EvalNet: index", index) |
| 608 | logits = self.gather(logits, index, 0) |
| 609 | logits = logits.view(bs, 1, -1) |
| 610 | log_probs = self.log_softmax(logits) |
| 611 | return log_probs |
no outgoing calls
no test coverage detected