MediumNLPnlp-151

Sentence-Transformer Mean Pooling

Problem

Compute masked mean pooling over token vectors. Average only the vectors whose mask entry equals 1.

Return an empty list when no position is active, and take the vector dimension from the first active vector.

Hint

Filter the vectors by mask first, then average each coordinate.

case 1

in: [[[1,2],[3,4],[5,6]],[1,0,1]]

expected: [3,4]

case 2

in: [[[1,1],[2,2]],[1,1]]

expected: [1.5,1.5]

case 3

in: [[[1,2]],[0]]

expected: []

case 4

in: [[[1,2],[3,4]],[1,1]]

expected: [2,3]

case 5

in: [[],[]]

expected: []

Starter codePython

def mean_pool_masked(vectors, mask):
    # Your code here
    pass

5 test cases · run your code against them in the editor