| """Public API: load a released ActionMan and turn (query, tools) into tool calls. |
| |
| from actionman.api import ActionMan |
| am = ActionMan.from_pretrained("shing-dev/ActionMan-0.1") # or a local export directory |
| am.predict("Weather in Paris in celsius?", tools) # tools: OpenAI-style function schemas |
| # -> [{"name": "get_weather", "arguments": {"city": "Paris", "unit": "celsius"}, "confidence": 0.97}] |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import os |
|
|
| import torch |
|
|
| from . import cached as C |
| from .features import load_tokenizer |
| from .infer import predict as _predict |
| from .model import ActionMan as _Model |
| from .schema import Example, from_openai_tool |
|
|
|
|
| class ActionMan: |
| def __init__(self, model, tok, flags: dict, device: str): |
| self.model, self.tok, self.flags, self.device = model, tok, flags, device |
|
|
| @classmethod |
| def from_pretrained(cls, path: str, device: str | None = None, compile: bool | None = None, token: str | None = None): |
| """path: a local export dir or a Hugging Face repo id. compile (default: on GPU): padded static encoders |
| under torch.compile, about 4x faster per call after a one-off warm-up.""" |
| from safetensors.torch import load_file |
| from transformers import ModernBertConfig |
| if not os.path.isdir(path): |
| from huggingface_hub import snapshot_download |
| path = snapshot_download(path, token=token) |
| cfg = json.load(open(os.path.join(path, "config.json"))) |
| device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| tok = load_tokenizer(path) |
| model = _Model.from_encoder(ModernBertConfig.from_dict(cfg["encoder"]), tok, v3=cfg["flags"]["v3"]) |
| model.load_state_dict({k: v.float() for k, v in load_file(os.path.join(path, "model.safetensors")).items()}) |
| model = model.to(device).eval() |
| if compile if compile is not None else device == "cuda": |
| C.encode_context, C.encode_slots = C.make_static(model.mlm.model) |
| return cls(model, tok, cfg["flags"], device) |
|
|
| @torch.no_grad() |
| def predict(self, query: str, tools: list[dict], T: int = 3, require: bool = False) -> list[dict]: |
| """Tool calls for `query`, best first. Empty list = no tool fits (unless require=True, like tool_choice="required"). |
| confidence = min(tool confidence, every argument's confidence). It ranks errors well enough to abstain on |
| (docs/results/benchmark.md: the least confident 10% abstaining takes our eval from 78.6 to 85.1).""" |
| ex = Example(query, [from_openai_tool(t) for t in tools], []) |
| pred = _predict(self.model, self.tok, ex, T=T, device=self.device, require=require, **self.flags) |
| return [{"name": c.name, "arguments": c.args, |
| "confidence": round(min([c.conf_tool, *c.conf_args.values()]), 4)} for c in pred.calls] |
|
|
