All files / embed-code-core/src pooler.ts

100% Statements 42/42
100% Branches 6/6
100% Functions 3/3
100% Lines 29/29

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75                        4x 4x 5x 5x 10x 6x 6x 6x 12x     5x 8x     4x                       2x 2x 3x 3x 6x   2x                           3x 3x 4x 4x 7x 3x 3x     4x 4x 8x   3x    
/**
 * Mean pooling — mask-aware aggregation of hidden states.
 *
 * Computes: pooled[b][d] = sum_s hidden[b][s][d] * mask[b][s] / sum_s mask[b][s]
 */
export function meanPool(
  hidden: Float32Array,
  attentionMask: Int32Array,
  batch: number,
  seqLen: number,
  dim: number,
): Float32Array {
  const out = new Float32Array(batch * dim);
  for (let b = 0; b < batch; b++) {
    let count = 0;
    for (let s = 0; s < seqLen; s++) {
      if (attentionMask[b * seqLen + s]! === 1) {
        count++;
        const src = (b * seqLen + s) * dim;
        const dst = b * dim;
        for (let d = 0; d < dim; d++) out[dst + d] += hidden[src + d]!;
      }
    }
    if (count > 0) {
      for (let d = 0; d < dim; d++) out[b * dim + d] /= count;
    }
  }
  return out;
}
 
/**
 * CLS pooling — returns the [CLS] token embedding (position 0).
 */
export function clsPool(
  hidden: Float32Array,
  batch: number,
  seqLen: number,
  dim: number,
): Float32Array {
  const out = new Float32Array(batch * dim);
  for (let b = 0; b < batch; b++) {
    const src = b * seqLen * dim; // position 0 of sequence b
    const dst = b * dim;
    for (let d = 0; d < dim; d++) out[dst + d] = hidden[src + d]!;
  }
  return out;
}
 
/**
 * Last-token pooling — returns the last non-padding token embedding.
 * Uses attentionMask to find the last real token in each sequence.
 */
export function lastTokenPool(
  hidden: Float32Array,
  attentionMask: Int32Array,
  batch: number,
  seqLen: number,
  dim: number,
): Float32Array {
  const out = new Float32Array(batch * dim);
  for (let b = 0; b < batch; b++) {
    let lastIdx = 0;
    for (let s = seqLen - 1; s >= 0; s--) {
      if (attentionMask[b * seqLen + s]! === 1) {
        lastIdx = s;
        break;
      }
    }
    const src = (b * seqLen + lastIdx) * dim;
    const dst = b * dim;
    for (let d = 0; d < dim; d++) out[dst + d] = hidden[src + d]!;
  }
  return out;
}