snap.py

"""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))