# %pip install setfit[onnx]from setfit import SetFitModel

model = SetFitModel.from_pretrained("lewtun/my-awesome-setfit-model")from setfit.exporters.onnx import export_onnx

# Export the sklearn based model
output_path = "sklearn_model.onnx"
export_onnx(model.model_body,
            model.model_head,
            opset=12,
            output_path=output_path)import os
assert output_path in os.listdir() # Run inference using the original model
input_text = ["i loved the spiderman movie!", "pineapple on pizza is the worst 🤮"]
pytorch_preds = model(input_text)
pytorch_preds import onnxruntime
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("lewtun/my-awesome-setfit-model")
inputs = tokenizer(
    input_text,
    padding=True,
    truncation=True,
    return_attention_mask=True,
    return_token_type_ids=True,
    return_tensors="np",
)

session = onnxruntime.InferenceSession(output_path)

onnx_preds = session.run(None, dict(inputs))[0]onnx_predsimport numpy as np
assert np.array_equal(onnx_preds, pytorch_preds)