def pool(self, df):
        """Pooling to aggregate token-wise embeddings to sentence embeddings

        Args:
            df: pd.DataFrame with columns text_index (int), token (str),
                layer_index (int), values (list[float])

        Returns:
            pd.DataFrame grouped by text index and layer index
        """

        def max_pool(x):
            values = np.array(
                [
                    np.reshape(np.array(x.values[i]), self.embedding_dim)
                    for i in range(x.values.shape[0])
                ]
            )
            m, _ = torch.max(torch.tensor(values, dtype=torch.float), 0)
            return m.numpy()

        def mean_pool(x):
            values = np.array(
                [
                    np.reshape(np.array(x.values[i]), self.embedding_dim)
                    for i in range(x.values.shape[0])
                ]
            )
            return torch.mean(torch.tensor(values, dtype=torch.float), 0).numpy()

        def cls_pool(x):
            values = np.array(
                [
                    np.reshape(np.array(x.values[i]), self.embedding_dim)
                    for i in range(x.values.shape[0])
                ]
            )
            return values[0]

        try:
            if self.pooling_strategy == "max":
                pool_func = max_pool
            elif self.pooling_strategy == "mean":
                pool_func = mean_pool
            elif self.pooling_strategy == "cls":
                pool_func = cls_pool
            else:
                raise ValueError("Please enter valid pooling strategy")
        except ValueError as ve:
            print(ve)

        return (
            df.groupby(["text_index", "layer_index"])["values"]
            .apply(lambda x: pool_func(x))
            .reset_index()
        )
Pooling to aggregate token-wise embeddings to sentence embeddings

        Args:
            df: pd.DataFrame with columns text_index (int), token (str),
                layer_index (int), values (list[float])

        Returns:
            pd.DataFrame grouped by text index and layer index
        
utils_nlp/models/bert/sequence_encoding.py