Method prepare_inputs

tensorrt_llm/models/enc_dec/model.py:2125–2192

(self, max_batch_size=16)

        return x

    def prepare_inputs(self, max_batch_size=16):

        bs_range = [1, (max_batch_size + 1) // 2, max_batch_size]
        min_feat_len, optimal_feat_len = 10, 1000  # 100ms, 10s
        inlen_range = [
            min_feat_len, optimal_feat_len, self.max_audio_feature_seq_len
        ]
        inlen_range_after_downsample = [
            min_feat_len // self.downsample_factor,
            optimal_feat_len // self.downsample_factor,
            self.max_audio_feature_seq_len // self.downsample_factor
        ]
        if not default_net().plugin_config.remove_input_padding:
            x = Tensor(name="input_features",
                       dtype=self._dtype,
                       shape=[-1, self.config.n_mels, -1],
                       dim_range=OrderedDict([
                           ("batch_size", [bs_range]),
                           ("feature_dim", [self.config.n_mels]),
                           ("feature_len_range", [inlen_range]),
                       ]))
            position_ids = Tensor(
                name='position_ids',
                dtype=trt.int32,
                shape=[-1, -1],
                dim_range=OrderedDict([('batch_size', [bs_range]),
                                       ('feature_len_downsample_range',
                                        [inlen_range_after_downsample])]),
            )
        else:
            batch_seqlen_range = [
                1,
                (self.max_audio_feature_seq_len * max_batch_size + 1) // 2,
                self.max_audio_feature_seq_len * max_batch_size,
            ]
            batch_seqlen_downsample_range = [
                1,
                (self.max_audio_feature_seq_len // self.downsample_factor *
                 max_batch_size + 1) // 2,
                self.max_audio_feature_seq_len // self.downsample_factor *
                max_batch_size,
            ]
            x = Tensor(name="input_features",
                       dtype=self._dtype,
                       shape=[-1, self.config.n_mels],
                       dim_range=OrderedDict([
                           ("batch_seqlen_range", [batch_seqlen_range]),
                           ("feature_dim", [self.config.n_mels]),
                       ]))
            position_ids = Tensor(
                name='position_ids',
                dtype=trt.int32,
                shape=[-1],
                dim_range=OrderedDict([('batch_seqlen_downsample_range',
                                        [batch_seqlen_downsample_range])]),
            )
        input_lengths = Tensor(
            name="input_lengths",