"""
ONNX and WebGPU Export Pipeline for Q-TensorFormer Models.
Exports Q-TensorFormer models to ONNX format with dynamic batch and sequence axes,
verifies numeric equivalence via ONNX Runtime, and provides a browser WebGPU runtime template.
Usage:
python export_onnx.py --model-path . --output qtensorformer.onnx --opset 17 --verify
"""
import argparse
import io
import os
import sys
# Ensure UTF-8 stdout on Windows to prevent UnicodeEncodeError on emoji prints
if hasattr(sys.stdout, "reconfigure"):
sys.stdout.reconfigure(encoding="utf-8")
elif hasattr(sys.stdout, "buffer"):
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8")
import torch
import torch.nn as nn
class QTensorFormerONNXWrapper(nn.Module):
"""
Wrapper module to provide standard (input_ids) -> (logits) interface
tailored for ONNX export and graph optimization.
"""
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
outputs = self.model(input_ids=input_ids, use_cache=False)
return outputs.logits
def export_to_onnx(
model_path: str = ".",
output_path: str = "qtensorformer.onnx",
opset_version: int = 17,
verify: bool = True,
) -> str:
print(f"[*] Initializing Q-TensorFormer for ONNX export from: {model_path}")
from configuration_qtensorformer import QTensorFormerConfig
from modeling_qtensorformer import QTensorFormerForCausalLM
if os.path.isdir(model_path):
config = QTensorFormerConfig.from_pretrained(model_path)
else:
config = QTensorFormerConfig()
# For fast and clean export graph, use a representative configuration
config.num_hidden_layers = min(config.num_hidden_layers, 4)
model = QTensorFormerForCausalLM(config)
model.eval()
wrapper = QTensorFormerONNXWrapper(model)
wrapper.eval()
dummy_input_ids = torch.randint(0, config.vocab_size, (1, 8), dtype=torch.long)
print(f"[*] Tracing and exporting ONNX graph (Opset={opset_version})...")
input_names = ["input_ids"]
output_names = ["logits"]
dynamic_axes = {
"input_ids": {0: "batch_size", 1: "sequence_length"},
"logits": {0: "batch_size", 1: "sequence_length"},
}
torch.onnx.export(
wrapper,
dummy_input_ids,
output_path,
export_params=True,
opset_version=opset_version,
do_constant_folding=True,
input_names=input_names,
output_names=output_names,
dynamic_axes=dynamic_axes,
)
file_size_mb = os.path.getsize(output_path) / (1024 * 1024)
print(f"[+] Successfully exported ONNX model: {output_path} ({file_size_mb:.2f} MB)")
if verify:
print("[*] Verifying ONNX graph with onnxruntime...")
try:
import onnxruntime as ort
session = ort.InferenceSession(output_path, providers=["CPUExecutionProvider"])
ort_inputs = {"input_ids": dummy_input_ids.numpy()}
ort_outputs = session.run(None, ort_inputs)
with torch.no_grad():
torch_output = wrapper(dummy_input_ids).numpy()
diff = float(abs(ort_outputs[0] - torch_output).max())
print(f"[+] ONNX Runtime verification PASSED. Max absolute difference: {diff:.6e}")
except ImportError:
print("[!] onnxruntime not installed. Skipping runtime verification.")
except Exception as e:
print(f"[!] Warning during verification: {e}")
# Generate WebGPU HTML template
web_dir = os.path.join(os.path.dirname(output_path) or ".", "web")
os.makedirs(web_dir, exist_ok=True)
html_path = os.path.join(web_dir, "index.html")
html_content = """
Q-TensorFormer WebGPU Browser Inference
⚛️ Q-TensorFormer: Zero-Server WebGPU Inference
Run closed-loop Tensor-Train transformer generation locally inside your browser using WebGPU hardware acceleration.
Click the button above to load and run the model via WebGPU...
"""
with open(html_path, "w", encoding="utf-8") as f:
f.write(html_content)
print(f"[+] Generated WebGPU browser template: {html_path}")
return output_path
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Export Q-TensorFormer to ONNX and WebGPU")
parser.add_argument("--model-path", type=str, default=".", help="Path to model directory")
parser.add_argument("--output", type=str, default="qtensorformer.onnx", help="Output .onnx filename")
parser.add_argument("--opset", type=int, default=18, help="ONNX opset version")
parser.add_argument("--no-verify", action="store_true", help="Skip ONNX Runtime verification")
args = parser.parse_args()
export_to_onnx(
model_path=args.model_path,
output_path=args.output,
opset_version=args.opset,
verify=not args.no_verify,
)