Class LogprobsLists

fastdeploy/worker/output.py:39–69

class LogprobsLists(NamedTuple):
    """ """

    # [num_reqs, max_num_logprobs + 1]
    logprob_token_ids: list[list[int]]
    # [num_reqs, max_num_logprobs + 1]
    logprobs: list[list[float]]
    # [num_reqs]
    sampled_token_ranks: list[int]

    def slice_columns(self, start: int, end: int):
        """
        Slice columns (per-row top-k logprobs and token IDs).
        Keeps the number of requests unchanged.
        """
        return LogprobsLists(
            [row[start:end] for row in self.logprob_token_ids],
            [row[start:end] for row in self.logprobs],
            self.sampled_token_ranks,  # unchanged
        )

    def slice_rows(self, start: int, end: int):
        """
        Slice rows.
        Keeps the number of max_num_logprobs unchanged.
        """
        return LogprobsLists(
            self.logprob_token_ids[start:end],
            self.logprobs[start:end],
            self.sampled_token_ranks[start:end],
        )


class LogprobsTensors(NamedTuple):