Skip to content
💻 🧠 Code 1001 > LLMs: Articles, Practical Materials, and Resources > ⚡ Optimized Version with optimum and onnxruntime

⚡ Optimized Version with optimum and onnxruntime

from transformers import AutoTokenizer, pipeline
from optimum.onnxruntime import ORTModelForSequenceClassification
from optimum.onnxruntime.configuration import AutoQuantizationConfig
from optimum.onnxruntime import ORTQuantizer
from pathlib import Path

# === CONFIGURATION ===
model_id: str = "hypo69/my_model_from_existing_datasets"
onnx_dir: Path = Path("./onnx-model")
onnx_dir.mkdir(exist_ok=True)

# === STEP 1: LOAD ORIGINAL MODEL AND TOKENIZER ===
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = ORTModelForSequenceClassification.from_pretrained(model_id, export=True)  # auto-export to ONNX

# === STEP 2: SAVE THE EXPORTED MODEL ===
model.save_pretrained(onnx_dir)
tokenizer.save_pretrained(onnx_dir)

# === STEP 3: QUANTIZATION (REDUCE SIZE, INCREASE SPEED) ===
quantized_dir = onnx_dir / "quantized"
quantized_dir.mkdir(exist_ok=True)

quantizer = ORTQuantizer.from_pretrained(model)
qconfig = AutoQuantizationConfig.avx512_vnni(is_static=False)  # choose dynamically if needed
quantizer.quantize(save_dir=quantized_dir, quantization_config=qconfig)

# === STEP 4: LOAD OPTIMIZED MODEL ===
optimized_model = ORTModelForSequenceClassification.from_pretrained(quantized_dir)
optimized_tokenizer = AutoTokenizer.from_pretrained(quantized_dir)

# === STEP 5: CREATE PIPELINE ===
classifier = pipeline("text-classification", model=optimized_model, tokenizer=optimized_tokenizer)

# === STEP 6: RUN INFERENCE ===
text = "ONNX Runtime with quantization makes inference super fast!"
result = classifier(text)
print(result)

📊 What This Code Does

StepDescription
1. LoadDownloads the model from Hugging Face Hub and automatically exports it to ONNX format.
2. SaveSaves both model and tokenizer to a local directory.
3. QuantizationApplies AVX512_VNNI optimization — reduces model size by 3–4× and boosts inference speed.
4. Load Optimized ModelLoads the optimized (quantized) ONNX model for inference.
5. PipelineBuilds a transformers.pipeline that uses the ONNX backend.
6. InferenceRuns text classification with significantly faster inference (2–5× faster than PyTorch/TensorFlow).

⚙️ Output Example

[{'label': 'POSITIVE', 'score': 0.99976}]

💡 Additional Tips

  • Replace AutoQuantizationConfig.avx512_vnni with AutoQuantizationConfig.arm64() for Apple Silicon (M1/M2) or other ARM systems.
  • Add "provider": "CUDAExecutionProvider" to from_pretrained() for GPU acceleration.
  • Set "use_io_binding=True" for large batch inference optimization.

Would you like me to extend this version with automatic GPU (CUDA) support and CPU fallback?

Leave a Reply

Your email address will not be published. Required fields are marked *