File size: 2,209 Bytes
29b9510 | 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 | """Minimal example for running Sakura Needle 3 ONNX inference."""
import json
import os
import sys
# Ensure parent directory is on sys.path when running directly
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
sys.path.insert(0, ROOT)
from runtime.needle_onnx import NeedleONNX
def main():
model_path = os.path.join(ROOT, "models", "model_int8_compact.onnx")
print(f"Loading Needle 3 ONNX engine from {model_path}...")
engine = NeedleONNX(model_path=model_path, num_threads=4)
print(f"Engine loaded in {engine.load_time_ms:.1f} ms.\n")
tools = [
{
"name": "control_lights",
"description": "Adjust lights in a specified room",
"parameters": {
"type": "object",
"properties": {
"room": {"type": "string", "description": "Room name"},
"action": {"type": "string", "enum": ["on", "off", "dim"]},
"brightness": {"type": "integer", "description": "Brightness percentage 0-100"},
},
"required": ["room", "action"],
},
},
{
"name": "set_timer",
"description": "Set a countdown timer",
"parameters": {
"type": "object",
"properties": {
"duration_minutes": {"type": "integer"},
"label": {"type": "string"},
},
"required": ["duration_minutes"],
},
},
]
query = "Turn off the kitchen lights and dim the living room lights to 35%."
print(f"Query: {query}\n")
system = "date: 2026-10-04 Sun 12:00; You control smart home devices."
print("Running completion (compatibility='cactus')...")
res = engine.complete(query, tools=tools, system=system, compatibility="cactus")
print(f"Prefill: {res['prefill_tps']:.1f} tok/s | Decode: {res['decode_tps']:.1f} tok/s | TTFT: {res['first_token_latency_ms']:.1f} ms")
print(f"Reasoning: {res['reasoning'].strip()}")
print("Function calls:")
print(json.dumps(res["function_calls"], indent=2))
if __name__ == "__main__":
main()
|