From 69cd5d7de7b7a5724f0a4e07d7abb99cbd73be86 Mon Sep 17 00:00:00 2001
From: xblwh <64543322+xblwh@users.noreply.github.com>
Date: Tue, 15 Sep 2026 20:13:05 +0800
Subject: [PATCH] fix: normalize float16 embeddings without NaNs

---
 model2vec/model.py  |  3 ++-
 tests/test_model.py | 21 +++++++++++++++++++--
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/model2vec/model.py b/model2vec/model.py
index 324062c..a172aaa 100644
--- a/model2vec/model.py
+++ b/model2vec/model.py
@@ -476,7 +476,8 @@ def _encode_batch(self, sentences: Sequence[str], normalize: bool) -> np.ndarray
                     out[i] = emb.mean(axis=0)
 
         if normalize:
-            norm = np.linalg.norm(out, axis=1, keepdims=True) + 1e-32
+            norm_input = out.astype(np.float32) if out.dtype == np.float16 else out
+            norm = np.linalg.norm(norm_input, axis=1, keepdims=True) + 1e-32
             np.divide(out, norm, out=out)
 
         return out
diff --git a/tests/test_model.py b/tests/test_model.py
index df9f25e..003808b 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -42,11 +42,12 @@ def test_encode_single_sentence(
     assert encoded.shape == (2,)
 
 
+@pytest.mark.parametrize("dtype", ["float16", "float32", "float64"])
 def test_encode_single_sentence_empty(
-    mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str]
+    mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str], dtype: str
 ) -> None:
     """Test encoding of a single empty sentence."""
-    model = StaticModel(vectors=mock_vectors, tokenizer=mock_tokenizer, config=mock_config)
+    model = StaticModel(vectors=mock_vectors.astype(dtype), tokenizer=mock_tokenizer, config=mock_config)
     model.normalize = True
     encoded = model.encode("")
     assert not np.isnan(encoded).any()
@@ -159,6 +160,22 @@ def test_normalize(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_con
     np.testing.assert_almost_equal(normalized, expected)
 
 
+@pytest.mark.parametrize("dtype", ["float16", "float32", "float64"])
+@pytest.mark.parametrize("scale", [1e-4, 1.0, 1e4])
+def test_normalize_zero_and_nonzero_embeddings(
+    mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, dtype: str, scale: float
+) -> None:
+    """Keep zero vectors finite and normalize nonzero vectors across precisions and scales."""
+    model = StaticModel(vectors=(mock_vectors * scale).astype(dtype), tokenizer=mock_tokenizer, normalize=True)
+    with np.errstate(divide="raise", invalid="raise", over="raise"):
+        encoded = model.encode(["", "unknown", "word3", "word1 word2"])
+
+    assert encoded.dtype == np.dtype(dtype)
+    assert np.isfinite(encoded).all()
+    np.testing.assert_array_equal(encoded[:3], np.zeros((3, model.dim)))
+    assert np.linalg.norm(encoded[3].astype(np.float64)) == pytest.approx(1.0, abs=1e-3)
+
+
 def test_save_pretrained(
     tmp_path: Path, mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str]
 ) -> None: