Class EvalNet

codegeex/mindspore/src/pangu_alpha.py:547–594

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

class EvalNet(nn.Cell):
    """
    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:
        outputs: Tensor, corresponding output for different tasks
    """

    def __init__(self, backbone, generate=False, pad_token=6, seq_length=2048):
        super(EvalNet, self).__init__(auto_prefix=False)
        self.backbone = backbone
        self.pad_token = pad_token
        self.argmax = P.Argmax()
        self.generate = generate
        self.topk = P.TopK(sorted=True).shard(((1, 1),))
        self.gather = P.Gather().shard(((1, 1), (1,)))
        self.log_softmax = P.LogSoftmax().shard(((1, 1, 1),))
        self.get_attention_mask = AttentionMask(seq_length)
        self.expand = P.ExpandDims().shard(((1, 1, 1),))
        self.all_ones_attention_mask = Tensor(np.ones((1, 1, seq_length)), mstype.float32)
        self.not_equal = P.NotEqual().shard(((1, 1), ()))

    def construct(self, input_ids, current_index, init_reset=True, batch_valid_length=None):
        """evaluation net"""
        # input_mask = F.cast(F.not_equal(input_ids, self.pad_token), mstype.float32)
        input_mask = F.cast(self.not_equal(input_ids, self.pad_token), mstype.float32)
        bs, seq_length = F.shape(input_ids)
        if self.is_first_iteration is False:
            attention_mask = P.Tile()(self.all_ones_attention_mask, (bs, 1, 1))
        else:
            attention_mask = self.get_attention_mask(input_mask)
        input_position = F.tuple_to_array(F.make_range(seq_length))
        input_position = P.Tile()(input_position, (bs, 1))
        logits = self.backbone(input_ids, input_position, attention_mask,
                               init_reset, batch_valid_length)
        index = current_index.view(-1, )
        # P.Print()("==logits_is:", logits, ",shape is:", logits.shape)
        # P.Print()("==index_is:", index, ",shape is:", index.shape)
        logits = self.gather(logits, index, 0)
        logits = logits.view(bs, 1, -1)
        log_probs = self.log_softmax(logits)
        return log_probs


class LogitsNet(nn.Cell):