File size: 2,974 Bytes
bde2f3a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | #!/usr/bin/env python3
"""Quantize bge-reranker-v2-m3 to ONNX INT8.
Reduces model from 2.1GB fp32 → ~500MB INT8, 2-3x faster inference.
Run once to generate the quantized model. The reranker auto-detects and prefers it.
Usage:
docker exec rmi-backend python3 scripts/quantize_reranker.py
"""
import os
import shutil
import sys
import time
MODEL_NAME = "BAAI/bge-reranker-v2-m3"
OUTPUT_DIR = "/app/data/models/bge-reranker-v2-m3-onnx"
def main():
print(f"Quantizing {MODEL_NAME} → ONNX INT8...")
start = time.time()
# Ensure output directory
os.makedirs(OUTPUT_DIR, exist_ok=True)
# 1. Export to ONNX using optimum
try:
from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer
print(" [1/3] Exporting to ONNX...")
model = ORTModelForSequenceClassification.from_pretrained(
MODEL_NAME, export=True, provider="CPUExecutionProvider"
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
# Save ONNX model
model.save_pretrained(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)
print(f" ✓ ONNX model saved to {OUTPUT_DIR}")
# 2. Quantize
print(" [2/3] Quantizing to INT8...")
from onnxruntime.quantization import QuantType, quantize_dynamic
onnx_path = os.path.join(OUTPUT_DIR, "model.onnx")
quant_path = os.path.join(OUTPUT_DIR, "model_quantized.onnx")
quantize_dynamic(
model_input=onnx_path,
model_output=quant_path,
weight_type=QuantType.QInt8,
)
# Replace original with quantized
shutil.move(quant_path, onnx_path)
print(" ✓ Quantized to INT8")
except ImportError:
print(" optimum/onnxruntime not available, trying direct sentence-transformers export...")
try:
from sentence_transformers import CrossEncoder
model = CrossEncoder(MODEL_NAME)
# sentence-transformers >= 3.3 supports ONNX export
if hasattr(model, "save"):
model.save(OUTPUT_DIR, safe_serialization=False)
print(f" ✓ Exported to {OUTPUT_DIR}")
else:
print(" ✗ sentence-transformers too old for ONNX export")
return 1
except Exception as e:
print(f" ✗ Export failed: {e}")
return 1
# 3. Verify
print(" [3/3] Verifying...")
size_mb = sum(
os.path.getsize(os.path.join(dirpath, f)) for dirpath, _, filenames in os.walk(OUTPUT_DIR) for f in filenames
) / (1024 * 1024)
elapsed = time.time() - start
print(f" ✓ Done in {elapsed:.1f}s")
print(f" Model size: {size_mb:.0f}MB (was ~2100MB fp32)")
print(f" Saved to: {OUTPUT_DIR}")
print("\n The reranker will auto-detect and use this optimized model.")
return 0
if __name__ == "__main__":
sys.exit(main())
|