| """ |
| 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 |
|
|
|
|
| |
| @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") |
|
|
| |
| 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).") |
|
|
| |
| query = "dim the living room light to 30 percent" |
| print(f"\nUser Query: '{query}'") |
|
|
| |
| 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}%)") |
|
|
| |
| if selected_tool == "set_lights": |
| result = set_lights(room="living room", brightness=30) |
| print(f"Execution Result: {result}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
