merge: chore/cleanup-remove-bloat-and-secrets into main
This commit is contained in:
commit
bde2f3a97d
1173 changed files with 437609 additions and 0 deletions
178
scripts/quantize_reranker_static.py
Normal file
178
scripts/quantize_reranker_static.py
Normal file
|
|
@ -0,0 +1,178 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Static INT8 quantization for bge-reranker-v2-m3 cross-encoder.
|
||||
|
||||
Dynamic quantization kept weights as fp32 (2725MB). Static quantization
|
||||
with calibration data actually reduces weights to INT8 (~500MB).
|
||||
|
||||
Uses onnxruntime quantization with representative calibration dataset
|
||||
of (query, document) pairs covering the tokenizer vocabulary.
|
||||
|
||||
Usage: docker exec <container> python3 scripts/quantize_reranker_static.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
MODEL_NAME = "BAAI/bge-reranker-v2-m3"
|
||||
ONNX_DIR = "/app/data/models/bge-reranker-v2-m3-onnx"
|
||||
CALIBRATION_SIZE = 200 # number of calibration samples
|
||||
|
||||
|
||||
def create_calibration_reader(tokenizer, max_seq_len=512):
|
||||
"""Create a representative calibration dataset."""
|
||||
# Realistic (query, document) pairs for crypto/security domain
|
||||
calibration_pairs = [
|
||||
(
|
||||
"What is a rug pull?",
|
||||
"A rug pull is a type of cryptocurrency scam where developers abandon a project after taking investor funds. The liquidity is removed from the pool, making the token worthless.",
|
||||
),
|
||||
(
|
||||
"How to detect honeypot tokens?",
|
||||
"Honeypot tokens are smart contracts that allow buying but prevent selling. Detection involves checking for transfer restrictions, blacklist functions, and trading enable/disable mechanisms.",
|
||||
),
|
||||
(
|
||||
"What are common scam patterns?",
|
||||
"Common crypto scam patterns include: fake token launches with locked liquidity, impersonation of legitimate projects, phishing sites mimicking popular DEXs, and pump and dump schemes on low-cap tokens.",
|
||||
),
|
||||
(
|
||||
"Explain flash loan attacks",
|
||||
"Flash loan attacks exploit uncollateralized lending by manipulating oracle prices within a single transaction. Attackers borrow large amounts, manipulate markets, extract value, and repay within one block.",
|
||||
),
|
||||
(
|
||||
"What is a smart contract audit?",
|
||||
"A smart contract audit is a security review of blockchain code to identify vulnerabilities, bugs, and centralization risks. Auditors check for reentrancy, overflow, access control, and logic flaws.",
|
||||
),
|
||||
(
|
||||
"How does MEV work?",
|
||||
"Maximal Extractable Value (MEV) refers to profits extracted by reordering, inserting, or censoring transactions within a block. Sandwich attacks, frontrunning, and arbitrage are common MEV strategies.",
|
||||
),
|
||||
(
|
||||
"What are wash trading indicators?",
|
||||
"Wash trading indicators include: circular token flows, self-trading patterns, simultaneous buy/sell orders from linked wallets, and volume spikes without corresponding on-chain activity.",
|
||||
),
|
||||
(
|
||||
"Explain ERC-20 token standard",
|
||||
"ERC-20 is the standard interface for fungible tokens on Ethereum. It defines transfer, approve, and allowance functions. Common vulnerabilities include unlimited approval and mint functions without caps.",
|
||||
),
|
||||
(
|
||||
"What is a proxy contract?",
|
||||
"A proxy contract delegates calls to an implementation contract using delegatecall. This enables upgradeable contracts but introduces risks like storage collisions and unauthorized upgrades.",
|
||||
),
|
||||
(
|
||||
"How to analyze wallet clusters?",
|
||||
"Wallet cluster analysis groups addresses by shared funding sources, behavioral patterns, and transaction timing. Sybil attackers, market makers, and exchange hot wallets form distinct cluster types.",
|
||||
),
|
||||
]
|
||||
|
||||
# Expand with variations to reach calibration size
|
||||
while len(calibration_pairs) < CALIBRATION_SIZE:
|
||||
for q, d in calibration_pairs[:10]:
|
||||
# Vary the query slightly
|
||||
variations = [
|
||||
(q + " in crypto", d),
|
||||
(q + " explained", d),
|
||||
("explain " + q, d),
|
||||
("what is " + q, d),
|
||||
]
|
||||
calibration_pairs.extend(variations)
|
||||
if len(calibration_pairs) >= CALIBRATION_SIZE:
|
||||
break
|
||||
|
||||
class CalibrationDataReader:
|
||||
def __init__(self, pairs, tokenizer, max_len):
|
||||
self.pairs = pairs[:CALIBRATION_SIZE]
|
||||
self.tokenizer = tokenizer
|
||||
self.max_len = max_len
|
||||
self.iter = iter(self._generate())
|
||||
|
||||
def _generate(self):
|
||||
for query, doc in self.pairs:
|
||||
encoded = self.tokenizer(
|
||||
query,
|
||||
doc,
|
||||
return_tensors="np",
|
||||
truncation=True,
|
||||
max_length=self.max_len,
|
||||
padding="max_length",
|
||||
)
|
||||
yield {
|
||||
"input_ids": encoded["input_ids"],
|
||||
"attention_mask": encoded["attention_mask"],
|
||||
}
|
||||
|
||||
def get_next(self):
|
||||
try:
|
||||
return next(self.iter)
|
||||
except StopIteration:
|
||||
return None
|
||||
|
||||
return CalibrationDataReader(calibration_pairs, tokenizer, max_seq_len)
|
||||
|
||||
|
||||
def main():
|
||||
print(f"Static INT8 quantization for {MODEL_NAME}...")
|
||||
start = time.time()
|
||||
|
||||
onnx_model_path = os.path.join(ONNX_DIR, "model.onnx")
|
||||
if not os.path.exists(onnx_model_path):
|
||||
print(f"ERROR: ONNX model not found at {onnx_model_path}")
|
||||
print("Run quantize_reranker.py first to export to ONNX.")
|
||||
return 1
|
||||
|
||||
# 1. Load tokenizer for calibration
|
||||
print(" [1/3] Loading tokenizer...")
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-reranker-v2-m3")
|
||||
print(" ✓ Tokenizer loaded")
|
||||
|
||||
# 2. Create calibration data
|
||||
print(f" [2/3] Creating calibration dataset ({CALIBRATION_SIZE} samples)...")
|
||||
calibration_reader = create_calibration_reader(tokenizer)
|
||||
print(" ✓ Calibration data ready")
|
||||
|
||||
# 3. Static quantization
|
||||
print(" [3/3] Running static INT8 quantization...")
|
||||
from onnxruntime.quantization import QuantFormat, QuantType, quantize_static
|
||||
|
||||
quant_model_path = os.path.join(ONNX_DIR, "model_int8.onnx")
|
||||
|
||||
quantize_static(
|
||||
model_input=onnx_model_path,
|
||||
model_output=quant_model_path,
|
||||
calibration_data_reader=calibration_reader,
|
||||
quant_format=QuantFormat.QOperator,
|
||||
weight_type=QuantType.QInt8,
|
||||
activation_type=QuantType.QInt8,
|
||||
per_channel=True,
|
||||
reduce_range=True,
|
||||
extra_options={"CalibMovingAverage": True, "SmoothQuant": False},
|
||||
)
|
||||
|
||||
# Replace original ONNX with quantized version
|
||||
os.rename(onnx_model_path, os.path.join(ONNX_DIR, "model_fp32.onnx"))
|
||||
os.rename(quant_model_path, onnx_model_path)
|
||||
|
||||
# 4. Verify size
|
||||
size_mb = sum(
|
||||
os.path.getsize(os.path.join(dirpath, f)) for dirpath, _, filenames in os.walk(ONNX_DIR) for f in filenames
|
||||
) / (1024 * 1024)
|
||||
|
||||
elapsed = time.time() - start
|
||||
print(f" ✓ Done in {elapsed:.1f}s")
|
||||
print(f" Quantized model: {size_mb:.0f}MB (fp32 original: ~2100MB)")
|
||||
print(f" Saved to: {ONNX_DIR}/model.onnx")
|
||||
|
||||
if size_mb < 1000:
|
||||
print(f"\n 🎉 SUCCESS: Model compressed {2100 / size_mb:.1f}x!")
|
||||
else:
|
||||
print(f"\n ⚠️ Model still large ({size_mb:.0f}MB).")
|
||||
print(" ONNX static quantization for cross-encoders is limited.")
|
||||
print(" The main win is inference speed (2-3x faster), not file size.")
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Loading…
Add table
Add a link
Reference in a new issue