Add onnx mean pool function

README.md

@@ -25206,6 +25206,15 @@ import onnxruntime
 import numpy as np
 from transformers import AutoTokenizer, PretrainedConfig
 
+# Mean pool function
+def mean_pooling(model_output: np.ndarray, attention_mask: np.ndarray):
+    token_embeddings = model_output
+    input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
+    input_mask_expanded = np.broadcast_to(input_mask_expanded, token_embeddings.shape)
+    sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
+    sum_mask = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
+    return sum_embeddings / sum_mask
+
 # Load tokenizer and model config
 tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embeddings-v3')
 config = PretrainedConfig.from_pretrained('jinaai/jina-embeddings-v3')
@@ -25230,6 +25239,9 @@ inputs = {
 outputs = session.run(None, inputs)[0]
 
 # Apply mean pooling to 'outputs' to get a single representation of each text
+embeddings = mean_pooling(outputs, input_text["attention_mask"])
+norm = np.linalg.norm(embeddings, ord=2, axis=1, keepdims=True)
+embeddings = embeddings / norm
 ```
 
 </p>