File size: 6,434 Bytes
61af2eb | 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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | #!/usr/bin/env python3
"""Export a q8 ONNX model and validate the browser-demo embeddings."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import numpy as np
import onnx
from onnx import numpy_helper
from optimum.onnxruntime.configuration import AutoQuantizationConfig
from sentence_transformers import SentenceTransformer
from sentence_transformers.backend import export_dynamic_quantized_onnx_model
from goktugtr.text import harrier_query
EXAMPLE_QUERIES = [
"Bir modeli küçük GPU'da nasıl verimli eğitebilirim?",
"Arama sistemlerinde kalite nasıl ölçülür?",
"LiDAR verisinden üç boyutlu harita üretmek",
"Yapay zekâ sonuçları neden yeniden üretilebilir olmalı?",
]
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=Path, default=Path("outputs/goktugtr-270m/final"))
parser.add_argument("--space-dir", type=Path, default=Path("space"))
parser.add_argument(
"--report", type=Path, default=Path("artifacts/onnx-browser-validation.json")
)
return parser.parse_args()
def encode(model: SentenceTransformer, texts: list[str]) -> np.ndarray:
return model.encode(
texts,
batch_size=8,
normalize_embeddings=True,
convert_to_numpy=True,
show_progress_bar=True,
).astype("float32")
def safe_quantization_config() -> object:
"""Quantize the large embedding table while preserving transformer fidelity."""
return AutoQuantizationConfig.avx2(
is_static=False,
operators_to_quantize=["Gather"],
)
def requantize_embedding_per_dimension(
onnx_path: Path, pytorch_model: SentenceTransformer
) -> None:
"""Replace the lossy single-scale embedding quantizer with 640 scales."""
embedding = (
pytorch_model[0]
.auto_model.get_input_embeddings()
.weight.detach()
.float()
.cpu()
.numpy()
)
scales = np.max(np.abs(embedding), axis=0).astype("float32") / 127.0
scales = np.maximum(scales, np.finfo("float32").tiny)
quantized = np.clip(np.rint(embedding / scales[None, :]), -127, 127).astype("int8")
zero_points = np.zeros(scales.shape, dtype="int8")
model = onnx.load(onnx_path)
replacements = {
"embed_tokens.weight_quantized": numpy_helper.from_array(
quantized, name="embed_tokens.weight_quantized"
),
"embed_tokens.weight_scale": numpy_helper.from_array(
scales, name="embed_tokens.weight_scale"
),
"embed_tokens.weight_zero_point": numpy_helper.from_array(
zero_points, name="embed_tokens.weight_zero_point"
),
}
replaced: set[str] = set()
for index, initializer in enumerate(model.graph.initializer):
if initializer.name in replacements:
model.graph.initializer[index].CopyFrom(replacements[initializer.name])
replaced.add(initializer.name)
if replaced != set(replacements):
raise RuntimeError(f"Embedding quantizer tensors not found: {set(replacements) - replaced}")
dequantizers = [
node
for node in model.graph.node
if node.op_type == "DequantizeLinear"
and node.name.endswith("Gather_output_0_DequantizeLinear")
]
if len(dequantizers) != 1:
raise RuntimeError(f"Expected one embedding dequantizer, found {len(dequantizers)}")
del dequantizers[0].attribute[:]
dequantizers[0].attribute.append(onnx.helper.make_attribute("axis", 2))
onnx.checker.check_model(model)
temporary = onnx_path.with_suffix(".tmp.onnx")
onnx.save(model, temporary)
temporary.replace(onnx_path)
def main() -> None:
args = parse_args()
corpus = json.loads((args.space_dir / "corpus.json").read_text(encoding="utf-8"))
documents = [item["text"] for item in corpus]
queries = [harrier_query(query) for query in EXAMPLE_QUERIES]
pytorch_model = SentenceTransformer(str(args.model), device="cpu")
pytorch_model.max_seq_length = 256
reference = encode(pytorch_model, documents + queries)
export_model = SentenceTransformer(str(args.model), backend="onnx", device="cpu")
export_dynamic_quantized_onnx_model(
export_model,
quantization_config=safe_quantization_config(),
model_name_or_path=str(args.model),
file_suffix="quantized",
)
onnx_path = args.model / "onnx/model_quantized.onnx"
requantize_embedding_per_dimension(onnx_path, pytorch_model)
q8_model = SentenceTransformer(
str(args.model),
backend="onnx",
device="cpu",
model_kwargs={"file_name": "onnx/model_quantized.onnx"},
)
q8_model.max_seq_length = 256
quantized = encode(q8_model, documents + queries)
document_count = len(documents)
reference_documents = reference[:document_count]
reference_queries = reference[document_count:]
quantized_documents = quantized[:document_count]
quantized_queries = quantized[document_count:]
embedding_cosines = np.sum(reference * quantized, axis=1)
reference_top = np.argmax(reference_queries @ reference_documents.T, axis=1)
quantized_top = np.argmax(quantized_queries @ quantized_documents.T, axis=1)
report = {
"format": "ONNX mixed int8 embedding / float transformer blocks",
"quantization": "per-dimension int8 embedding; float transformer blocks",
"documents": document_count,
"queries": len(queries),
"embedding_dimension": int(quantized.shape[1]),
"mean_pytorch_to_q8_cosine": float(np.mean(embedding_cosines)),
"minimum_pytorch_to_q8_cosine": float(np.min(embedding_cosines)),
"maximum_absolute_component_difference": float(np.max(np.abs(reference - quantized))),
"example_query_top1_agreement": float(np.mean(reference_top == quantized_top)),
"reference_top1_indices": reference_top.tolist(),
"quantized_top1_indices": quantized_top.tolist(),
}
args.report.parent.mkdir(parents=True, exist_ok=True)
args.report.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
(args.space_dir / "embeddings.json").write_text(
json.dumps(np.round(quantized_documents, 7).tolist(), separators=(",", ":")) + "\n",
encoding="utf-8",
)
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|