Spaces:
Runtime error
Runtime error
| import os | |
| import torch | |
| from transformers import AutoFeatureExtractor, Wav2Vec2ForSequenceClassification | |
| from optimum.onnxruntime import ORTModelForAudioClassification | |
| from optimum.onnxruntime.configuration import AutoQuantizationConfig | |
| from optimum.onnxruntime import ORTQuantizer | |
| import shutil | |
| # --- CONFIG --- | |
| MODEL_PATH = "models/wav2vec2-finetuned" | |
| ONNX_PATH = "models/onnx" | |
| QUANTIZED_PATH = "models/onnx_quantized" | |
| def export_to_onnx(): | |
| print(f"Exporting PyTorch model to ONNX...") | |
| # Load PyTorch model | |
| feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_PATH) | |
| # Export to ONNX using Optimum (handles complex configs automatically) | |
| model = ORTModelForAudioClassification.from_pretrained(MODEL_PATH, export=True) | |
| # Save ONNX model | |
| if os.path.exists(ONNX_PATH): | |
| shutil.rmtree(ONNX_PATH) | |
| model.save_pretrained(ONNX_PATH) | |
| feature_extractor.save_pretrained(ONNX_PATH) | |
| print(f"ONNX model saved to: {ONNX_PATH}") | |
| def quantize_onnx(): | |
| print(f"Quantizing ONNX model to INT8...") | |
| # Load ONNX model for quantization | |
| quantizer = ORTQuantizer.from_pretrained(ONNX_PATH, file_name="model.onnx") | |
| # Define quantization config (INT8 dynamic quantization) | |
| # Exclude Conv layers to prevent 'initializer' errors in Wav2Vec2 | |
| qconfig = AutoQuantizationConfig.arm64( | |
| is_static=False, | |
| per_channel=False, | |
| operators_to_quantize=["MatMul", "Attention", "LSTM", "Gather", "Transpose", "EmbedLayerNormalization"] | |
| ) | |
| # Apply quantization | |
| if os.path.exists(QUANTIZED_PATH): | |
| shutil.rmtree(QUANTIZED_PATH) | |
| quantizer.quantize(save_dir=QUANTIZED_PATH, quantization_config=qconfig) | |
| # Copy feature extractor config to quantized folder so it's self-contained | |
| feature_extractor = AutoFeatureExtractor.from_pretrained(ONNX_PATH) | |
| feature_extractor.save_pretrained(QUANTIZED_PATH) | |
| print(f"INT8 Quantized model saved to: {QUANTIZED_PATH}") | |
| if __name__ == "__main__": | |
| export_to_onnx() | |
| quantize_onnx() | |