Q-TensorFormer / scripts /run_onnx_edge.py
Premchandyadav369
feat: real-world production pillars, OpenAI streaming API, and reasoning benchmarks
000c57b
Raw History Blame Contribute Delete
3.03 kB
#!/usr/bin/env python3
"""
scripts/run_onnx_edge.py
Standalone Edge Embedded Deployment Runner for Q-TensorFormer.
Executes token generation via ONNX Runtime with zero PyTorch runtime dependency:
- Runs dynamic batch and sequence axes on edge CPUs/SoCs (ARM Cortex, Jetson, Raspberry Pi).
- Profiles Inter-Token Latency (ITL) and memory footprints.
"""
import sys
import time
from pathlib import Path
import numpy as np
ROOT_DIR = Path(__file__).resolve().parent.parent
if str(ROOT_DIR) not in sys.path:
sys.path.insert(0, str(ROOT_DIR))
try:
import onnxruntime as ort
HAS_ORT = True
except ImportError:
HAS_ORT = False
from src.tokenizer import QTensorFormerTokenizer
def run_onnx_inference(onnx_model_path: Path, max_new_tokens: int = 16):
print("=" * 70)
print("Q-TensorFormer: Standalone Edge ONNX Runtime Deployment")
print(f"Target Model: {onnx_model_path}")
print("=" * 70)
if not HAS_ORT:
print("[Error] onnxruntime is not installed. Run: pip install onnxruntime")
return
tokenizer = QTensorFormerTokenizer(vocab_size=256)
prompt = "Edge computing requires adaptive computation"
token_ids = tokenizer.encode(prompt, add_bos=True)
print(f"Input Prompt: '{prompt}' ({len(token_ids)} tokens)")
# Initialize ONNX Runtime Session with optimization
opts = ort.SessionOptions()
opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
opts.intra_op_num_threads = 4
session = ort.InferenceSession(str(onnx_model_path), sess_options=opts, providers=["CPUExecutionProvider"])
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
print("\n[+] Starting Autoregressive Generation on Edge ONNX Engine...")
curr_ids = list(token_ids)
latencies = []
for step in range(max_new_tokens):
t0 = time.perf_counter()
inp = np.array([curr_ids], dtype=np.int64)
outputs = session.run([output_name], {input_name: inp})
logits = outputs[0] # [1, seq_len, vocab_size]
next_token = int(np.argmax(logits[0, -1, :]))
curr_ids.append(next_token)
lat = (time.perf_counter() - t0) * 1000.0
latencies.append(lat)
gen_ids = curr_ids[len(token_ids):]
generated_text = tokenizer.decode(gen_ids)
print(f"\n[Generated Output]: {generated_text}")
print(f"Mean Per-Token Latency: {np.mean(latencies):.2f} ms")
print(f"P99 Per-Token Latency: {np.percentile(latencies, 99):.2f} ms")
print(f"Edge Token Throughput: {1000.0 / np.mean(latencies):.2f} tokens/second")
print("=" * 70)
def main():
onnx_file = ROOT_DIR / "qtensorformer.onnx"
if not onnx_file.exists():
print(f"[*] ONNX model not found at {onnx_file}. Triggering automated export...")
from export_onnx import export_to_onnx
export_to_onnx(model_path=str(ROOT_DIR), output_path=str(onnx_file), verify=False)
run_onnx_inference(onnx_file, max_new_tokens=12)
if __name__ == "__main__":
main()