| """Minimal example for running Sakura Needle 3 ONNX inference.""" |
|
|
| import json |
| import os |
| import sys |
|
|
| |
| 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() |
|
|
