"""Snap 1: typed decisions in one forward pass. pip install "laya==0.3.4" torch transformers safetensors huggingface_hub from snap import Snap snap = Snap("Otron/snap-1") snap.decide( state={"message": "my card got charged twice for the same order"}, questions={ "topic": {"type": "choice", "instructions": "What is the message about?", "criteria": {"billing": "Payments, charges, refunds", "shipping": "Delivery", "other": "Anything else"}}, "urgent": {"type": "noul", "instructions": "The customer needs a reply today"}, }, ) # {"topic": {"choice": "billing", "confidence": 0.997, "probabilities": {...}}, # "urgent": {"noul": 0.42, "confidence": 0.58}} Questions use the common typed-decision wire format: `noul` (probability a statement is true), `choice` (one key from `criteria`, a dict of key -> description, or a list of keys) and `score` (one of the ordered levels in `criteria`). All questions about one state are answered together: they are packed into as few sequences as fit, each read once by the encoder. """ import json import math import os from typing import Dict, List, Optional, Union import torch from laya.agent import Agent from laya.common import QTYPE_NAMES, QTYPES, render_options, serialize_state NONE_TYPE = 3 # token type of [CLS], the state and the final [SEP] OPTION_TOKENS = 48 # per-option token cap QUESTION_TOKENS = 256 # per-question instruction cap class Snap: def __init__(self, repo_or_path: str = "Otron/snap-1", device: Optional[str] = None, token: Optional[str] = None): path = repo_or_path if not os.path.isdir(path): from huggingface_hub import snapshot_download path = snapshot_download(repo_or_path, token=token, ignore_patterns=["coreml/*"]) agent = Agent(path, device=device) self.cfg = agent.cfg self.tok = agent.tok self.model = agent.model.eval() self.device = next(self.model.parameters()).device self.max_len = int(self.cfg.get("max_len", 4096)) self.head_max_len = int(self.cfg.get("head_max_len", 3584)) # the embedding Snap gives [CLS] and state tokens in its packed layout from safetensors.torch import load_file none = next(iter(load_file(os.path.join(path, "extra.safetensors")).values())) self.type_table = torch.cat([self.model.type_emb.weight.detach(), none[None].to(self.device)], 0) self.temps = dict(self.cfg.get("temperature_by_options") or {}) self.temps.update({QTYPE_NAMES[i]: t for i, t in enumerate(self.cfg.get("temperature") or [])}) # ------------------------------------------------------------------ public @torch.no_grad() def decide(self, state: Union[str, dict, list], questions: Dict[str, Dict]) -> Dict[str, Dict]: """Answer every question about `state`. Returns {question id: answer}.""" mask = self.tok.mask_token state_ids = self._ids(serialize_state(state).replace(mask, " ")) blocks = [] for qid, qdef in questions.items(): q = Agent._to_internal(qdef) tokens, markers = self._block(q) if len(tokens) <= self.head_max_len: blocks.append((qid, q, tokens, markers)) answers = {} for group in self._pack(blocks): ids, types, where = [self.tok.cls_token_id], [NONE_TYPE], [] for qid, q, tokens, markers in group: start = len(ids) ids += tokens types += [QTYPES[q["t"]]] * len(tokens) where.append((qid, q, [start + m for m in markers])) room = max(0, self.max_len - len(ids) - 1) ids += state_ids[:room] + [self.tok.sep_token_id] types += [NONE_TYPE] * (len(ids) - len(types)) scores = self._scores(ids, types) for qid, q, positions in where: answers[qid] = self._answer(q, [float(scores[p]) for p in positions]) return answers # ------------------------------------------------------------------ internals def _ids(self, text: str) -> List[int]: return self.tok(text, add_special_tokens=False)["input_ids"] def _block(self, q: Dict): mask = self.tok.mask_token head = self._ids("%s question: %s" % (q["t"], str(q["ins"]).replace(mask, " ")))[:QUESTION_TOKENS] body, markers = [], [] for text in render_options(q): markers.append(len(head) + 1 + len(body)) body += [self.tok.mask_token_id] + self._ids(" " + text.replace(mask, " "))[:OPTION_TOKENS] return head + [self.tok.sep_token_id] + body + [self.tok.sep_token_id], markers def _pack(self, blocks): groups, current, used = [], [], 0 for b in blocks: if current and used + len(b[2]) > self.head_max_len: groups.append(current) current, used = [], 0 current.append(b) used += len(b[2]) if current: groups.append(current) return groups def _scores(self, ids: List[int], types: List[int]) -> torch.Tensor: m = self.model x = torch.tensor([ids], device=self.device) attn = torch.ones_like(x) h = m.encoder(input_ids=x, attention_mask=attn).last_hidden_state h = h + self.type_table[torch.tensor([types], device=self.device)].to(h.dtype) if m.head is not None: for layer in m.head.layers: h = layer(h, src_key_padding_mask=~attn.bool()) return m.scorer(h).squeeze(-1)[0].float().cpu() def _answer(self, q: Dict, logits: List[float]) -> Dict: k = len(logits) size = "2" if k <= 2 else "3-5" if k <= 5 else "6-10" if k <= 10 else "11+" t = self.temps.get("%s:%s" % (q["t"], size), self.temps.get(q["t"], 1.0)) z = [v / max(1e-3, t) for v in logits] top = max(z) e = [math.exp(v - top) for v in z] s = sum(e) p = [v / s for v in e] best = max(range(k), key=p.__getitem__) if q["t"] == "noul": return {"noul": p[1], "confidence": max(p)} keys = list(q["crit"].keys()) if q["t"] == "choice" else [str(i) for i in range(k)] out = {"confidence": p[best], "probabilities": dict(zip(keys, p))} out["choice" if q["t"] == "choice" else "score"] = keys[best] if q["t"] == "choice" else best return out if __name__ == "__main__": import sys r = Snap(sys.argv[1] if len(sys.argv) > 1 else "Otron/snap-1") print(json.dumps(r.decide( {"message": "can you turn the volume down a bit"}, {"about_audio": {"type": "noul", "instructions": "The message is about sound or volume"}, "intent": {"type": "choice", "instructions": "What does the user want?", "criteria": {"volume_down": "Lower the volume", "volume_up": "Raise the volume", "other": "Something else"}}}), indent=1))