Class LogprobsTensors

fastdeploy/worker/output.py:72–126

class LogprobsTensors(NamedTuple):
    """ """

    # [num_reqs, max_num_logprobs + 1]
    logprob_token_ids: paddle.Tensor
    # [num_reqs, max_num_logprobs + 1]
    logprobs: paddle.Tensor
    # [num_reqs]
    selected_token_ranks: paddle.Tensor

    def tolists(self):
        """Convert to lists."""
        return LogprobsLists(
            self.logprob_token_ids.tolist(),
            self.logprobs.tolist(),
            self.selected_token_ranks.tolist(),
        )

    @staticmethod
    def empty_cpu(num_positions: int, num_tokens_per_position: int) -> "LogprobsTensors":
        """Create empty LogprobsTensors on CPU."""

        logprob_token_ids = paddle.empty([num_positions, num_tokens_per_position], device="cpu", dtype=paddle.int64)
        logprobs = paddle.empty_like(logprob_token_ids, device="cpu", dtype=paddle.float32)
        selected_token_ranks = paddle.empty([num_positions], device="cpu", dtype=paddle.int64)
        return LogprobsTensors(
            logprob_token_ids=logprob_token_ids,
            logprobs=logprobs,
            selected_token_ranks=selected_token_ranks,
        )

    @staticmethod
    def empty(num_positions: int, num_tokens_per_position: int) -> "LogprobsTensors":
        """Create empty LogprobsTensors on default device."""

        logprob_token_ids = paddle.empty([num_positions, num_tokens_per_position], dtype=paddle.int64)
        logprobs = paddle.empty_like(logprob_token_ids, dtype=paddle.float32)
        selected_token_ranks = paddle.empty([num_positions], dtype=paddle.int64)
        return LogprobsTensors(
            logprob_token_ids=logprob_token_ids,
            logprobs=logprobs,
            selected_token_ranks=selected_token_ranks,
        )

    def slice_rows(self, start: int, end: int):
        """
        Slice rows.
        Keeps the number of max_num_logprobs unchanged.
        """
        with paddle.no_grad():
            return LogprobsTensors(
                paddle.to_tensor(self.logprob_token_ids.cpu()[start:end], place="cpu"),
                paddle.to_tensor(self.logprobs.cpu()[start:end], place="cpu"),
                paddle.to_tensor(self.selected_token_ranks.cpu()[start:end], place="cpu"),
            )


PromptLogprobs = LogprobsTensors | list[dict[int, Logprob] | None]