File size: 1,678 Bytes
8eeaf44 82737a1 8eeaf44 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
import onnxruntime as ort
import torch
import numpy as np
from transformers import AutoModel, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-small-en-v1.5')
model = AutoModel.from_pretrained('BAAI/bge-small-en-v1.5')
model.eval()
model_ort = ort.InferenceSession('./bge-small-en-v1.5.onnx',providers=['CPUExecutionProvider'])
sentences_1 = ["I really love math"]
sentences_2 = ["I pretty like mathematics"]
encoded_input1 = tokenizer(sentences_1, padding='max_length', max_length=512, truncation=True, return_tensors='pt')
encoded_input2 = tokenizer(sentences_2, padding='max_length', max_length=512, truncation=True, return_tensors='pt')
model_features_ort1 = model_ort.run(None,{'input_ids':np.array(encoded_input1.input_ids)})
model_features_ort2 = model_ort.run(None,{'input_ids':np.array(encoded_input2.input_ids)})
# Perform pooling. In this case, cls pooling.
embeddings_1 = model_features_ort1[0][:, 0]
embeddings_1 /= np.linalg.norm(embeddings_1, axis=1, keepdims=True)
embeddings_2 = model_features_ort2[0][:, 0]
embeddings_2 /= np.linalg.norm(embeddings_2, axis=1, keepdims=True)
similarity = embeddings_1 @ embeddings_2.T
print("similarity:",similarity)
with torch.no_grad():
model_output1 = model(**encoded_input1)
model_output2 = model(**encoded_input2)
embeddings_gt1 = model_output1[0].detach().cpu().numpy()[:, 0]
embeddings_gt1 /= np.linalg.norm(embeddings_gt1, axis=1, keepdims=True)
embeddings_gt2 = model_output2[0].detach().cpu().numpy()[:, 0]
embeddings_gt2 /= np.linalg.norm(embeddings_gt2, axis=1, keepdims=True)
similarity = embeddings_gt1 @ embeddings_gt2.T
print("gt similarity:",similarity) |