[jaysheeldodia]

Fixes #752

Models that use mean pooling have been returning float64 embeddings, even though the ONNX model outputs float32. The values are correct, but the embeddings take twice the memory. tobytes() also writes the wrong byte size, and anything expecting float32 rejects them.

Why it happened

mean_pooling multiplies the float32 embeddings by the int64 attention mask, and NumPy turns the result into float64. A final .astype(np.float32) used to hide this, but #492 removed it. Custom mean-pooling models from #479 had the same problem.

The fix

The result is now converted back to the model's output format in PooledEmbedding, PooledNormalizedEmbedding and CustomTextEmbedding. For custom models this only happens for float outputs, so integer outputs aren't cut off. The conversion comes after normalize(), because converting a float16 result first can overflow the norm and turn the embedding into all zeros.

The values don't change. The new output is exactly the old output in float32, with no speed difference and half the memory.

Tests

• test_embedding and test_mock_add_custom_models now also check the output format.
• New test_mean_pooling_keeps_model_dtype checks the format and the values for both built-in mean-pooling classes, with float32 and float16 input.
• New test_custom_mean_pooling_keeps_float16_without_overflow checks that a float16 custom model with large values still gets the right result.

All Submissions:

• Have you followed the guidelines in our Contributing document?
• Have you checked to ensure there aren't other open Pull Requests for the same update/change?

New Feature Submissions:

• Does your submission pass the existing tests?
• Have you added tests for your feature?
• Have you installed pre-commit with pip3 install pre-commit and set up hooks with pre-commit install?

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: 1d7e4208-13d0-436c-bc17-d9c2b9383b3b

📥 Commits

Reviewing files that changed from the base of the PR and between 4b8dd8a and 2b752db.

📒 Files selected for processing (3)

• fastembed/text/custom_text_embedding.py
• tests/test_custom_models.py
• tests/test_text_onnx_embeddings.py

🚧 Files skipped from review as they are similar to previous changes (3)

• fastembed/text/custom_text_embedding.py
• tests/test_text_onnx_embeddings.py
• tests/test_custom_models.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 6 remain after this review.

---

📝 Walkthrough

Walkthrough

Mean-pooling post-processing now casts embeddings back to the model output dtype after pooling and normalization. The custom embedding path casts only floating-point outputs. Tests check float32 and float16 dtype preservation and verify normalized values for float16 inputs.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix · Severity of issue fixed: Medium

Merge Risk: ⚪ Minimal · up to 2b752

The change preserves model output types without introducing an identified numerical regression. It is mergeable subject to normal test checks.

Security Architecture Review

Security architecture risk: ⚪ Minimal · up to 4b8dd

The inspected changes restore floating-point output precision after pooling without adding access, privileges, or external side effects. No material security risk was identified in these changes.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

• inferred — The demonstrated effect is on embedding representations returned to callers. The inspected changes do not expand authority or introduce a new sensitive sink; any downstream security-sensitive use of those embeddings is outside the established evidence.

Trust Boundaries and Controls

• observed — Both built-in pooled paths still reject a missing attention mask. Custom mean pooling also rejects a missing mask before dtype restoration; the added cast does not bypass these preconditions.

Resilience and Maintainability Implications

• inferred — The inspected mean-pooling transition is a computation followed by a return, not a shared-state update. Failure before return does not leave a new registry mutation or persistent intermediate state requiring rollback.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 36.36% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 11 functions across 5 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main change: preserving float32 output for mean pooling models. It is concise and related to the changeset.
Description check	✅ Passed	The description accurately explains the float64 promotion issue, the dtype-preserving fix, normalization order, affected classes, and regression tests.
Linked Issues check	✅ Passed	The PR satisfies the coding requirements in issue #752. PooledEmbedding and PooledNormalizedEmbedding cast mean-pooled results to the model output dtype after normalization where applicable. `Cust…
Out of Scope Changes check	✅ Passed	The source changes address the dtype behavior described in issue #752. The test changes verify the affected built-in and custom pooling paths. No unrelated functional changes appear in the reviewed ch…

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[joein]

Hey @jaysheeldodia
Thank you for your contribution!

Unfortunately, I didn't see the PR and implemented the fix on my own in #752

[jaysheeldodia]

Hey @joein
I'm happy to contribute.