BAAI/bge-small-en-v1.5: Mean Pooling + Normalize

[niltonheck]

Hi everyone,

I recently had to use the BAAI/bge-small-en-v1.5 model for text embeddings and noticed that the Fastembed API for this model provides embeddings that are already normalized, without returning the attention mask, or allowing other pooling methods. As a result, I had to rely on internal functions to achieve the desired outcome.

Is there a specific reason why the API returns embeddings already normalized and doesn’t provide an option for mean pooling for example?

Am I overlooking something?

Here’s how I managed to:

• Obtain raw embeddings
• Perform mean pooling
• Normalize the pooled embeddings

def generate_embeddings(documents, batch_size):
    for batch in iter_batch(documents, batch_size):
        onnx_output = embedding_model.model.onnx_embed(batch)
        model_output = onnx_output.model_output

        attention_mask = onnx_output.attention_mask
        input_ids = onnx_output.input_ids

        for i, v in enumerate(model_output):
            mean_polled = torch.tensor(PooledEmbedding.mean_pooling(model_output[i], attention_mask[i]))
            yield documents[i], torch.nn.functional.normalize(mean_polled, p=2, dim=1)