optimum_onnx_runtime_bert_text_classification_inference.py

1from optimum.onnxruntime import ORTModelForSequenceClassification
2from transformers import AutoTokenizer, pipeline
3
4model_id = "distilbert-base-uncased-finetuned-sst-2-english"
5save_directory = "onnx_output"
6
7# Load a model from transformers and export it to ONNX
8model = ORTModelForSequenceClassification.from_pretrained(model_id, export=True)
9tokenizer = AutoTokenizer.from_pretrained(model_id)
10
11# Save the converted model
12model.save_pretrained(save_directory)
13tokenizer.save_pretrained(save_directory)
14
15# Run inference using the ONNX Runtime
16classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)
17results = classifier("Optimum is a great tool for model optimization.")
18
19print(results)