| |
| """Export experimental EmbeddingGemma 2 embedding ONNX variants only.""" |
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
| from export_common import ( |
| MODEL_ID, REVISION, ROOT, OPSET, EMBEDDING_DIM, EXTERNAL_DATA, |
| export_patches, load_reference, quantize_int8, quantize_int4, |
| print_artifact, sha256_file, |
| ) |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from transformers import AutoModel |
|
|
| EMBEDDING_OUT = ROOT / "audio_embedding.onnx" |
| EMBEDDING_TEMP = ROOT / "audio_embedding.tmp.onnx" |
| EMBEDDING_INT8_OUT = ROOT / "audio_embedding_int8.onnx" |
| EMBEDDING_INT8_TEMP = ROOT / "audio_embedding_int8.tmp.onnx" |
| EMBEDDING_INT4_OUT = ROOT / "audio_embedding_int4.onnx" |
| EMBEDDING_INT4_TEMP = ROOT / "audio_embedding_int4.tmp.onnx" |
|
|
|
|
| class AudioEmbedding(torch.nn.Module): |
| def __init__(self, model: torch.nn.Module): |
| super().__init__() |
| self.model = model |
|
|
| def forward(self, input_ids, attention_mask, input_features, input_features_mask): |
| result = self.model( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| input_features=input_features, |
| input_features_mask=input_features_mask, |
| return_dict=True, |
| ) |
| states = result.last_hidden_state.to(torch.float32) |
| mask = attention_mask.unsqueeze(-1).to(torch.float32) |
| pooled = (states * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0) |
| return F.normalize(pooled[:, :EMBEDDING_DIM], p=2.0, dim=-1, eps=1e-12) |
|
|
|
|
| def _numpy_inputs(inputs: dict[str, torch.Tensor]) -> dict[str, np.ndarray]: |
| dtypes = { |
| "input_ids": np.int64, |
| "attention_mask": np.int64, |
| "input_features": np.float32, |
| "input_features_mask": np.bool_, |
| } |
| return {name: inputs[name].detach().cpu().numpy().astype(dtype) for name, dtype in dtypes.items()} |
|
|
|
|
| def check_embedding( |
| model_path: Path, |
| inputs: dict[str, torch.Tensor], |
| reference: np.ndarray, |
| *, |
| min_cosine: float, |
| max_abs_diff: float | None, |
| ) -> dict[str, float]: |
| import onnx |
| import onnxruntime as ort |
|
|
| onnx.checker.check_model(str(model_path)) |
| options = ort.SessionOptions() |
| options.log_severity_level = 3 |
| options.intra_op_num_threads = 1 |
| options.inter_op_num_threads = 1 |
| session = ort.InferenceSession(str(model_path), options, providers=["CPUExecutionProvider"]) |
| actual = session.run(["embedding"], _numpy_inputs(inputs))[0].astype(np.float32) |
| if actual.shape != (1, EMBEDDING_DIM): |
| raise RuntimeError(f"{model_path.name} output shape {actual.shape}, expected (1, {EMBEDDING_DIM})") |
| if not np.isfinite(actual).all() or not np.isfinite(reference).all(): |
| raise RuntimeError(f"PyTorch or {model_path.name} produced non-finite embeddings") |
| cosine = float(np.sum(reference * actual)) |
| max_diff = float(np.max(np.abs(reference - actual))) |
| if cosine < min_cosine or (max_abs_diff is not None and max_diff > max_abs_diff): |
| raise RuntimeError(f"{model_path.name} parity failed: cosine={cosine}, max_abs_diff={max_diff}") |
| return {"cosine": cosine, "max_abs_diff": max_diff} |
|
|
|
|
| def inline_external_constants(model_path: Path) -> None: |
| """Inline small exported Constant tensors before removing their sidecar files.""" |
| import onnx |
| from onnx.external_data_helper import load_external_data_for_tensor |
|
|
| def attribute_tensors(graph): |
| for node in graph.node: |
| for attribute in node.attribute: |
| if attribute.type == onnx.AttributeProto.TENSOR: |
| yield attribute.t |
| elif attribute.type == onnx.AttributeProto.TENSORS: |
| yield from attribute.tensors |
| elif attribute.type == onnx.AttributeProto.GRAPH: |
| yield from attribute_tensors(attribute.g) |
| elif attribute.type == onnx.AttributeProto.GRAPHS: |
| for nested in attribute.graphs: |
| yield from attribute_tensors(nested) |
|
|
| model = onnx.load(str(model_path), load_external_data=False) |
| source_paths = set() |
| for tensor in attribute_tensors(model.graph): |
| if tensor.data_location != onnx.TensorProto.EXTERNAL: |
| continue |
| metadata = {entry.key: entry.value for entry in tensor.external_data} |
| location = metadata.get("location", "") |
| source = model_path.parent / location |
| if (not location or source.is_symlink() |
| or source.resolve().parent != model_path.parent.resolve() |
| or not source.is_file()): |
| raise RuntimeError(f"Unsafe Constant sidecar: {location!r}") |
| length = int(metadata.get("length", str(source.stat().st_size))) |
| if length > 1024 * 1024: |
| raise RuntimeError(f"Constant sidecar exceeds 1 MiB: {location!r}") |
| load_external_data_for_tensor(tensor, str(model_path.parent)) |
| tensor.data_location = onnx.TensorProto.DEFAULT |
| del tensor.external_data[:] |
| source_paths.add(source) |
| if not source_paths: |
| return |
|
|
| |
| retained = { |
| (model_path.parent / entry.value).resolve() |
| for tensor in model.graph.initializer |
| for entry in tensor.external_data if entry.key == "location" |
| } |
| temporary = model_path.with_name(model_path.stem + ".constants.tmp.onnx") |
| try: |
| onnx.save_model(model, str(temporary), save_as_external_data=False) |
| onnx.checker.check_model(str(temporary)) |
| temporary.replace(model_path) |
| for source in source_paths: |
| if source.resolve() not in retained: |
| source.unlink() |
| finally: |
| temporary.unlink(missing_ok=True) |
|
|
|
|
| def consolidate_external_data(model_path: Path, data_filename: str) -> None: |
| """Pack the fp32 embedding graph's per-tensor files into one aligned sidecar.""" |
| import onnx |
|
|
| model = onnx.load(str(model_path), load_external_data=False) |
| tensors = [t for t in model.graph.initializer if t.data_location == onnx.TensorProto.EXTERNAL] |
| if not tensors: |
| return |
|
|
| data_path = model_path.parent / data_filename |
| temp_data_path = data_path.with_name(data_path.name + ".tmp") |
| temp_model_path = model_path.with_name(model_path.stem + ".consolidated.tmp.onnx") |
| alignment = 64 * 1024 |
| offset = 0 |
| source_locations: set[str] = set() |
|
|
| try: |
| with temp_data_path.open("wb") as destination: |
| for tensor in tensors: |
| metadata = {entry.key: entry.value for entry in tensor.external_data} |
| location = metadata.get("location") |
| if not location or location in source_locations: |
| raise RuntimeError(f"Unexpected ONNX external-data location: {location!r}") |
| source_locations.add(location) |
| source_path = model_path.parent / location |
| if not source_path.is_file(): |
| raise RuntimeError(f"Missing ONNX external-data file: {source_path}") |
|
|
| source_offset = int(metadata.get("offset", "0")) |
| source_length = int(metadata.get("length", str(source_path.stat().st_size - source_offset))) |
| aligned_offset = (offset + alignment - 1) // alignment * alignment |
| destination.write(b"\0" * (aligned_offset - offset)) |
| with source_path.open("rb") as source: |
| source.seek(source_offset) |
| remaining = source_length |
| while remaining: |
| chunk = source.read(min(8 * 1024 * 1024, remaining)) |
| if not chunk: |
| raise RuntimeError(f"Short read from ONNX external-data file: {source_path}") |
| destination.write(chunk) |
| remaining -= len(chunk) |
|
|
| del tensor.external_data[:] |
| for key, value in ( |
| ("location", data_filename), |
| ("offset", str(aligned_offset)), |
| ("length", str(source_length)), |
| ): |
| entry = tensor.external_data.add() |
| entry.key = key |
| entry.value = value |
| tensor.data_location = onnx.TensorProto.EXTERNAL |
| offset = aligned_offset + source_length |
|
|
| onnx.save_model(model, str(temp_model_path), save_as_external_data=False) |
| temp_data_path.replace(data_path) |
| reloaded = onnx.load(str(temp_model_path), load_external_data=False) |
| references = [t for t in reloaded.graph.initializer if t.data_location == onnx.TensorProto.EXTERNAL] |
| if len(references) != len(tensors): |
| raise RuntimeError("External-data consolidation changed the initializer count") |
| data_size = data_path.stat().st_size |
| for tensor in references: |
| metadata = {entry.key: entry.value for entry in tensor.external_data} |
| start, length = int(metadata.get("offset", "-1")), int(metadata.get("length", "-1")) |
| if metadata.get("location") != data_filename or start < 0 or length < 0 or start + length > data_size: |
| raise RuntimeError(f"Invalid consolidated external-data reference: {tensor.name}") |
| temp_model_path.replace(model_path) |
| for location in source_locations: |
| (model_path.parent / location).unlink(missing_ok=True) |
| finally: |
| temp_data_path.unlink(missing_ok=True) |
| temp_model_path.unlink(missing_ok=True) |
|
|
|
|
| def export_embeddings(config, processor) -> None: |
| config.vision_config = None |
| model = AutoModel.from_pretrained( |
| MODEL_ID, |
| revision=REVISION, |
| config=config, |
| dtype=torch.float32, |
| ).eval() |
| wrapper = AudioEmbedding(model).eval() |
| sample_rate = int(processor.feature_extractor.sampling_rate) |
| audio3 = (0.4 * np.sin(2 * np.pi * 220 * np.arange(3 * sample_rate) / sample_rate)).astype(np.float32) |
| audio1 = (0.4 * np.sin(2 * np.pi * 220 * np.arange(sample_rate) / sample_rate)).astype(np.float32) |
| inputs3 = processor(audio=[audio3], sampling_rate=sample_rate, return_tensors="pt") |
| inputs1 = processor(audio=[audio1], sampling_rate=sample_rate, return_tensors="pt") |
| required = {"input_ids", "attention_mask", "input_features", "input_features_mask"} |
| for label, inputs in (("3 s", inputs3), ("1 s", inputs1)): |
| missing = required - set(inputs) |
| if missing: |
| raise RuntimeError(f"Processor omitted {label} graph inputs: {sorted(missing)}") |
| for inputs in (inputs3, inputs1): |
| inputs["input_features"] = inputs["input_features"].to(torch.float32) |
| inputs["input_features_mask"] = inputs["input_features_mask"].to(torch.bool) |
| with torch.inference_mode(): |
| ref3 = wrapper( |
| inputs3["input_ids"], inputs3["attention_mask"], |
| inputs3["input_features"], inputs3["input_features_mask"], |
| ).cpu().numpy().astype(np.float32) |
| ref1 = wrapper( |
| inputs1["input_ids"], inputs1["attention_mask"], |
| inputs1["input_features"], inputs1["input_features_mask"], |
| ).cpu().numpy().astype(np.float32) |
|
|
| dynamic_axes = { |
| "input_ids": {1: "text_tokens"}, |
| "attention_mask": {1: "text_tokens"}, |
| "input_features": {1: "feature_frames"}, |
| "input_features_mask": {1: "feature_frames"}, |
| } |
| with export_patches(include_text=True): |
| torch.onnx.export( |
| wrapper, |
| ( |
| inputs3["input_ids"], |
| inputs3["attention_mask"], |
| inputs3["input_features"], |
| inputs3["input_features_mask"], |
| ), |
| str(EMBEDDING_TEMP), |
| input_names=["input_ids", "attention_mask", "input_features", "input_features_mask"], |
| output_names=["embedding"], |
| opset_version=OPSET, |
| dynamo=False, |
| do_constant_folding=True, |
| dynamic_axes=dynamic_axes, |
| ) |
| raw3 = check_embedding( |
| EMBEDDING_TEMP, inputs3, ref3, min_cosine=0.99999, max_abs_diff=5e-4 |
| ) |
| EMBEDDING_TEMP.replace(EMBEDDING_OUT) |
|
|
| quantize_int8(EMBEDDING_OUT, EMBEDDING_INT8_TEMP) |
| int8_3 = check_embedding( |
| EMBEDDING_INT8_TEMP, inputs3, ref3, min_cosine=0.97, max_abs_diff=None |
| ) |
| EMBEDDING_INT8_TEMP.replace(EMBEDDING_INT8_OUT) |
| quantize_int4(EMBEDDING_OUT, EMBEDDING_INT4_TEMP) |
| int4_3 = check_embedding( |
| EMBEDDING_INT4_TEMP, inputs3, ref3, min_cosine=0.80, max_abs_diff=None |
| ) |
| EMBEDDING_INT4_TEMP.replace(EMBEDDING_INT4_OUT) |
|
|
| short = { |
| path: check_embedding(path, inputs1, ref1, min_cosine=threshold, max_abs_diff=None) |
| for path, threshold in ( |
| (EMBEDDING_OUT, 0.99999), |
| (EMBEDDING_INT8_OUT, 0.97), |
| (EMBEDDING_INT4_OUT, 0.80), |
| ) |
| } |
| inline_external_constants(EMBEDDING_OUT) |
| consolidate_external_data(EMBEDDING_OUT, EXTERNAL_DATA) |
| print_artifact(EMBEDDING_OUT, raw3, "Embedding fp32; 3 s") |
| print_artifact(EMBEDDING_INT8_OUT, int8_3, "Embedding int8; 3 s") |
| print_artifact(EMBEDDING_INT4_OUT, int4_3, "Embedding int4; 3 s") |
| data_path = ROOT / EXTERNAL_DATA |
| print(f"External weights: {data_path.name} | {data_path.stat().st_size:,} bytes | sha256={sha256_file(data_path)}") |
| print(f"Embedding short-audio cosine: fp32={short[EMBEDDING_OUT]['cosine']:.7f}, " |
| f"int8={short[EMBEDDING_INT8_OUT]['cosine']:.7f}, int4={short[EMBEDDING_INT4_OUT]['cosine']:.7f}") |
|
|
|
|
| def main() -> None: |
| from validation import parser, run_validation |
|
|
| args = parser(__doc__).parse_args() |
| print("EXPERIMENTAL: all artifacts from this exporter require downstream qualification.", flush=True) |
| if not args.validate_only: |
| config, processor = load_reference() |
| export_embeddings(config, processor) |
| if not args.skip_validation: |
| raise SystemExit(0 if run_validation("embedding", args) else 1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
