92 lines
2.9 KiB
Python
92 lines
2.9 KiB
Python
#!/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())
|