Class AllPool

fastdeploy/model_executor/layers/pooler.py:365–378

class AllPool(PoolingMethod):
    def get_supported_tasks(self) -> Set[PoolingTask]:
        return {"encode"}

    def forward_all(
        self,
        hidden_states: paddle.Tensor,
        pooling_cursor: PoolingCursor,
    ) -> Union[list[paddle.Tensor], paddle.Tensor]:

        assert not pooling_cursor.is_partial_prefill(), "partial prefill not supported with ALL pooling"
        hidden_states_lst = list(hidden_states.split(pooling_cursor.num_scheduled_tokens_cpu.tolist()))

        return [hidden_states_lst[i] for i in pooling_cursor.index]


class MeanPool(PoolingMethod):