From 2a859a5c338b744baa82d05548f8c6a6b360c771 Mon Sep 17 00:00:00 2001 From: wToTw123 <196762320+wToTw123@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:40:22 +0800 Subject: [PATCH 1/2] =?UTF-8?q?fix:=20=E9=98=BB=E6=AD=A2=E9=87=8D=E6=8E=92?= =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E5=A4=A7=E5=B0=8F=E5=86=99=E9=87=8D=E5=A4=8D?= =?UTF-8?q?=E6=B3=A8=E5=86=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../cross_encoder/text_cross_encoder.py | 3 +- tests/test_cross_encoder_registration.py | 47 +++++++++++++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) create mode 100644 tests/test_cross_encoder_registration.py diff --git a/fastembed/rerank/cross_encoder/text_cross_encoder.py b/fastembed/rerank/cross_encoder/text_cross_encoder.py index 6f98cb24a..095284ffd 100644 --- a/fastembed/rerank/cross_encoder/text_cross_encoder.py +++ b/fastembed/rerank/cross_encoder/text_cross_encoder.py @@ -143,9 +143,10 @@ def add_custom_model( size_in_gb: float = 0.0, additional_files: list[str] | None = None, ) -> None: + """Register a custom reranker under a case-insensitively unique model name.""" registered_models = cls._list_supported_models() for registered_model in registered_models: - if model == registered_model.model: + if model.lower() == registered_model.model.lower(): raise ValueError( f"Model {model} is already registered in CrossEncoderModel, if you still want to add this model, " f"please use another model name" diff --git a/tests/test_cross_encoder_registration.py b/tests/test_cross_encoder_registration.py new file mode 100644 index 000000000..47c056f9c --- /dev/null +++ b/tests/test_cross_encoder_registration.py @@ -0,0 +1,47 @@ +"""Keep custom cross-encoder registration consistent with model lookup.""" + +from pathlib import Path + +import pytest + +from fastembed.common.model_description import ModelSource +from fastembed.rerank.cross_encoder import TextCrossEncoder +from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder +from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder + + +@pytest.mark.parametrize("existing", ["builtin", "custom"]) +@pytest.mark.parametrize("case", [str, str.lower, str.upper, str.swapcase]) +def test_reject_duplicate_names_regardless_of_case(monkeypatch, existing: str, case) -> None: + """Names resolving to an existing model must not register an unreachable replacement.""" + monkeypatch.setattr(CustomTextCrossEncoder, "SUPPORTED_MODELS", []) + if existing == "builtin": + name = OnnxTextCrossEncoder._list_supported_models()[0].model + else: + name = "example/MyReranker" + TextCrossEncoder.add_custom_model(name, sources=ModelSource(hf="example/original")) + before = TextCrossEncoder.list_supported_models() + + with pytest.raises(ValueError, match="already registered"): + TextCrossEncoder.add_custom_model(case(name), sources=ModelSource(hf="example/other")) + + assert TextCrossEncoder.list_supported_models() == before + + +def test_distinct_custom_models_remain_available(monkeypatch, tmp_path: Path) -> None: + """Construct each distinct model using a case variant without loading ONNX weights.""" + monkeypatch.setattr(CustomTextCrossEncoder, "SUPPORTED_MODELS", []) + for name in ("example/FirstReranker", "example/SecondReranker"): + TextCrossEncoder.add_custom_model(name, sources=ModelSource(hf=name)) + encoder = TextCrossEncoder( + name.upper(), + cache_dir=str(tmp_path), + specific_model_path=str(tmp_path), + lazy_load=True, + local_files_only=True, + cuda=False, + ) + assert isinstance(encoder.model, CustomTextCrossEncoder) + assert encoder.model.model_description.sources.hf == name + + assert len(CustomTextCrossEncoder.SUPPORTED_MODELS) == 2 From a8dfe9490afb52991e33a1113b199490cb4f7155 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Mon, 12 Oct 2026 20:09:10 +0700 Subject: [PATCH 2/2] tests: move cross-encoder case-duplicate tests into test_custom_models --- tests/test_cross_encoder_registration.py | 47 ------------------------ tests/test_custom_models.py | 43 ++++++++++++++++------ 2 files changed, 31 insertions(+), 59 deletions(-) delete mode 100644 tests/test_cross_encoder_registration.py diff --git a/tests/test_cross_encoder_registration.py b/tests/test_cross_encoder_registration.py deleted file mode 100644 index 47c056f9c..000000000 --- a/tests/test_cross_encoder_registration.py +++ /dev/null @@ -1,47 +0,0 @@ -"""Keep custom cross-encoder registration consistent with model lookup.""" - -from pathlib import Path - -import pytest - -from fastembed.common.model_description import ModelSource -from fastembed.rerank.cross_encoder import TextCrossEncoder -from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder -from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder - - -@pytest.mark.parametrize("existing", ["builtin", "custom"]) -@pytest.mark.parametrize("case", [str, str.lower, str.upper, str.swapcase]) -def test_reject_duplicate_names_regardless_of_case(monkeypatch, existing: str, case) -> None: - """Names resolving to an existing model must not register an unreachable replacement.""" - monkeypatch.setattr(CustomTextCrossEncoder, "SUPPORTED_MODELS", []) - if existing == "builtin": - name = OnnxTextCrossEncoder._list_supported_models()[0].model - else: - name = "example/MyReranker" - TextCrossEncoder.add_custom_model(name, sources=ModelSource(hf="example/original")) - before = TextCrossEncoder.list_supported_models() - - with pytest.raises(ValueError, match="already registered"): - TextCrossEncoder.add_custom_model(case(name), sources=ModelSource(hf="example/other")) - - assert TextCrossEncoder.list_supported_models() == before - - -def test_distinct_custom_models_remain_available(monkeypatch, tmp_path: Path) -> None: - """Construct each distinct model using a case variant without loading ONNX weights.""" - monkeypatch.setattr(CustomTextCrossEncoder, "SUPPORTED_MODELS", []) - for name in ("example/FirstReranker", "example/SecondReranker"): - TextCrossEncoder.add_custom_model(name, sources=ModelSource(hf=name)) - encoder = TextCrossEncoder( - name.upper(), - cache_dir=str(tmp_path), - specific_model_path=str(tmp_path), - lazy_load=True, - local_files_only=True, - cuda=False, - ) - assert isinstance(encoder.model, CustomTextCrossEncoder) - assert encoder.model.model_description.sources.hf == name - - assert len(CustomTextCrossEncoder.SUPPORTED_MODELS) == 2 diff --git a/tests/test_custom_models.py b/tests/test_custom_models.py index 38751bd62..f6bc36327 100644 --- a/tests/test_custom_models.py +++ b/tests/test_custom_models.py @@ -284,6 +284,21 @@ def test_custom_text_model_lookup_is_case_insensitive(): assert model.model._normalization is True +def test_custom_cross_encoder_lookup_is_case_insensitive(): + model_name = "Org/Model" + + TextCrossEncoder.add_custom_model( + model_name, + sources=ModelSource(hf="artificial"), + size_in_gb=0.1, + ) + + model = TextCrossEncoder("org/model", lazy_load=True, specific_model_path="./") + + assert isinstance(model.model, CustomTextCrossEncoder) + assert model.model.model_description.sources.hf == "artificial" + + def test_do_not_add_existing_model(): existing_base_model = "sentence-transformers/all-MiniLM-L6-v2" custom_model_name = "intfloat/multilingual-e5-small" @@ -322,12 +337,13 @@ def test_do_not_add_existing_cross_encoder(): existing_base_model = "Xenova/ms-marco-MiniLM-L-6-v2" custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2" - with pytest.raises(ValueError, match=f"Model {existing_base_model} is already registered"): - TextCrossEncoder.add_custom_model( - existing_base_model, - sources=ModelSource(hf=existing_base_model), - size_in_gb=0.08, - ) + for name in (existing_base_model, existing_base_model.lower()): + with pytest.raises(ValueError, match=f"Model {name} is already registered"): + TextCrossEncoder.add_custom_model( + name, + sources=ModelSource(hf=existing_base_model), + size_in_gb=0.08, + ) TextCrossEncoder.add_custom_model( custom_model_name, @@ -335,9 +351,12 @@ def test_do_not_add_existing_cross_encoder(): size_in_gb=0.08, ) - with pytest.raises(ValueError, match=f"Model {custom_model_name} is already registered"): - TextCrossEncoder.add_custom_model( - custom_model_name, - sources=ModelSource(hf=custom_model_name), - size_in_gb=0.08, - ) + for name in (custom_model_name, custom_model_name.upper()): + with pytest.raises(ValueError, match=f"Model {name} is already registered"): + TextCrossEncoder.add_custom_model( + name, + sources=ModelSource(hf=custom_model_name), + size_in_gb=0.08, + ) + + assert len(CustomTextCrossEncoder.SUPPORTED_MODELS) == 1