From e69d3ff4be482f86fe23dd6f7fa46dd3a66167a8 Mon Sep 17 00:00:00 2001 From: tuanzirwar <1281227988@qq.com> Date: Thu, 1 Oct 2026 16:59:18 +0800 Subject: [PATCH 1/2] feat: select a custom text model ONNX output by name --- docs/Custom_Model_Outputs.md | 29 +++++ fastembed/text/custom_text_embedding.py | 11 +- fastembed/text/text_embedding.py | 3 + tests/test_custom_output_selection.py | 135 ++++++++++++++++++++++++ 4 files changed, 177 insertions(+), 1 deletion(-) create mode 100644 docs/Custom_Model_Outputs.md create mode 100644 tests/test_custom_output_selection.py diff --git a/docs/Custom_Model_Outputs.md b/docs/Custom_Model_Outputs.md new file mode 100644 index 000000000..4e2483f6d --- /dev/null +++ b/docs/Custom_Model_Outputs.md @@ -0,0 +1,29 @@ +# Selecting an output from a custom text model + +An ONNX graph can expose multiple outputs, such as token hidden states and a pooled +sentence embedding. By default FastEmbed uses the first output, preserving existing +custom model behavior. To select a different output, pass its exact ONNX name when +registering the model: + +```python +from fastembed import TextEmbedding +from fastembed.common.model_description import ModelSource, PoolingType + +TextEmbedding.add_custom_model( + model="my-org/my-embedding-model", + pooling=PoolingType.DISABLED, + normalization=True, + sources=ModelSource(hf="my-org/my-embedding-model"), + dim=384, + output_name="sentence_embedding", +) +model = TextEmbedding("my-org/my-embedding-model") +embeddings = list(model.embed(["Example document"])) +``` + +Use `onnxruntime.InferenceSession(...).get_outputs()` to inspect the graph's output +names. Choose `DISABLED` for an already pooled `[batch, dim]` output, or a token +pooling mode for a `[batch, tokens, dim]` output. Normalization runs after pooling. +The same selection applies to lazy loading and parallel workers. An unknown name +raises ONNX Runtime's normal invalid-output error when inference runs. This selects +one embedding output; it does not return all graph outputs. diff --git a/fastembed/text/custom_text_embedding.py b/fastembed/text/custom_text_embedding.py index c3d5cf5f0..7eb711559 100644 --- a/fastembed/text/custom_text_embedding.py +++ b/fastembed/text/custom_text_embedding.py @@ -20,6 +20,7 @@ class PostprocessingConfig: pooling: PoolingType normalization: bool + output_name: str | None = None class CustomTextEmbedding(OnnxTextEmbedding): @@ -54,6 +55,8 @@ def __init__( postprocessing_config = self.POSTPROCESSING_MAPPING[self.model_description.model] self._pooling = postprocessing_config.pooling self._normalization = postprocessing_config.normalization + if postprocessing_config.output_name is not None: + self.ONNX_OUTPUT_NAMES = [postprocessing_config.output_name] @classmethod def _list_supported_models(cls) -> list[DenseModelDescription]: @@ -108,10 +111,15 @@ def add_model( model_description: DenseModelDescription, pooling: PoolingType, normalization: bool, + output_name: str | None = None, ) -> None: + if output_name is not None and ( + not isinstance(output_name, str) or not output_name.strip() + ): + raise ValueError("output_name must be a non-empty string or None") cls.SUPPORTED_MODELS.append(model_description) cls.POSTPROCESSING_MAPPING[model_description.model] = PostprocessingConfig( - pooling=pooling, normalization=normalization + pooling=pooling, normalization=normalization, output_name=output_name ) @@ -135,6 +143,7 @@ def init_embedding( model_description, pooling=postprocessing_config.pooling, normalization=postprocessing_config.normalization, + output_name=postprocessing_config.output_name, ) return CustomTextEmbedding( model_name=model_name, diff --git a/fastembed/text/text_embedding.py b/fastembed/text/text_embedding.py index 8848389b2..a1d441edc 100644 --- a/fastembed/text/text_embedding.py +++ b/fastembed/text/text_embedding.py @@ -58,7 +58,9 @@ def add_custom_model( license: str = "", size_in_gb: float = 0.0, additional_files: list[str] | None = None, + output_name: str | None = None, ) -> None: + """注册自定义模型；output_name 选择送入池化的单个 ONNX 输出。""" registered_models = cls._list_supported_models() for registered_model in registered_models: if model.lower() == registered_model.model.lower(): @@ -80,6 +82,7 @@ def add_custom_model( ), pooling=pooling, normalization=normalization, + output_name=output_name, ) def __init__( diff --git a/tests/test_custom_output_selection.py b/tests/test_custom_output_selection.py new file mode 100644 index 000000000..4df8cef18 --- /dev/null +++ b/tests/test_custom_output_selection.py @@ -0,0 +1,135 @@ +import json + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper +from onnxruntime.capi.onnxruntime_pybind11_state import InvalidArgument +from tokenizers import Tokenizer, models, pre_tokenizers + +from fastembed import TextEmbedding +from fastembed.common.model_description import ModelSource, PoolingType +from fastembed.text.custom_text_embedding import CustomTextEmbedding + + +@pytest.fixture +def local_model(tmp_path, monkeypatch): + monkeypatch.setattr(CustomTextEmbedding, "SUPPORTED_MODELS", []) + monkeypatch.setattr(CustomTextEmbedding, "POSTPROCESSING_MAPPING", {}) + tokenizer = Tokenizer(models.WordLevel({"[PAD]": 0, "[UNK]": 1, "hello": 2})) + tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() + tokenizer.save(str(tmp_path / "tokenizer.json")) + (tmp_path / "tokenizer_config.json").write_text( + json.dumps({"model_max_length": 16, "pad_token": "[PAD]"}) + ) + (tmp_path / "config.json").write_text(json.dumps({"pad_token_id": 0})) + graph = helper.make_graph( + [ + helper.make_node("Cast", ["input_ids"], ["first"], to=TensorProto.FLOAT), + helper.make_node("Add", ["first", "offset"], ["sentence_embedding"]), + helper.make_node("Unsqueeze", ["sentence_embedding", "axis"], ["hidden_states"]), + ], + "two_outputs", + [helper.make_tensor_value_info("input_ids", TensorProto.INT64, ["batch", "seq"])], + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, ["batch", "seq"]), + helper.make_tensor_value_info( + "sentence_embedding", TensorProto.FLOAT, ["batch", "seq"] + ), + helper.make_tensor_value_info("hidden_states", TensorProto.FLOAT, ["batch", "seq", 1]), + ], + [ + helper.make_tensor("offset", TensorProto.FLOAT, [1], [10.0]), + helper.make_tensor("axis", TensorProto.INT64, [1], [2]), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)], ir_version=8) + onnx.checker.check_model(model) + onnx.save(model, tmp_path / "model.onnx") + return tmp_path + + +@pytest.mark.parametrize("parallel", [None, 2]) +@pytest.mark.parametrize("lazy_load", [False, True]) +@pytest.mark.parametrize("output_name,expected", [(None, 2.0), ("sentence_embedding", 12.0)]) +def test_output_selection_in_real_onnx_session( + local_model, parallel, lazy_load, output_name, expected +): + TextEmbedding.add_custom_model( + "local/output-test", + PoolingType.DISABLED, + False, + ModelSource(hf="local/output-test"), + dim=1, + model_file="model.onnx", + output_name=output_name, + ) + model = TextEmbedding( + "local/output-test", + specific_model_path=str(local_model), + cache_dir=str(local_model), + lazy_load=lazy_load, + threads=1, + cuda=False, + ) + vectors = np.stack(list(model.embed(["hello"] * 4, batch_size=2, parallel=parallel))) + np.testing.assert_array_equal(vectors, np.full((4, 1), expected, dtype=np.float32)) + + +def test_unknown_output_name_raises_from_onnx(local_model): + TextEmbedding.add_custom_model( + "local/missing", + PoolingType.DISABLED, + False, + ModelSource(hf="local/missing"), + dim=1, + model_file="model.onnx", + output_name="missing", + ) + model = TextEmbedding( + "local/missing", + specific_model_path=str(local_model), + cache_dir=str(local_model), + cuda=False, + ) + with pytest.raises(InvalidArgument, match="Invalid output name"): + list(model.embed("hello")) + + +@pytest.mark.parametrize("output_name", ["", " ", 123, []]) +def test_invalid_output_name_does_not_register_model(local_model, output_name): + with pytest.raises(ValueError, match="output_name"): + TextEmbedding.add_custom_model( + "local/invalid", + PoolingType.DISABLED, + False, + ModelSource(hf="local/invalid"), + dim=1, + output_name=output_name, + ) + assert CustomTextEmbedding.SUPPORTED_MODELS == [] + assert CustomTextEmbedding.POSTPROCESSING_MAPPING == {} + + +@pytest.mark.parametrize("pooling", [PoolingType.CLS, PoolingType.MEAN, PoolingType.LAST_TOKEN]) +@pytest.mark.parametrize("normalization", [False, True]) +def test_selected_output_uses_registered_postprocessing(local_model, pooling, normalization): + TextEmbedding.add_custom_model( + "local/pooled", + pooling, + normalization, + ModelSource(hf="local/pooled"), + dim=1, + model_file="model.onnx", + output_name="hidden_states", + ) + model = TextEmbedding( + "local/pooled", + specific_model_path=str(local_model), + cache_dir=str(local_model), + cuda=False, + ) + expected = 1.0 if normalization else 12.0 + np.testing.assert_allclose( + list(model.embed(["hello", "hello hello"])), [[expected], [expected]] + ) From 5cb3ebe5f7f0d524f6798dcf8abb7458710e7dbd Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Sun, 4 Oct 2026 01:55:46 +0700 Subject: [PATCH 2/2] fix: validate output_name on model load, simplify tests, drop extra docs page --- docs/Custom_Model_Outputs.md | 29 ----- fastembed/text/custom_text_embedding.py | 15 ++- fastembed/text/text_embedding.py | 1 - tests/test_custom_models.py | 37 +++++++ tests/test_custom_output_selection.py | 135 ------------------------ 5 files changed, 48 insertions(+), 169 deletions(-) delete mode 100644 docs/Custom_Model_Outputs.md delete mode 100644 tests/test_custom_output_selection.py diff --git a/docs/Custom_Model_Outputs.md b/docs/Custom_Model_Outputs.md deleted file mode 100644 index 4e2483f6d..000000000 --- a/docs/Custom_Model_Outputs.md +++ /dev/null @@ -1,29 +0,0 @@ -# Selecting an output from a custom text model - -An ONNX graph can expose multiple outputs, such as token hidden states and a pooled -sentence embedding. By default FastEmbed uses the first output, preserving existing -custom model behavior. To select a different output, pass its exact ONNX name when -registering the model: - -```python -from fastembed import TextEmbedding -from fastembed.common.model_description import ModelSource, PoolingType - -TextEmbedding.add_custom_model( - model="my-org/my-embedding-model", - pooling=PoolingType.DISABLED, - normalization=True, - sources=ModelSource(hf="my-org/my-embedding-model"), - dim=384, - output_name="sentence_embedding", -) -model = TextEmbedding("my-org/my-embedding-model") -embeddings = list(model.embed(["Example document"])) -``` - -Use `onnxruntime.InferenceSession(...).get_outputs()` to inspect the graph's output -names. Choose `DISABLED` for an already pooled `[batch, dim]` output, or a token -pooling mode for a `[batch, tokens, dim]` output. Normalization runs after pooling. -The same selection applies to lazy loading and parallel workers. An unknown name -raises ONNX Runtime's normal invalid-output error when inference runs. This selects -one embedding output; it does not return all graph outputs. diff --git a/fastembed/text/custom_text_embedding.py b/fastembed/text/custom_text_embedding.py index 7eb711559..22482a8ab 100644 --- a/fastembed/text/custom_text_embedding.py +++ b/fastembed/text/custom_text_embedding.py @@ -58,6 +58,17 @@ def __init__( if postprocessing_config.output_name is not None: self.ONNX_OUTPUT_NAMES = [postprocessing_config.output_name] + def load_onnx_model(self) -> None: + super().load_onnx_model() + # with eager loading this runs inside super().__init__(), before ONNX_OUTPUT_NAMES is set, + # so the output name is taken from the registered config + output_name = self.POSTPROCESSING_MAPPING[self.model_description.model].output_name + output_names = [output.name for output in self.model.get_outputs()] # type: ignore[union-attr] + if output_name is not None and output_name not in output_names: + raise ValueError( + f"Output {output_name!r} not found in the model, available outputs: {output_names}" + ) + @classmethod def _list_supported_models(cls) -> list[DenseModelDescription]: return cls.SUPPORTED_MODELS @@ -113,10 +124,6 @@ def add_model( normalization: bool, output_name: str | None = None, ) -> None: - if output_name is not None and ( - not isinstance(output_name, str) or not output_name.strip() - ): - raise ValueError("output_name must be a non-empty string or None") cls.SUPPORTED_MODELS.append(model_description) cls.POSTPROCESSING_MAPPING[model_description.model] = PostprocessingConfig( pooling=pooling, normalization=normalization, output_name=output_name diff --git a/fastembed/text/text_embedding.py b/fastembed/text/text_embedding.py index a1d441edc..14966383e 100644 --- a/fastembed/text/text_embedding.py +++ b/fastembed/text/text_embedding.py @@ -60,7 +60,6 @@ def add_custom_model( additional_files: list[str] | None = None, output_name: str | None = None, ) -> None: - """注册自定义模型；output_name 选择送入池化的单个 ONNX 输出。""" registered_models = cls._list_supported_models() for registered_model in registered_models: if model.lower() == registered_model.model.lower(): diff --git a/tests/test_custom_models.py b/tests/test_custom_models.py index 1cf47f7f1..c20d63f05 100644 --- a/tests/test_custom_models.py +++ b/tests/test_custom_models.py @@ -101,6 +101,43 @@ def test_text_custom_model_parallel_processing(): delete_model_cache(model.model._model_dir) +def test_text_custom_model_output_name(): + is_ci = os.getenv("CI") + custom_model_name = "custom/granite-embedding-small-english-r2" + + # onnx outputs are `last_hidden_state` and `sentence_embedding`, the builtin model uses the latter + TextEmbedding.add_custom_model( + custom_model_name, + pooling=PoolingType.DISABLED, + normalization=False, + sources=ModelSource(hf="onnx-community/granite-embedding-small-english-r2-ONNX"), + dim=384, + additional_files=["onnx/model.onnx_data"], + output_name="sentence_embedding", + ) + + model = TextEmbedding(custom_model_name) + builtin_model = TextEmbedding("ibm-granite/granite-embedding-small-english-r2") + docs = ["hello world", "flag embedding"] + expected = np.stack(list(builtin_model.embed(docs)), axis=0) + + embeddings = np.stack(list(model.embed(docs)), axis=0) + assert np.allclose(embeddings, expected, atol=1e-5) + + # workers re-register the model, so they have to receive the output name too + embeddings = np.stack(list(model.embed(docs, batch_size=1, parallel=2)), axis=0) + assert np.allclose(embeddings, expected, atol=1e-5) + + CustomTextEmbedding.POSTPROCESSING_MAPPING[custom_model_name] = PostprocessingConfig( + pooling=PoolingType.DISABLED, normalization=False, output_name="sentence_embeddings" + ) + with pytest.raises(ValueError, match="available outputs"): + TextEmbedding(custom_model_name) + + if is_ci: + delete_model_cache(model.model._model_dir) + + def test_cross_encoder_custom_model(): is_ci = os.getenv("CI") custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2" diff --git a/tests/test_custom_output_selection.py b/tests/test_custom_output_selection.py deleted file mode 100644 index 4df8cef18..000000000 --- a/tests/test_custom_output_selection.py +++ /dev/null @@ -1,135 +0,0 @@ -import json - -import numpy as np -import onnx -import pytest -from onnx import TensorProto, helper -from onnxruntime.capi.onnxruntime_pybind11_state import InvalidArgument -from tokenizers import Tokenizer, models, pre_tokenizers - -from fastembed import TextEmbedding -from fastembed.common.model_description import ModelSource, PoolingType -from fastembed.text.custom_text_embedding import CustomTextEmbedding - - -@pytest.fixture -def local_model(tmp_path, monkeypatch): - monkeypatch.setattr(CustomTextEmbedding, "SUPPORTED_MODELS", []) - monkeypatch.setattr(CustomTextEmbedding, "POSTPROCESSING_MAPPING", {}) - tokenizer = Tokenizer(models.WordLevel({"[PAD]": 0, "[UNK]": 1, "hello": 2})) - tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() - tokenizer.save(str(tmp_path / "tokenizer.json")) - (tmp_path / "tokenizer_config.json").write_text( - json.dumps({"model_max_length": 16, "pad_token": "[PAD]"}) - ) - (tmp_path / "config.json").write_text(json.dumps({"pad_token_id": 0})) - graph = helper.make_graph( - [ - helper.make_node("Cast", ["input_ids"], ["first"], to=TensorProto.FLOAT), - helper.make_node("Add", ["first", "offset"], ["sentence_embedding"]), - helper.make_node("Unsqueeze", ["sentence_embedding", "axis"], ["hidden_states"]), - ], - "two_outputs", - [helper.make_tensor_value_info("input_ids", TensorProto.INT64, ["batch", "seq"])], - [ - helper.make_tensor_value_info("first", TensorProto.FLOAT, ["batch", "seq"]), - helper.make_tensor_value_info( - "sentence_embedding", TensorProto.FLOAT, ["batch", "seq"] - ), - helper.make_tensor_value_info("hidden_states", TensorProto.FLOAT, ["batch", "seq", 1]), - ], - [ - helper.make_tensor("offset", TensorProto.FLOAT, [1], [10.0]), - helper.make_tensor("axis", TensorProto.INT64, [1], [2]), - ], - ) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)], ir_version=8) - onnx.checker.check_model(model) - onnx.save(model, tmp_path / "model.onnx") - return tmp_path - - -@pytest.mark.parametrize("parallel", [None, 2]) -@pytest.mark.parametrize("lazy_load", [False, True]) -@pytest.mark.parametrize("output_name,expected", [(None, 2.0), ("sentence_embedding", 12.0)]) -def test_output_selection_in_real_onnx_session( - local_model, parallel, lazy_load, output_name, expected -): - TextEmbedding.add_custom_model( - "local/output-test", - PoolingType.DISABLED, - False, - ModelSource(hf="local/output-test"), - dim=1, - model_file="model.onnx", - output_name=output_name, - ) - model = TextEmbedding( - "local/output-test", - specific_model_path=str(local_model), - cache_dir=str(local_model), - lazy_load=lazy_load, - threads=1, - cuda=False, - ) - vectors = np.stack(list(model.embed(["hello"] * 4, batch_size=2, parallel=parallel))) - np.testing.assert_array_equal(vectors, np.full((4, 1), expected, dtype=np.float32)) - - -def test_unknown_output_name_raises_from_onnx(local_model): - TextEmbedding.add_custom_model( - "local/missing", - PoolingType.DISABLED, - False, - ModelSource(hf="local/missing"), - dim=1, - model_file="model.onnx", - output_name="missing", - ) - model = TextEmbedding( - "local/missing", - specific_model_path=str(local_model), - cache_dir=str(local_model), - cuda=False, - ) - with pytest.raises(InvalidArgument, match="Invalid output name"): - list(model.embed("hello")) - - -@pytest.mark.parametrize("output_name", ["", " ", 123, []]) -def test_invalid_output_name_does_not_register_model(local_model, output_name): - with pytest.raises(ValueError, match="output_name"): - TextEmbedding.add_custom_model( - "local/invalid", - PoolingType.DISABLED, - False, - ModelSource(hf="local/invalid"), - dim=1, - output_name=output_name, - ) - assert CustomTextEmbedding.SUPPORTED_MODELS == [] - assert CustomTextEmbedding.POSTPROCESSING_MAPPING == {} - - -@pytest.mark.parametrize("pooling", [PoolingType.CLS, PoolingType.MEAN, PoolingType.LAST_TOKEN]) -@pytest.mark.parametrize("normalization", [False, True]) -def test_selected_output_uses_registered_postprocessing(local_model, pooling, normalization): - TextEmbedding.add_custom_model( - "local/pooled", - pooling, - normalization, - ModelSource(hf="local/pooled"), - dim=1, - model_file="model.onnx", - output_name="hidden_states", - ) - model = TextEmbedding( - "local/pooled", - specific_model_path=str(local_model), - cache_dir=str(local_model), - cuda=False, - ) - expected = 1.0 if normalization else 12.0 - np.testing.assert_allclose( - list(model.embed(["hello", "hello hello"])), [[expected], [expected]] - )