Merge pull request #4 from arrmansa/main

Fixed dimension of normalized embeddings
davidberenstein1957 · Nov 14, 2022 · aaff8fd · aaff8fd
2 parents aa90a62 + fdd0b97
commit aaff8fd
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/fast_sentence_transformers/FastSentenceTransformer.py b/fast_sentence_transformers/FastSentenceTransformer.py
@@ -289,7 +289,8 @@ def encode(
                 embeddings = embeddings.cpu().detach().numpy()
 
             if normalize_embeddings:
-                embeddings = np.linalg.norm(embeddings, ord=2, axis=1, keepdims=False)
+                norms = np.linalg.norm(embeddings, ord=2, axis=1, keepdims=True)
+                embeddings = embeddings/np.where(norms<1e-12, 1e-12, norms)
 
             all_embeddings.extend(embeddings)