Method forward

tensorrt_llm/models/bert/model.py:302–367

(self,
                input_ids=None,
                input_lengths=None,
                position_ids=None,
                token_type_ids=None,
                hidden_states=None,
                max_input_length=None)

        ])

    def forward(self,
                input_ids=None,
                input_lengths=None,
                position_ids=None,
                token_type_ids=None,
                hidden_states=None,
                max_input_length=None):
        # remove_input_padding requires these fields as explicit input
        mask = None
        if not default_net().plugin_config.remove_input_padding:
            seq_len_2d = concat([1, shape(input_ids, 1)])

            # create position ids
            position_ids_buffer = constant(
                np.expand_dims(
                    np.arange(self.max_position_embeddings).astype(np.int32),
                    0))
            tmp_position_ids = slice(position_ids_buffer,
                                     starts=[0, 0],
                                     sizes=seq_len_2d)
            tmp_position_ids = expand(tmp_position_ids, shape(input_ids))  #BxL
            tmp_input_lengths = unsqueeze(input_lengths, 1)  #Bx1
            tmp_input_lengths = expand(tmp_input_lengths,
                                       shape(input_ids))  #BxL
            mask = tmp_position_ids < tmp_input_lengths  # BxL
            mask = mask.cast('int32')

            if position_ids is None:
                if self.is_roberta:
                    # see create_position_ids_from_input_ids() in https://github.com/huggingface/transformers/blob/main/src/transformers/models/roberta/modeling_roberta.py
                    position_ids = (tmp_position_ids + 1) * mask
                    position_ids = position_ids + self.padding_idx
                else:
                    position_ids = slice(position_ids_buffer,
                                         starts=[0, 0],
                                         sizes=seq_len_2d)
                    position_ids = expand(position_ids, shape(input_ids))

            # create token_type_ids
            if token_type_ids is None:
                token_type_ids_buffer = constant(
                    np.expand_dims(
                        np.zeros(self.max_position_embeddings).astype(np.int32),
                        0))
                token_type_ids = slice(token_type_ids_buffer,
                                       starts=[0, 0],
                                       sizes=seq_len_2d)
                token_type_ids = expand(token_type_ids, shape(input_ids))

        hidden_states = self.embedding(input_ids, position_ids, token_type_ids)
        self.register_network_output('embedding_output', hidden_states)

        for idx, layer in enumerate(self.layers):
            hidden_states = layer(hidden_states=hidden_states,
                                  input_lengths=input_lengths,
                                  attention_mask=mask,
                                  max_input_length=max_input_length)
            # keep the last layer output name as hidden_states