Class RewardPoolerHead

fastdeploy/model_executor/layers/pooler.py:295–311

class RewardPoolerHead(PoolerHead):

    def __init__(self, model_config: Optional["ModelConfig"] = None) -> None:
        super().__init__(activation=PoolerClassify(static_num_labels=False))
        self.model_config = model_config

    def forward(self, pooled_data: Union[list[paddle.Tensor], paddle.Tensor], pooling_metadata: PoolingMetadata):
        pooling_params = get_pooling_params(pooling_metadata)

        flags = [p.softmax for p in pooling_params]
        if len(set(flags)) == 1:
            if flags[0]:
                pooled_data = self.activation(pooled_data)
        else:
            pooled_data = [self.activation(vecs) if f else vecs for vecs, f in zip(pooled_data, flags)]

        return pooled_data


class PoolingMethod(nn.Layer, ABC):