Method generate

tensorrt_llm/runtime/model_runner.py:836–990

Generates sequences of token ids. The generation-controlling parameters are set in the sampling_config; it will be set to a default one if not passed. You can override any sampling_config's attributes by passing corresponding parameters. Args: batch_inpu

(self,
                 batch_input_ids: List[torch.Tensor],
                 position_ids: List[torch.Tensor] = None,
                 sampling_config: Optional[SamplingConfig] = None,
                 prompt_table: Optional[Union[str, torch.Tensor]] = None,
                 prompt_tasks: Optional[str] = None,
                 lora_uids: Optional[list] = None,
                 streaming: bool = False,
                 output_generation_logits: bool = False,
                 stopping_criteria: Optional[StoppingCriteria] = None,
                 logits_processor: Optional[LogitsProcessor] = None,
                 medusa_choices: Optional[List[List[int]]] = None,
                 encoder_max_input_length: int = None,
                 encoder_input_features: List[torch.Tensor] = None,
                 encoder_output_lengths: List[torch.Tensor] = None,
                 cross_attention_masks: List[torch.Tensor] = None,
                 **kwargs)

        return self.session.gather_generation_logits

    def generate(self,
                 batch_input_ids: List[torch.Tensor],
                 position_ids: List[torch.Tensor] = None,
                 sampling_config: Optional[SamplingConfig] = None,
                 prompt_table: Optional[Union[str, torch.Tensor]] = None,
                 prompt_tasks: Optional[str] = None,
                 lora_uids: Optional[list] = None,
                 streaming: bool = False,
                 output_generation_logits: bool = False,
                 stopping_criteria: Optional[StoppingCriteria] = None,
                 logits_processor: Optional[LogitsProcessor] = None,
                 medusa_choices: Optional[List[List[int]]] = None,
                 encoder_max_input_length: int = None,
                 encoder_input_features: List[torch.Tensor] = None,
                 encoder_output_lengths: List[torch.Tensor] = None,
                 cross_attention_masks: List[torch.Tensor] = None,
                 **kwargs) -> Union[torch.Tensor, dict]:
        """
        Generates sequences of token ids.
        The generation-controlling parameters are set in the sampling_config; it will be set to a default one if not passed.
        You can override any sampling_config's attributes by passing corresponding parameters.

        Args:
            batch_input_ids (List[torch.Tensor]):
                A list of input id tensors. Each tensor is of shape (sequence_length, ).
            sampling_config (SamplingConfig):
                The sampling configuration to be used as base parametrization for the generation call.
                The passed **kwargs matching the sampling_config's attributes will override them.
                If the sampling_config is not provided, a default will be used.
            prompt_table (str or torch.Tensor):
                The file path of prompt table (.npy format, exported by nemo_prompt_convert.py) or the prompt table itself.
            prompt_tasks (str):
                The prompt tuning task ids for the input batch, in format of comma-separated list (e.g., 0,3,1,0).
            lora_uids (list):
                The uids of LoRA weights for the input batch. Use -1 to disable the LoRA module.
            streaming (bool):
                Whether or not to use streaming mode for generation.
            stopping_criteria (StoppingCriteria):
                Custom stopping criteria.
            logits_processor (LogitsProcessor):
                Custom logits processors.
            medusa_choices (List[List[int]]):
                Medusa decoding choices.
            kwargs (Dict[str, Any]:
                Ad hoc parametrization of sampling_config.
                The passed **kwargs matching the sampling_config's attributes will override them.
        Returns:
            torch.Tensor or dict:
                If return_dict=False, the method returns generated output_ids.
                If return_dict=True, the method returns a dict of output_ids,
                sequence_lengths (if sampling_config.output_sequence_lengths=True),
                context_logits and generation_logits (if self.gather_context_logits=True
                and self.gather_generation_logits=True, respectively).
        """
        # Use sampling_config like HF's generation_config
        if sampling_config is None:
            sampling_config = SamplingConfig(end_id=None, pad_id=None)
        else: