File size: 2,947 Bytes
6fd3740 | 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 | """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": # process-wide swap, as in evaluate --compile
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]
|