Method execute

models_jittor/model.py:255–344

(
        self,
        input_ids: Optional[jt.Var] = None,
        past_key_values: Optional[Tuple[Tuple[jt.Var]]] = None,
        attention_mask: Optional[jt.Var] = None,
        token_type_ids: Optional[jt.Var] = None,
        position_ids: Optional[jt.Var] = None,
        head_mask: Optional[jt.Var] = None,
        inputs_embeds: Optional[jt.Var] = None,
        use_cache: Optional[bool] = None,
    )

        self.apply(partial(_init_weights, config))

    def execute(
        self,
        input_ids: Optional[jt.Var] = None,
        past_key_values: Optional[Tuple[Tuple[jt.Var]]] = None,
        attention_mask: Optional[jt.Var] = None,
        token_type_ids: Optional[jt.Var] = None,
        position_ids: Optional[jt.Var] = None,
        head_mask: Optional[jt.Var] = None,
        inputs_embeds: Optional[jt.Var] = None,
        use_cache: Optional[bool] = None,
    ):
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = input_ids.size()
            input_ids = input_ids.view(-1, input_shape[-1])
            batch_size = input_ids.shape[0]
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            batch_size = inputs_embeds.shape[0]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        if token_type_ids is not None:
            token_type_ids = token_type_ids.view(-1, input_shape[-1])

        if position_ids is not None:
            position_ids = position_ids.view(-1, input_shape[-1])

        if past_key_values is None:
            past_length = 0
            past_key_values = tuple([None] * len(self.h))
        else:
            past_length = past_key_values[0][0].size(-2)

        if position_ids is None:
            position_ids = jt.arange(past_length, input_shape[-1] + past_length, dtype='int64')
            position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])

        # Attention mask.
        if attention_mask is not None:
            if batch_size <= 0:
                raise ValueError("batch_size has to be defined and > 0")
            attention_mask = attention_mask.view(batch_size, -1)
            # [batch_size, 1, 1, to_seq_length]
            attention_mask = attention_mask[:, None, None, :]

            if jt.flags.amp_level >= 3:
                attention_mask = attention_mask.half()  # fp16 compatibility
                attention_mask = (1.0 - attention_mask) * -65504.0
            else:
                # finfo.min
                attention_mask = (1.0 - attention_mask) * -3.402e38

        # n_layer x batch x num_attention_heads x N x N
        head_mask = get_head_mask(head_mask, self.config.n_layer)