mirror of
https://github.com/qdrant/fastembed.git
synced 2026-07-23 11:20:51 -05:00
331207976e
* MiniLM fix * Added MiniLM to text embedding Fixed MiniLM source destination Black + isort for repo * Fixed model all-MiniLM-L6-v2 description Recomputed canonical vector for all-MiniLM-L6-v2 in test --------- Co-authored-by: d.rudenko <dimitriyrudenk@gmail.com>
34 lines
1.0 KiB
Python
34 lines
1.0 KiB
Python
import numpy as np
|
|
import onnx
|
|
import onnxruntime
|
|
from transformers import AutoTokenizer
|
|
|
|
model_id = "sentence-transformers/paraphrase-MiniLM-L6-v2"
|
|
output_dir = f"models/{model_id.replace('/', '_')}"
|
|
model_kwargs = {"output_attentions": True, "return_dict": True}
|
|
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
|
|
|
model_path = f"{output_dir}/model.onnx"
|
|
onnx_model = onnx.load(model_path)
|
|
ort_session = onnxruntime.InferenceSession(model_path)
|
|
text = "This is a test sentence"
|
|
tokenizer_output = tokenizer(text, return_tensors="np")
|
|
input_ids = tokenizer_output["input_ids"]
|
|
attention_mask = tokenizer_output["attention_mask"]
|
|
print(attention_mask)
|
|
# Prepare the input
|
|
input_ids = np.array(input_ids).astype(
|
|
np.int64
|
|
) # Replace your_input_ids with actual input data
|
|
|
|
# Run the ONNX model
|
|
outputs = ort_session.run(
|
|
None, {"input_ids": input_ids, "attention_mask": attention_mask}
|
|
)
|
|
|
|
# Get the attention weights
|
|
attentions = outputs[-1]
|
|
|
|
# Print the attention weights for the first layer and first head
|
|
print(attentions[0][0])
|