fix: look up custom model postprocessing by resolved name

[Ramnath0521]

Fixes #650.

What was wrong

CustomTextEmbedding.__init__ looked up its postprocessing config with the model_name argument:

self._pooling = self.POSTPROCESSING_MAPPING[model_name].pooling
self._normalization = self.POSTPROCESSING_MAPPING[model_name].normalization

But POSTPROCESSING_MAPPING is keyed by model_description.model — the name the model was registered under (add_model) — while model lookup is case-insensitive (_get_model_description compares model_name.lower() == model.model.lower()). So the caller's spelling need not match the registered key.

Reproduction

On main:

from fastembed.common.model_description import PoolingType, ModelSource
from fastembed.text.text_embedding import TextEmbedding

TextEmbedding.add_custom_model(
    "Org/Model", pooling=PoolingType.MEAN, normalization=True,
    sources=ModelSource(hf="intfloat/multilingual-e5-small"), dim=384,
)
TextEmbedding("org/model", lazy_load=True)

  File "fastembed/text/custom_text_embedding.py", line 53, in __init__
    self._pooling = self.POSTPROCESSING_MAPPING[model_name].pooling
KeyError: 'org/model'

Resolution itself is fine — TextEmbedding._get_model_description("org/model").model correctly returns "Org/Model". Only the mapping lookup uses the raw argument.

The change

self.model_description is already set by the base class from that same case-insensitive matcher, so .model is the canonical registered name. The lookup now uses it, and reads the entry once instead of twice.

Tests

Two cases in tests/test_custom_models.py:

Test	Purpose
test_custom_model_postprocessing_lookup_is_case_insensitive	register Org/Model, instantiate org/model — fails on main with the KeyError above
test_custom_model_postprocessing_lookup_with_matching_case_still_works	control — exact casing worked before and must keep working, including a different pooling/normalization pair so the assertion is not vacuous

Both use lazy_load=True, so they exercise the constructor without downloading weights.

One extra change, called out

The autouse restore_custom_models_fixture reset SUPPORTED_MODELS but not POSTPROCESSING_MAPPING, so a registration leaked into every later test in the file. Adding a test that registers a model made that visible, so the fixture now clears both. Happy to split it out if you would rather it landed separately.

Checks run locally

Windows 11, Python 3.12:

• pytest tests/test_custom_models.py — 7 passed (5 existing + 2 new)
• mypy fastembed --disallow-incomplete-defs --disallow-untyped-defs --disable-error-code=import-untyped — 5 errors, all pre-existing in vocab_resolver.py, preprocessor_utils.py and colbert.py; none in the changed file, and identical on an unmodified main
• ruff check on both changed files — 5 findings, all pre-existing (UP035/I001 import style), byte-identical on an unmodified main; ruff format --check — already formatted

I did not run the full pytest tests/ here: it downloads ONNX weights for every supported model and takes ~20 hours on this machine. This change touches only the custom-model constructor path, which the tests above cover directly.

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?
• Have you added tests for your feature?
• pre-commit — not installed locally; ran ruff check and ruff format --check directly instead

---

Authored by Claude (an AI coding agent) on the account owner's machine and with their authorization; they have reviewed the change and confirmed it. The reproduction, the failing-test-first sequence and every check above were genuinely executed here rather than asserted. Flagging the AI authorship plainly rather than leaving it to be inferred — happy to take any correction in review.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Team

Run ID: 328820b9-f4e0-45ef-93b7-61a74b13e7d3

📥 Commits

Reviewing files that changed from the base of the PR and between a34e7bc and bbdad26.

📒 Files selected for processing (2)

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

Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review.

---

📝 Walkthrough

Walkthrough

CustomTextEmbedding now retrieves postprocessing configuration with the canonical model name from the resolved model description. Tests reset custom-model mappings and verify both case-insensitive and matching-case lookups for pooling and normalization settings.

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

Merge Risk: ⚪ Minimal · up to bbdad

Custom embedding models now retain their registered pooling and normalization settings when instantiated with different casing. The supported cases are covered, with no remaining merge-blocking risk identified.

Suggested reviewers: joein

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% 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	✅ Passed	The title clearly and concisely describes the primary fix: using the resolved model name for custom model postprocessing lookup.
Description check	✅ Passed	The description directly explains the case-sensitivity bug, the implementation change, the tests, and validation results.
Linked Issues check	✅ Passed	The change satisfies issue #650 by using self.model_description.model for postprocessing lookup and adding tests for case-insensitive and exact-case construction.
Out of Scope Changes check	✅ Passed	The changes remain within scope. The fixture reset supports test isolation for the new custom-model tests and is directly related to the fix.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.}