Method forward

codegeex/oneflow/codegeex_model.py:704–765

(
        self,
        hidden_states,
        query_hidden_state,
        attention_mask,
        layer_past=None,
        get_key_value=False,
        prompt_length=None,
        context_length=None,
    )

        return self.layers[self._get_layer_index(layer_number)]

    def forward(
        self,
        hidden_states,
        query_hidden_state,
        attention_mask,
        layer_past=None,
        get_key_value=False,
        prompt_length=None,
        context_length=None,
    ):
        # data format change to avoid explicit tranposes : [b s h] --> [s b h]
        hidden_states = hidden_states.transpose(0, 1).contiguous()
        query_hidden_state = query_hidden_state.transpose(0, 1).contiguous()

        origin_attention_mask = attention_mask
        if get_key_value:
            presents = []
        for index in range(self.num_layers):
            layer = self._get_layer(index)
            past = None
            if layer_past is not None:
                past = layer_past[index]
            hidden_states, attention_mask = layer(hidden_states,
                                  attention_mask,
                                  layer_past=past,
                                  get_key_value=get_key_value,
                                  prompt_length=prompt_length,
                                  context_length=context_length,
                                  layer_id=index)
            if get_key_value:
                hidden_states, present = hidden_states
                presents.append(present)

        # Use FP32 for Layernorm
        # hidden_states_ = self.final_layernorm(hidden_states.float()).half()
        hidden_states_ = self.final_layernorm(hidden_states)

        #################################
        # top query layer
        #################################
        past = None
        if layer_past is not None:
            past = layer_past[self.num_layers]
        hidden_states = self.topQueryLayer(hidden_states_,
                                           query_hidden_state,
                                           origin_attention_mask,
                                           layer_past=past,
                                           get_key_value=get_key_value,
                                           prompt_length=prompt_length,
                                           context_length=context_length)

        if get_key_value:
            hidden_states, present = hidden_states
            presents.append(present)

        # reverting data format change [s b h] --> [b s h]
        output = hidden_states.transpose(0, 1).contiguous()