jinaai

/

jina-embeddings-v3

like 600

Follow

Jina AI 428

The existing mean pooling code currently takes tensor as input, but the ONNX model expects numpy as input, so it's necessary to add a numpy version of mean pooling.

import numpy as np

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

# Run onnx model
...
outputs = session.run(None, inputs)[0]

embeddings = mean_pooling(outputs, input_text["attention_mask"])
norm = np.linalg.norm(embeddings, ord=2, axis=1, keepdims=True)
embeddings = embeddings / norm