diff --git a/fastembed/common/onnx_model.py b/fastembed/common/onnx_model.py index d357f2c1..3f99a210 100644 --- a/fastembed/common/onnx_model.py +++ b/fastembed/common/onnx_model.py @@ -156,6 +156,19 @@ def add_extra_session_options( if "enable_cpu_mem_arena" in extra_options: session_options.enable_cpu_mem_arena = extra_options["enable_cpu_mem_arena"] + def _check_output_finite(self, output: NumpyArray) -> None: + """Validate that the model output contains only finite numbers (no NaN or Inf). + + Raises: + RuntimeError: If non-finite values (NaN/Inf) are detected in the output. + """ + if not np.all(np.isfinite(output)): + providers = self.model.get_providers() if self.model is not None else None + raise RuntimeError( + f"Model '{self.model_name}' produced non-finite (NaN/Inf) embeddings " + f"with provider(s) '{providers}'." + ) + def load_onnx_model(self) -> None: raise NotImplementedError("Subclasses must implement this method") diff --git a/fastembed/text/builtin_sentence_embedding.py b/fastembed/text/builtin_sentence_embedding.py index 430ea7ca..bcf83370 100644 --- a/fastembed/text/builtin_sentence_embedding.py +++ b/fastembed/text/builtin_sentence_embedding.py @@ -51,7 +51,9 @@ def _post_process_onnx_output( def _run_model( self, onnx_input: dict[str, Any], onnx_output_names: list[str] | None = None ) -> NumpyArray: - return self.model.run(onnx_output_names, onnx_input)[1] # type: ignore[union-attr] + model_output = self.model.run(onnx_output_names, onnx_input)[1] # type: ignore[union-attr] + self._check_output_finite(model_output) + return model_output class BuiltinSentenceEmbeddingWorker(OnnxTextEmbeddingWorker): diff --git a/fastembed/text/onnx_text_model.py b/fastembed/text/onnx_text_model.py index 10a4aa17..cf6b11f7 100644 --- a/fastembed/text/onnx_text_model.py +++ b/fastembed/text/onnx_text_model.py @@ -105,7 +105,9 @@ def onnx_embed( def _run_model( self, onnx_input: dict[str, Any], onnx_output_names: list[str] | None = None ) -> NumpyArray: - return self.model.run(onnx_output_names, onnx_input)[0] # type: ignore[union-attr] + model_output = self.model.run(onnx_output_names, onnx_input)[0] # type: ignore[union-attr] + self._check_output_finite(model_output) + return model_output def _embed_documents( self, diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 8744b617..dad79acb 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -308,3 +308,71 @@ def test_token_count(model_cache, model_name) -> None: doc_token_count = model.token_count(documents) assert first_doc_token_count + second_doc_token_count == doc_token_count assert doc_token_count == model.token_count(documents, batch_size=1) + + +# --------------------------------------------------------------------------- +# Regression tests for issue #688 - non-finite embedding guard +# --------------------------------------------------------------------------- +# All four tests use monkeypatch to replace the underlying ort.InferenceSession.run +# so they are deterministic, fast, and independent of any real model download or +# specific ONNX Runtime version. +# --------------------------------------------------------------------------- + +_BGE_SMALL = "BAAI/bge-small-en-v1.5" + + +def _load_bge_small() -> TextEmbedding: + """Load a small, widely-cached model. Uses lazy_load=False to ensure the + ONNX session is initialised before monkeypatching.""" + return TextEmbedding(model_name=_BGE_SMALL, lazy_load=False, cuda=False) + + +def test_non_finite_nan_raises(monkeypatch) -> None: + """NaN values must raise RuntimeError with model name and provider (issue #688). + + Verifies the full diagnostic message so users can immediately identify which + model/provider combination produced the bad output. + """ + model = _load_bge_small() + monkeypatch.setattr( + model.model.model, + "run", + lambda *args, **kwargs: [np.full((1, 5, 384), np.nan, dtype=np.float32)], + ) + with pytest.raises( + RuntimeError, + match=r"Model 'BAAI/bge-small-en-v1\.5' produced non-finite.*CPUExecutionProvider", + ): + list(model.embed(["test"])) + + +def test_non_finite_pos_inf_raises(monkeypatch) -> None: + """+Inf values in ONNX output must raise RuntimeError (issue #688).""" + model = _load_bge_small() + monkeypatch.setattr( + model.model.model, + "run", + lambda *args, **kwargs: [np.full((1, 5, 384), np.inf, dtype=np.float32)], + ) + with pytest.raises(RuntimeError, match="non-finite"): + list(model.embed(["test"])) + + +def test_non_finite_neg_inf_raises(monkeypatch) -> None: + """-Inf values in ONNX output must raise RuntimeError (issue #688).""" + model = _load_bge_small() + monkeypatch.setattr( + model.model.model, + "run", + lambda *args, **kwargs: [np.full((1, 5, 384), -np.inf, dtype=np.float32)], + ) + with pytest.raises(RuntimeError, match="non-finite"): + list(model.embed(["test"])) + + +def test_non_finite_finite_output_succeeds() -> None: + """Sanity check: finite ONNX output must not raise (issue #688).""" + model = _load_bge_small() + results = list(model.embed(["test"])) + assert len(results) == 1 + assert np.all(np.isfinite(results[0]))