fix: Normalize the output of OnnxEmbeddingFunction. (#21)

Signed-off-by: wxywb <[email protected]>
milvus-io · May 17, 2024 · 3ee078f · 3ee078f
1 parent c8a11db
commit 3ee078f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/milvus_model/dense/onnx.py b/milvus_model/dense/onnx.py
@@ -43,7 +43,8 @@ def _to_embedding(self, data: str, **_):
         ort_outputs = self.ort_session.run(None, ort_inputs)
         ort_feat = ort_outputs[0]
         emb = self._post_proc(ort_feat, ort_inputs["attention_mask"])
-        return emb.flatten()
+        emb =  emb.flatten()
+        return emb / np.linalg.norm(emb)
 
     def _post_proc(self, token_embeddings, attention_mask):
         input_mask_expanded = (