[Bug]: Mean-pooling text models return float64 instead of float32 #752

What happened?

Text models that use mean pooling return float64 embeddings, even though their ONNX models output float32. Models with other pooling, such as BAAI/bge-small-en-v1.5, return float32.

The values are right (they match float32 up to rounding), but the type is wrong. For every affected model this means:

• Twice the memory. 1M embeddings from all-MiniLM-L6-v2 (384 dims) take 3.1 GB instead of 1.5 GB.
• Wrong raw bytes. embedding.tobytes() writes 8 bytes per value, so a reader that expects float32 gets garbage.
• Silent upcasts. Stacking them with float32 embeddings from another model turns the whole array into float64.
• Failed type checks. Code or schemas that expect float32 reject them.

Cause. mean_pooling multiplies the float32 token embeddings by the int64 attention mask, and NumPy promotes the result to float64. For built-in models, a final .astype(np.float32) used to hide this; #492 removed it, so they return float64 since v0.6.1. Custom models added with PoolingType.MEAN have returned float64 since #479 added them.

Affected on main (a56e1e1): the 16 built-in models below, and custom models added with PoolingType.MEAN.

What is the expected behaviour?

Embeddings keep the dtype the model outputs (float32 here), as in v0.6.0 and as #492 intended.

Proposed fix. Cast the result back to the model's dtype in the three mean-pooling paths, after normalize():

# fastembed/text/pooled_normalized_embedding.py
-        return normalize(self.mean_pooling(embeddings, attn_mask))
+        return normalize(self.mean_pooling(embeddings, attn_mask)).astype(embeddings.dtype, copy=False)

# fastembed/text/pooled_embedding.py
-        return self.mean_pooling(embeddings, attn_mask)
+        return self.mean_pooling(embeddings, attn_mask).astype(embeddings.dtype, copy=False)

# fastembed/text/custom_text_embedding.py
-        return self._normalize(self._pool(output.model_output, output.attention_mask))
+        embeddings = self._normalize(self._pool(output.model_output, output.attention_mask))
+        if np.issubdtype(output.model_output.dtype, np.floating):  # don't truncate integer outputs
+            return embeddings.astype(output.model_output.dtype, copy=False)
+        return embeddings

The cast must come after normalize(): for a float16 model, casting first can overflow the norm (1,024 values of 10.0 square-sum to 102,400, above the float16 max of 65,504), and the embedding becomes all zeros. With this change, values move only by float32 rounding (at most 1.4e-07 against v0.8.1 on all-MiniLM-L6-v2, paraphrase-multilingual-MiniLM-L12-v2, nomic-embed-text-v1 and gte-base). Happy to open a PR with tests.

A minimal reproducible example

from fastembed import TextEmbedding

for name in ["sentence-transformers/all-MiniLM-L6-v2", "BAAI/bge-small-en-v1.5"]:
    embedding = next(iter(TextEmbedding(name).embed(["hello world"])))
    print(name, embedding.dtype)

What Python version are you on? e.g. python --version

Python 3.12, installed with uv (pip). NumPy 2.5.3, onnxruntime 1.30.0.

FastEmbed version

v0.8.1, and main at a56e1e1. Compared with v0.6.0 and v0.6.1.

What os are you seeing the problem on?

MacOS

Relevant stack traces and/or logs

version   all-MiniLM-L6-v2   bge-small-en-v1.5
0.6.0     float32            float32
0.6.1     float64            float32
0.8.1     float64            float32