diff --git a/model2vec/model.py b/model2vec/model.py index 324062c..a13ab46 100644 --- a/model2vec/model.py +++ b/model2vec/model.py @@ -476,8 +476,9 @@ 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 - np.divide(out, norm, out=out) + out32 = out.astype(np.float32) + norm = np.linalg.norm(out32, axis=1, keepdims=True) + 1e-32 + return (out32 / norm).astype(out.dtype) return out diff --git a/tests/test_model.py b/tests/test_model.py index df9f25e..b7a062a 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -53,6 +53,18 @@ def test_encode_single_sentence_empty( assert np.all(encoded == 0) +def test_encode_single_sentence_empty_float16( + mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str] +) -> None: + """Test encoding of a single empty sentence with float16 embeddings.""" + model = StaticModel(vectors=mock_vectors.astype(np.float16), tokenizer=mock_tokenizer, config=mock_config) + model.normalize = True + encoded = model.encode("") + assert not np.isnan(encoded).any() + assert np.all(encoded == 0) + assert encoded.dtype == model.embedding.dtype + + def test_encode_multiple_sentences( mock_vectors: np.ndarray, mock_tokenizer: Tokenizer, mock_config: dict[str, str] ) -> None: