import onnxruntime
import numpy as np
from transformers import AutoTokenizer, PretrainedConfig


def onnx_embed(session: onnxruntime.InferenceSession, task_id: int):
	inputs = {
		'input_ids': input_text['input_ids'],
		'attention_mask': input_text['attention_mask'],
		'task_id': np.array(task_id, dtype=np.int64)
	}
	outputs = session.run(None, inputs)[0]
	embeddings = mean_pooling(outputs, input_text["attention_mask"])
	embeddings = embeddings / np.linalg.norm(embeddings, ord=2, axis=1, keepdims=True)
	return embeddings

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

docs = [
    "Hello World",
    "Follow the white rabbit."
]
model_name = 'jinaai/jina-embeddings-v3'
tokenizer = AutoTokenizer.from_pretrained(model_name)

config = PretrainedConfig.from_pretrained(model_name)

input_text = tokenizer(docs, return_tensors='np', padding=True, truncation=True)
input_ids = input_text['input_ids']
attention_mask = input_text['attention_mask']

model_path = 'models/models--jinaai--jina-embeddings-v3/snapshots/62a81741b58448ed8f691764cec7aa5d3c045e4c/onnx/model.onnx'
session = onnxruntime.InferenceSession(model_path)

for task_id in range(5):
	embeddings = onnx_embed(session, task_id)
	print([[round(value, 4) for value in e.tolist()[:5]] for e in embeddings])