pub fn pool_embeddings(
hidden: &[f32],
attention_mask: &[&[u32]],
batch: usize,
seq: usize,
hidden_size: usize,
pooling: Pooling,
) -> Vec<Vec<f32>>Expand description
Pool [batch, seq, hidden] hidden states into [batch, hidden] and L2-normalize.