fix: normalize custom token embeddings along the embedding axis

[ktz03]

When a custom ONNX model returns token embeddings with pooling=PoolingType.DISABLED and normalization=True, normalization runs across the token axis. A single token vector [3, 4] becomes [1, 1] instead of [0.6, 0.8]; repeating the token also changes its normalized vector. This affects custom token-level embedding exports used for multi-vector retrieval.

Normalize along the final embedding axis. This preserves the existing behavior for two-dimensional sentence outputs. Add regressions for token and sentence outputs, enabled/disabled normalization, zero vectors, repeated tokens, output dtypes, and input preservation.

The issue was found through code inspection and reproduced through TextEmbedding.add_custom_model, construction, and embed, using an actual local CPU ONNX lookup graph and tokenizer. The fixture emits known vectors and checks independent per-vector mathematical expectations; it requires no pretrained model download. Before the fix, "hello" produces [1, 1] and "hello hello" produces approximately [0.7071, 0.7071] for each token. After the fix, both inputs produce [0.6, 0.8] for each hello token.

Validation on Windows, Python 3.14.3, FastEmbed 0.9.0, NumPy 2.5.3, and ONNX Runtime 1.31.0 CPU:

• pytest tests/test_custom_models.py -k 'mock_add or lookup or existing or disabled_pooling' -q: 16 passed, including 12 new postprocessing cases. The unmodified baseline fails the three new normalized token-output cases and passes the other nine.
• pytest tests/test_common.py tests/test_image_transform.py tests/test_parallel_processor.py tests/test_postprocess.py -k 'not single_input' -q: 24 passed; one model-loading case deselected.
• A separate public-API ONNX probe checks 12 cases: token/sentence output, normalization enabled/disabled, single/repeated tokens, mixed-length batches, and zero vectors. All pass with the fix; three token-normalization cases fail on the baseline.
• Independent public-API verification with actual float16 and float32 CPU ONNX outputs passed 72 cases; the baseline failed 18 token-normalization cases and passed 54 normal controls.
• Repository mypy command: no issues in 66 source files. pyright tests/type_stub.py: zero errors or warnings.
• Pinned Ruff 0.3.4 pre-commit lint and formatting hooks passed.

No full pretrained-model suite, GPU execution, retrieval-quality benchmark, or full CI matrix was run locally.

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?

[coderabbitai]

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: 139f616b-80a5-4842-9ef7-6042c9614a82

📥 Commits

Reviewing files that changed from the base of the PR and between 539499b and 33fb600.

📒 Files selected for processing (2)

• fastembed/text/custom_text_embedding.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; 7 remain after this review.

---

📝 Walkthrough

Walkthrough

When normalization is enabled, CustomTextEmbedding._normalize now passes dim=-1 to normalize. When normalization is disabled, it returns the embeddings unchanged. Parameterized tests cover output modes, floating-point dtypes, normalization settings, zero vectors, dtype retention, and input immutability.

Priority: ⬇️ Low

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

Change: Bug fix

Merge Risk: ⚪ Minimal · up to 33fb6

No actionable merge-blocking issue is identified; this change appears mergeable after normal checks.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 16.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check		The title clearly and concisely describes the main change: normalizing custom token embeddings along the embedding axis.
Description check		The description directly explains the normalization bug, the fix, regression coverage, and validation results for the changeset.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

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 @ktz03

Thanks for the detailed report and the repro.

PoolingType.DISABLED is meant for models that already return a sentence embedding, so the ONNX output is expected to have shape (batch, dim). For that shape, dim=1 and dim=-1 give the same result, so normalization is already correct for the supported case.

Per-token output (batch, seq, dim) isn't supported through TextEmbedding.

That said, nothing documents or enforces this today, so your setup passed through silently. We'll make it explicit separately, so per-token output with DISABLED fails with a clear error. Closing this one, but thanks again for digging into it.