bge-base-en-v1.5-onnx-o4-o2-gpu

ONNX Conversion of BAAI/bge-base-en-v1.5

• ONNX model for GPU with O4-O2 optimisation
• We exported the model with use_raw_attention_mask=True due to this issue

Usage

import torch.nn.functional as F
from optimum.onnxruntime import ORTModelForFeatureExtraction
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.",
]

model_name = "EmbeddedLLM/bge-base-en-v1.5-onnx-o4-o2-gpu"
device = "cuda"
provider = "CUDAExecutionProvider"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = ORTModelForFeatureExtraction.from_pretrained(
    model_name, use_io_binding=True, provider=provider, device_map=device
)
inputs = tokenizer(
    sentences,
    padding=True,
    truncation=True,
    return_tensors="pt",
    max_length=model.config.max_position_embeddings,
)
inputs = inputs.to(device)
embeddings = model(**inputs).last_hidden_state[:, 0]
embeddings = F.normalize(embeddings, p=2, dim=1)
print(embeddings.cpu().numpy().shape)