nullsilver/alpha-sys-1-1.6B

1"""alpha-sys-1 inference client: one file, no dependency on this repository, shipped in the2Hugging Face repos as `alpha_sys_1.py`. It renders questions exactly as the model was trained3on them and reads the answer distribution from one forward pass.4 5    from alpha_sys_1 import SystemOne6    m = SystemOne("nullsilver/alpha-sys-1-1.6B")7    m.ask({"type": "choice", "instructions": "Which team should handle this?",8           "criteria": {"billing": "payments, refunds", "technical": "bugs, outages", "sales": "pricing"}},9          state="Our API started returning 500 errors this morning.")10    # -> {"choice": "technical", "probabilities": {...}, "confidence": 0.71}11 12    m.system_one({"state": ..., "images": [...], "questions": {"q1": {...}, "q2": {...}}})13    # -> {"model": ..., "answers": {"q1": {...}, "q2": {...}}}   (TypeSafe's System One shape)14 15The questions of one request share their images and state, so that prefix is computed once16and only the questions run against its cache (`shared_prefix=False` turns this off).17Question types: choice (criteria = {option: description or None} or a list of options, up to1826), noul (a statement; criteria = {"true": ..., "false": ...} optional), score (criteria = the19levels, lowest first; the score is the expected level index). Images: a PIL image, a path, or20a data URL; small images are upscaled to 256 px as in training.21"""22 23from __future__ import annotations24 25import base6426import io27import math28import string29from typing import Any30 31import torch32from PIL import Image33from transformers import AutoModelForImageTextToText, AutoProcessor34 35IMAGE_SIDE = 25636 37 38def render_state(state: Any) -> str:39    if state is None:40        return ""41    if isinstance(state, str):42        return state43    if isinstance(state, dict):44        return "\n".join(f"{k}: {v}" for k, v in state.items())45    return str(state)46 47 48def render(state: Any, q: dict) -> tuple[str, list[str], list[str]]:49    """-> (user text, label tokens in listed order, answer-space keys in the same order)."""50    parts = [s for s in [render_state(state)] if s]51    t = q["type"]52    if t == "noul":53        c = q.get("criteria") or {}54        clar = "".join(f"\n{lab} means: {c[k]}" for lab, k in (("Yes", "true"), ("No", "false")) if c.get(k))55        parts.append(f"Statement: {q['instructions']}{clar}\nIs the statement true? Answer with Yes or No only.")56        return "\n\n".join(parts), ["No", "Yes"], ["no", "yes"]57    crit = q["criteria"]58    if t == "choice":59        items = list(crit.items()) if isinstance(crit, dict) else [(o, None) for o in crit]60        keys = [k for k, _ in items]61    else:62        items, keys = [(lvl, None) for lvl in crit], [str(i) for i in range(len(crit))]63    if len(items) > 26:64        raise ValueError("at most 26 options or levels per question")65    labels = list(string.ascii_uppercase[: len(items)])66    lines = [f"{lab}. {o}" + (f": {d}" if d else "") for lab, (o, d) in zip(labels, items)]67    parts.append(f"{q['instructions']}\n" + "\n".join(lines) + "\nAnswer with the letter only.")68    return "\n\n".join(parts), labels, keys69 70 71def load_image(im: Any) -> Image.Image:72    if isinstance(im, Image.Image):73        img = im74    elif isinstance(im, str) and im.startswith("data:"):75        img = Image.open(io.BytesIO(base64.b64decode(im.split(",", 1)[1])))76    else:77        img = Image.open(im)78    img = img.convert("RGB")79    if max(img.size) < IMAGE_SIDE:80        img = img.resize((IMAGE_SIDE, IMAGE_SIDE), Image.BICUBIC)81    return img82 83 84def confidence(p: list[float]) -> float:85    n = len(p)86    if n < 2:87        return 1.088    h = -sum(x * math.log(x) for x in p if x > 0)89    return round(max(0.0, 1 - h / math.log(n)), 4)90 91 92class SystemOne:93    def __init__(self, repo: str, revision: str | None = None, device: str | None = None, dtype=torch.bfloat16, temperature: float = 1.0,94                 shared_prefix: bool = True):95        """temperature: the label logits are divided by it (1.0 = the model as released; see fit_temperature).96        shared_prefix: when several questions share their images and state, run that prefix once and97        only the questions against its cache (same probabilities up to bfloat16 noise, several times98        faster with an image). False runs every question as its own full sequence."""99        self.repo, self.revision, self.temperature, self.shared_prefix = repo, revision, temperature, shared_prefix100        self.processor = AutoProcessor.from_pretrained(repo, revision=revision)101        self.processor.tokenizer.padding_side = "left"102        self.model = AutoModelForImageTextToText.from_pretrained(repo, revision=revision, dtype=dtype)103        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")104        self.model.to(self.device).eval()105        self._ids: dict[str, int] = {}106 107    def _label_id(self, label: str) -> int:108        if label not in self._ids:109            ids = self.processor.tokenizer.encode(label, add_special_tokens=False)110            assert len(ids) == 1, label111            self._ids[label] = ids[0]112        return self._ids[label]113 114    def _messages(self, items: list[tuple[Any, dict, list | None]]) -> tuple[list, list[list[str]]]:115        msgs, labels_per = [], []116        for state, q, images in items:117            text, labels, _ = render(state, q)118            content = [{"type": "image", "image": load_image(im)} for im in (images or [])] + [{"type": "text", "text": text}]119            msgs.append([{"role": "user", "content": content}])120            labels_per.append(labels)121        return msgs, labels_per122 123    def _probs(self, logits: torch.Tensor, labels_per: list[list[str]]) -> list[list[float]]:124        out = []125        for i, labels in enumerate(labels_per):126            ids = torch.tensor([self._label_id(lab) for lab in labels], device=logits.device)127            out.append(torch.softmax(logits[i, ids] / self.temperature, -1).tolist())128        return out129 130    @torch.inference_mode()131    def distributions(self, items: list[tuple[Any, dict, list | None]]) -> list[list[float]]:132        """items: (state, question, images or None) -> probabilities in the answer-space order.133        Every item is its own sequence, read at its last position; items that share images and134        state share the prefix's computation when `shared_prefix` is on."""135        msgs, labels_per = self._messages(items)136        if self.shared_prefix and len(items) > 1 and all(it[0] == items[0][0] and (it[2] or None) == (items[0][2] or None) for it in items):137            logits = self._shared_prefix_logits(msgs)138            if logits is not None:139                return self._probs(logits, labels_per)140        inputs = self.processor.apply_chat_template(141            msgs, add_generation_prompt=True, tokenize=True, return_dict=True,142            processor_kwargs={"return_tensors": "pt", "padding": True}).to(self.device)143        return self._probs(self.model(**inputs, logits_to_keep=1).logits[:, -1].float(), labels_per)144 145    def _shared_prefix_logits(self, msgs: list) -> torch.Tensor | None:146        """One pass over the longest common token prefix (images, state), then the question suffixes,147        right-padded, against that cache repeated across the batch. None when there is too little to share."""148        encs = [self.processor.apply_chat_template([m], add_generation_prompt=True, tokenize=True, return_dict=True,149                                                   processor_kwargs={"return_tensors": "pt"}) for m in msgs]150        ids = [e["input_ids"][0] for e in encs]151        L = min(len(x) for x in ids) - 1  # at least one token per suffix152        for x in ids[1:]:153            diff = (x[:L] != ids[0][:L]).nonzero()154            if len(diff):155                L = min(L, int(diff[0]))156        image_id = getattr(self.model.config, "image_token_id", None)157        if L < 64 or (image_id is not None and any((x[L:] == image_id).any() for x in ids)):158            return None159        n = len(ids)160        image_kw = {k: v.to(self.device) for k, v in encs[0].items() if k in ("pixel_values", "spatial_shapes", "pixel_attention_mask")}161        cache = self.model(input_ids=ids[0][:L][None].to(self.device), **image_kw, use_cache=True).past_key_values162        cache.reorder_cache(torch.zeros(n, dtype=torch.long, device=self.device))163        sufs = [x[L:] for x in ids]164        lens = torch.tensor([len(s) for s in sufs])165        width = int(lens.max())166        suffix = torch.full((n, width), self.processor.tokenizer.pad_token_id, dtype=torch.long)167        attention = torch.zeros((n, L + width), dtype=torch.long)168        attention[:, :L] = 1169        for i, s in enumerate(sufs):170            suffix[i, : len(s)] = s171            attention[i, L : L + len(s)] = 1172        out = self.model(input_ids=suffix.to(self.device), attention_mask=attention.to(self.device), past_key_values=cache,173                         cache_position=torch.arange(L, L + width, device=self.device), use_cache=True)174        return out.logits[torch.arange(n), (lens - 1).to(self.device)].float()175 176    def answer(self, q: dict, p: list[float]) -> dict:177        _, _, keys = render(None, q)178        if q["type"] == "choice":179            return {"type": "choice", "choice": keys[max(range(len(p)), key=p.__getitem__)],180                    "probabilities": dict(zip(keys, p)), "confidence": confidence(p)}181        if q["type"] == "noul":182            return {"type": "noul", "noul": p[1]}183        return {"type": "score", "score": sum(i * x for i, x in enumerate(p)),184                "legend": dict(zip(keys, q["criteria"])), "probabilities": p, "confidence": confidence(p)}185 186    def ask(self, q: dict, state: Any = None, images: list | None = None) -> dict:187        return self.answer(q, self.distributions([(state, q, images)])[0])188 189    def system_one(self, request: dict, batch: int = 16) -> dict:190        """A request in TypeSafe's System One shape: {state, images?, questions: {id: q}}."""191        state, images = request.get("state"), request.get("images")192        ids = list(request["questions"])193        answers = {}194        for s in range(0, len(ids), batch):195            chunk = ids[s : s + batch]196            ps = self.distributions([(state, request["questions"][i], images) for i in chunk])197            for i, p in zip(chunk, ps):198                answers[i] = self.answer(request["questions"][i], p)199        return {"model": self.repo + (f"@{self.revision}" if self.revision else ""), "answers": answers}200 201 202def fit_temperature(model: SystemOne, examples: list[tuple[Any, dict, list | None, int]], batch: int = 16) -> float:203    """One scalar that minimises NLL on labelled examples (state, question, images, index of the204    true answer in the answer space: option position, 0/1 for noul, level index for score).205    A few hundred examples are enough. Use it as SystemOne(..., temperature=T)."""206    old, model.temperature = model.temperature, 1.0207    try:208        probs, truth = [], []209        for s in range(0, len(examples), batch):210            chunk = examples[s : s + batch]211            probs += model.distributions([(st, q, im) for st, q, im, _ in chunk])212            truth += [t for _, _, _, t in chunk]213    finally:214        model.temperature = old215    logs = [[math.log(max(x, 1e-12)) for x in p] for p in probs]216 217    def nll(t: float) -> float:218        total = 0.0219        for lp, y in zip(logs, truth):220            z = [v / t for v in lp]221            m = max(z)222            total -= z[y] - (m + math.log(sum(math.exp(v - m) for v in z)))223        return total / len(logs)224 225    lo, hi = math.log(0.05), math.log(20.0)  # golden-section search on log T226    g = (math.sqrt(5) - 1) / 2227    a, b = hi - g * (hi - lo), lo + g * (hi - lo)228    fa, fb = nll(math.exp(a)), nll(math.exp(b))229    for _ in range(60):230        if fa < fb:231            hi, b, fb = b, a, fa232            a = hi - g * (hi - lo)233            fa = nll(math.exp(a))234        else:235            lo, a, fa = a, b, fb236            b = lo + g * (hi - lo)237            fb = nll(math.exp(b))238    return round(math.exp((lo + hi) / 2), 3)239