Method sample

tensorrt_llm/serve/scripts/benchmark_dataset.py:426–554

(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        prefix_len: int = DEFAULT_PREFIX_LEN,
        range_ratio: float = DEFAULT_RANGE_RATIO,
        input_len: int = DEFAULT_INPUT_LEN,
        output_len: int = DEFAULT_OUTPUT_LEN,
        **kwargs,
    )

                RandomDataset.SHAREGPT_URL.split("/")[-1], download_timeout)

    def sample(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_requests: int,
        prefix_len: int = DEFAULT_PREFIX_LEN,
        range_ratio: float = DEFAULT_RANGE_RATIO,
        input_len: int = DEFAULT_INPUT_LEN,
        output_len: int = DEFAULT_OUTPUT_LEN,
        **kwargs,
    ) -> list[SampleRequest]:
        # Enforce range_ratio < 1
        if range_ratio >= 1.0:
            raise ValueError(
                "random_range_ratio must be < 1.0 to ensure a valid sampling range"
            )

        vocab_size = tokenizer.vocab_size

        prefix_token_ids = (torch.randint(
            0, vocab_size, size=(prefix_len, ), generator=self.rng).tolist()
                            if prefix_len > 0 else [])

        # New sampling logic: [X * (1 - b), X * (1 + b)]
        input_low = int(input_len * (1 - range_ratio))
        input_high = int(input_len * (1 + range_ratio))
        output_low = int(output_len * (1 - range_ratio))
        output_high = int(output_len * (1 + range_ratio))

        # Add logging for debugging
        logger.debug("Sampling input_len from [%s, %s]", input_low, input_high)
        logger.debug("Sampling output_len from [%s, %s]", output_low,
                     output_high)

        input_lens = torch.randint(input_low,
                                   input_high + 1,
                                   size=(num_requests, ),
                                   generator=self.rng).tolist()
        output_lens = torch.randint(output_low,
                                    output_high + 1,
                                    size=(num_requests, ),
                                    generator=self.rng).tolist()
        offsets = torch.randint(0,
                                vocab_size,
                                size=(num_requests, ),
                                generator=self.rng).tolist()

        requests = []
        if self.sample_from_sharegpt:
            with open(self.dataset_path) as f:
                dataset = json.load(f)
            # Filter out the conversations with less than 2 turns.
            dataset = [
                data for data in dataset
                if len(data.get("conversations", data.get("conversation", [])))
                >= 2
            ]
            # Only keep the first turn of each conversation.
            dataset = [