EmbeddedLLM/ms-marco-TinyBERT-L-2-v2-onnx-o3-cpu

from itertools import product
from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer

sentences = [
"The llama (/ˈlɑːmə/) (Lama glama) is a domesticated South American camelid.",
"The alpaca (Lama pacos) is a species of South American camelid mammal.",
"The vicuña (Lama vicugna) (/vɪˈkuːnjə/) is one of the two wild South American camelids.",
]
queries = ["What is a llama?", "What is a harimau?", "How to fly a kite?"]
pairs = list(product(queries, sentences))

model_name = "EmbeddedLLM/ms-marco-TinyBERT-L-2-v2-onnx-o3-cpu"
device = "cpu"
provider = "CPUExecutionProvider"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = ORTModelForSequenceClassification.from_pretrained(
model_name, use_io_binding=True, provider=provider, device_map=device
)
inputs = tokenizer(
pairs,
padding=True,
truncation=True,
return_tensors="pt",
max_length=model.config.max_position_embeddings,
)
inputs = inputs.to(device)
scores = model(**inputs).logits.cpu().numpy()
# Sort most similar to least
pairs = sorted(zip(pairs, scores), key=lambda x: x[1], reverse=True)
for ps in pairs:
print(ps)