#!/usr/bin/env python """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 # Initializer sidecars remain external and must never be deleted here. 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()