def mean_pooling()

in src/feature_extractor.py [0:0]

• 16 lines of code
• 8 McCabe index (conditional complexity)

    def mean_pooling(last_hidden_state, attention_mask):
        """
        Perform mean pooling of the last hidden state followed by a normalization step.

        Args:
            last_hidden_state (np.ndarray): Tensor of shape [batchSize, seqLength, embedDim].
            attention_mask (np.ndarray): Tensor of shape [batchSize, seqLength].

        Returns:
            np.ndarray: A new array of shape [batchSize, embedDim].
        """
        batch_size, seq_length, embed_dim = last_hidden_state.shape

        # Create an empty array for the output
        returned_data = np.zeros((batch_size, embed_dim), dtype=last_hidden_state.dtype)

        # Iterate over the batch and embedding dimensions
        for i in range(batch_size):
            for k in range(embed_dim):
                sum_ = 0.0
                count = 0

                # Iterate over the sequence length
                for j in range(seq_length):
                    attn = float(attention_mask[i, j])  # Convert attention mask to float (0 or 1)
                    sum_ += last_hidden_state[i, j, k] * attn
                    count += attn

                # Avoid division by zero
                avg = sum_ / count if count > 0 else 0.0
                returned_data[i, k] = avg

        # Normalize the returned data using L2 norm
        norms = np.linalg.norm(returned_data, axis=1, keepdims=True)
        normalized_data = returned_data / np.clip(norms, a_min=1e-10, a_max=None)  # Prevent division by zero

        return normalized_data