File size: 2,033 Bytes
8d737ff | 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 63 64 65 66 67 68 | """
Mara AFM Quickstart: Minimal On-Device Tool Calling Example.
"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
from mara.afm import tool, default_registry
from mara.model import Mara, MaraConfig
from mara.tokenizer import load_tokenizer
from mara.data import TOKENIZER_PATH
# 1. Define tools using standard Python typing and docstrings
@tool
def set_lights(room: str, brightness: int = 100) -> str:
"""Adjust brightness for a specific room light."""
return f"Set {room} brightness to {brightness}%"
@tool
def lock_door(door: str, locked: bool = True) -> str:
"""Lock or unlock an exterior security door."""
status = "locked" if locked else "unlocked"
return f"Door {door} is now {status}"
def main():
device = torch.device("cpu")
# 2. Load model & tokenizer
tok = load_tokenizer(TOKENIZER_PATH)
ckpt = torch.load("checkpoints/mara_afm.pt", map_location=device, weights_only=False)
cfg = MaraConfig(**ckpt["config"])
model = Mara(cfg).to(device)
model.load_state_dict(ckpt["model"])
model.eval()
print(f"Mara AFM loaded successfully ({model.num_params():,} parameters).")
# 3. User request
query = "dim the living room light to 30 percent"
print(f"\nUser Query: '{query}'")
# 4. Compile fast-path decision record and infer routing probabilities
record = default_registry.compile_fast_path_record(query)
with torch.no_grad():
probs = model.probs(tok, record, device=device)
tool_names = list(default_registry.tools.keys()) + ["none"]
pred_idx = probs[0].argmax().item()
confidence = probs[0][pred_idx].item()
selected_tool = tool_names[pred_idx]
print(f"Selected Tool: {selected_tool} (Confidence: {confidence * 100:.1f}%)")
# 5. Dispatch execution
if selected_tool == "set_lights":
result = set_lights(room="living room", brightness=30)
print(f"Execution Result: {result}")
if __name__ == "__main__":
main()
|