fastembed

1
import numpy as np
2
import onnx
3
import onnxruntime
4
from transformers import AutoTokenizer
5

6
model_id = "sentence-transformers/paraphrase-MiniLM-L6-v2"
7
output_dir = f"models/{model_id.replace('/', '_')}"
8
model_kwargs = {"output_attentions": True, "return_dict": True}
9
tokenizer = AutoTokenizer.from_pretrained(model_id)
10

11
model_path = f"{output_dir}/model.onnx"
12
onnx_model = onnx.load(model_path)
13
ort_session = onnxruntime.InferenceSession(model_path)
14
text = "This is a test sentence"
15
tokenizer_output = tokenizer(text, return_tensors="np")
16
input_ids = tokenizer_output["input_ids"]
17
attention_mask = tokenizer_output["attention_mask"]
18
print(attention_mask)
19
# Prepare the input
20
input_ids = np.array(input_ids).astype(np.int64)  # Replace your_input_ids with actual input data
21

22
# Run the ONNX model
23
outputs = ort_session.run(None, {"input_ids": input_ids, "attention_mask": attention_mask})
24

25
# Get the attention weights
26
attentions = outputs[-1]
27

28
# Print the attention weights for the first layer and first head
29
print(attentions[0][0])
30