mini-beatrix-2.5s: the 2s core (bit-identical) packaged with its arm library — 13 mountable rows over 11 weight files, a self-contained arm runtime verified bit-exact against amoe-lora, and an automodel that loads under current transformers
fc9a663 verified Download arms.py from AbstractPhil/mini-beatrix-2.5s: direct link, hf CLI and curl.
- Browser
- Download file 23.3 kB
-
https://huggingface.co/AbstractPhil/mini-beatrix-2.5s/resolve/main/arms.py
- Command line
-
hf download hf://AbstractPhil/mini-beatrix-2.5s/arms.py
-
curl -L -o arms.py https://huggingface.co/AbstractPhil/mini-beatrix-2.5s/resolve/main/arms.py
23.3 kB
| """arms.py — the detachable arm runtime for mini-beatrix-2.5s. | |
| An *arm* is a small aleph-addressed patch head placed after every decoder | |
| block of one exact frozen core. It is not a low-rank matrix: each site | |
| projects the block output into `n_slots` slots of dimension D, reads them | |
| through a K-atom aleph address (a closed-form signed read, never a | |
| selector), and adds a gated patch back into the residual stream: | |
| slots = proj(x).view(B, n, n_slots, D) | |
| m_hat = sum_k sinh(u_k) A_k / sum_k cosh(u_k), u = (x_hat . A) / tau | |
| x = x + sigmoid(gate) * consume(m_hat) | |
| `mount()` makes the live core carry exactly one recipe and returns a | |
| handle; `handle.detach(verify=True)` restores the untouched blocks and | |
| asserts the restored logits are BIT-EXACT with a pre-attach fingerprint. | |
| That exactness is the point of the design: an arm is a guest, never a | |
| migration of the core. | |
| Recipes carried by this package: | |
| anchor one arm, always on | |
| stack several anchor arms applied always-on, nested inner to outer, | |
| with no mixer between them (the ungated composition) | |
| routed a dense signed dispatch over member arms, optionally with a | |
| detached bounded per-member gain read from a causal usage EMA | |
| THE DAMPING LAW (do not "fix" it): the dispatch denominator sums cosh | |
| over ALL members, including masked ones, and masking a member removes | |
| its delta WITHOUT renormalizing. Masked is not solo. | |
| This file is self-contained (torch + safetensors). It is a vendored, | |
| verified re-expression of the amoe-lora runtime the arms were trained | |
| with (github.com/AbstractEyes/amoe-lora): the packaged arms were checked | |
| to produce bit-identical logits through both paths before release. | |
| """ | |
| from __future__ import annotations | |
| import contextlib | |
| import json | |
| import os | |
| from dataclasses import dataclass | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| ARM_TYPES = ("anchor", "stack", "routed") | |
| # --------------------------------------------------------------- geometry | |
| class AdapterSpec: | |
| n_slots: int = 16 | |
| K: int = 64 | |
| D: int = 4 | |
| tau: float = 0.1 | |
| hidden: int = 178 | |
| gate_init: float = -3.0 | |
| class AlephAddress(nn.Module): | |
| """m_hat(x) = sum_k sinh(u_k) A_k / sum_k cosh(u_k), max-|u| stabilized. | |
| `home` snapshots the codebook at init and is required in every arm | |
| state (the drift gauge).""" | |
| def __init__(self, K: int = 64, D: int = 4, tau: float = 0.1): | |
| super().__init__() | |
| self.K, self.D, self.tau = K, D, tau | |
| self.codebook = nn.Parameter(F.normalize(torch.randn(K, D), dim=-1)) | |
| self.register_buffer("home", self.codebook.detach().clone()) | |
| def m_hat(self, x: torch.Tensor) -> torch.Tensor: | |
| A = F.normalize(self.codebook, dim=-1) | |
| u = (F.normalize(x, dim=-1) @ A.transpose(-1, -2)) / self.tau | |
| m = u.abs().amax(dim=-1, keepdim=True) | |
| ep, en = torch.exp(u - m), torch.exp(-u - m) | |
| return ((ep - en) @ A) / (ep + en).sum(dim=-1, keepdim=True) | |
| def drift(self) -> float: | |
| """Mean angular drift (radians) of the codebook from `home`.""" | |
| A = F.normalize(self.codebook, dim=-1) | |
| H = F.normalize(self.home, dim=-1) | |
| return float(torch.arccos((A * H).sum(-1).clamp(-1.0, 1.0)).mean()) | |
| class SquaredReLU(nn.Module): | |
| def forward(self, x): | |
| return F.relu(x) ** 2 | |
| class RelayPatchwork(nn.Module): | |
| """One arm site: projection -> aleph address -> squared-ReLU patch head, | |
| added through a sigmoid gate.""" | |
| def __init__(self, d: int, spec: AdapterSpec | None = None): | |
| super().__init__() | |
| s = spec or AdapterSpec() | |
| self.spec = s | |
| self.n_slots = s.n_slots | |
| self.proj = nn.Linear(d, s.n_slots * s.D, bias=False) | |
| self.addr = AlephAddress(s.K, s.D, s.tau) | |
| self.consume = nn.Sequential( | |
| nn.Linear(s.n_slots * s.D, s.hidden), SquaredReLU(), | |
| nn.LayerNorm(s.hidden), nn.Linear(s.hidden, d)) | |
| self.gate = nn.Parameter(torch.tensor(float(s.gate_init))) | |
| def forward(self, x): | |
| B, n, _ = x.shape | |
| slots = self.proj(x).view(B, n, self.n_slots, self.spec.D) | |
| feats = self.addr.m_hat(slots).reshape(B, n, -1) | |
| return x + torch.sigmoid(self.gate) * self.consume(feats) | |
| def spec_from_state(state: dict, prefix: str | None = None) -> AdapterSpec: | |
| """Arm geometry read from the arm's own tensors (the library holds more | |
| than one shape, and modules are built before the state loads).""" | |
| if prefix is None: | |
| blocks = sorted(int(k.split(".")[0]) for k in state | |
| if k.endswith(".addr.codebook") and k.split(".")[0].isdigit()) | |
| prefix = f"{blocks[0]}." if blocks else "0." | |
| cb = state.get(f"{prefix}addr.codebook") | |
| proj = state.get(f"{prefix}proj.weight") | |
| cons = state.get(f"{prefix}consume.0.weight") | |
| if cb is None or proj is None or cons is None: | |
| return AdapterSpec() | |
| D = int(cb.shape[1]) | |
| return AdapterSpec(n_slots=int(proj.shape[0] // D), K=int(cb.shape[0]), | |
| D=D, hidden=int(cons.shape[0])) | |
| # --------------------------------------------------------------- wrappers | |
| class BlockWithAdapter(nn.Module): | |
| """Wraps one decoder block. `enabled=False` returns the block output | |
| untouched — the single-arm half of the toggle law.""" | |
| def __init__(self, block: nn.Module, adapter: RelayPatchwork): | |
| super().__init__() | |
| self.block = block | |
| self.adapter = adapter | |
| self.enabled = True | |
| def forward(self, *args, **kwargs): | |
| out = self.block(*args, **kwargs) | |
| if not self.enabled: | |
| return out | |
| if isinstance(out, tuple): | |
| return (self.adapter(out[0]),) + out[1:] | |
| return self.adapter(out) | |
| def prefill(self, *args, **kwargs): | |
| out, cache = self.block.prefill(*args, **kwargs) | |
| return (self.adapter(out) if self.enabled else out), cache | |
| def step(self, *args, **kwargs): | |
| out = self.block.step(*args, **kwargs) | |
| if not self.enabled: | |
| return out | |
| if isinstance(out, tuple): | |
| return (self.adapter(out[0]),) + out[1:] | |
| return self.adapter(out) | |
| class AnchorDispatch(nn.Module): | |
| """Dense signed sinh/cosh dispatch over frozen member arms. | |
| out = x + sum_a (sinh u_a / sum_all cosh u) (arm_a(x) - x) | |
| No top-k, no softmax over members, no load balancing. The denominator | |
| sums over every member, masked or not (the damping law).""" | |
| def __init__(self, anchors: nn.ModuleList, d: int, emb: int = 64, | |
| tau: float = 0.1): | |
| super().__init__() | |
| self.anchors = anchors | |
| self.tau = tau | |
| A = len(anchors) | |
| self.register_buffer("key_proj", torch.zeros(emb, d)) | |
| self.dispatch = nn.Parameter(torch.zeros(A, emb)) | |
| self.enabled = [True] * A | |
| def forward(self, x): | |
| emb = F.normalize(x @ self.key_proj.T, dim=-1) | |
| C = F.normalize(self.dispatch, dim=-1) | |
| u = (emb @ C.T) / self.tau | |
| m = u.abs().amax(dim=-1, keepdim=True) | |
| ep, en = torch.exp(u - m), torch.exp(-u - m) | |
| w = ep - en | |
| z = (ep + en).sum(dim=-1, keepdim=True) # ALL members (law) | |
| out = x | |
| for k, a in enumerate(self.anchors): | |
| if self.enabled[k]: | |
| out = out + (w[..., k:k + 1] / z) * (a(x) - x) | |
| return out | |
| class GainDispatch(AnchorDispatch): | |
| """AnchorDispatch with a detached bounded per-member gain on the | |
| numerator (the denominator is untouched, the signs are untouched): | |
| s_a(t) = cosh u_a / sum_all cosh u (the member's share) | |
| ubar_a = causal EMA of s_a over bytes, rate rho, start 1/A | |
| alpha_a = clamp(1 + gamma (A ubar_a - 1), a_min, a_max) | |
| out = x + sum_a alpha_a (sinh u_a / sum_all cosh u)(arm_a(x) - x) | |
| mode 'growth' lets the mean gain rise toward the cap; 'conserved' | |
| renormalizes alpha to mean 1 over members at every byte. `coherent` | |
| reads alpha once at the first site and shares it ('alpha'), and with | |
| 'full' shares the first site's signed routing weight too.""" | |
| def configure(self, *, mode="growth", rho=1 / 16, gamma=1.5, a_min=0.5, | |
| a_max=2.0, coherent=None, role="ref", holder=None): | |
| self.mode, self.rho, self.gamma = mode, float(rho), float(gamma) | |
| self.a_min, self.a_max = float(a_min), float(a_max) | |
| self.coherent, self.role = coherent, role | |
| self.holder = {} if holder is None else holder | |
| self._state, self._call = None, "forward" | |
| return self | |
| def reset_state(self): | |
| self._state = None | |
| def _alpha(self, ep, en, z): | |
| with torch.no_grad(): | |
| s = ((ep + en) / z).float() | |
| B, T, A = s.shape | |
| st = (self._state if (self._call == "step" and self._state is not None | |
| and self._state.shape[0] == B) else None) | |
| if st is None: | |
| st = torch.full((B, A), 1.0 / A, device=s.device, dtype=torch.float32) | |
| if T == 1: | |
| st = st + self.rho * (s[:, 0] - st) | |
| ubar = st.unsqueeze(1) | |
| elif T <= 4200: | |
| q = 1.0 - self.rho | |
| t = torch.arange(1, T + 1, device=s.device, dtype=torch.float64) | |
| acc = torch.cumsum((q ** -t).view(1, T, 1) * s.double(), dim=1) | |
| ubar = ((q ** t).view(1, T, 1) | |
| * (st.double().unsqueeze(1) + self.rho * acc)).float() | |
| st = ubar[:, -1] | |
| else: | |
| ubar = torch.empty_like(s) | |
| for i in range(T): | |
| st = st + self.rho * (s[:, i] - st) | |
| ubar[:, i] = st | |
| self._state = st.clone() if self._call in ("prefill", "step") else None | |
| alpha = (1.0 + self.gamma * (A * ubar - 1.0)).clamp(self.a_min, self.a_max) | |
| if self.mode == "conserved": | |
| alpha = alpha * (A / alpha.sum(dim=-1, keepdim=True)) | |
| return alpha.to(z.dtype) | |
| def forward(self, x): | |
| shared_route = self.coherent == "full" and self.role != "ref" | |
| if not shared_route: | |
| emb = F.normalize(x @ self.key_proj.T, dim=-1) | |
| C = F.normalize(self.dispatch, dim=-1) | |
| u = (emb @ C.T) / self.tau | |
| m = u.abs().amax(dim=-1, keepdim=True) | |
| ep, en = torch.exp(u - m), torch.exp(-u - m) | |
| z = (ep + en).sum(dim=-1, keepdim=True) | |
| g = (ep - en) / z | |
| if self.coherent and self.role != "ref": | |
| alpha = self.holder["alpha"] | |
| if shared_route: | |
| g = self.holder["g"] | |
| else: | |
| alpha = self._alpha(ep, en, z) | |
| if self.coherent: | |
| self.holder["alpha"], self.holder["g"] = alpha, g | |
| out = x | |
| for k, a in enumerate(self.anchors): | |
| if self.enabled[k]: | |
| out = out + (g[..., k:k + 1] * alpha[..., k:k + 1]) * (a(x) - x) | |
| return out | |
| class BlockWithDispatch(nn.Module): | |
| def __init__(self, block: nn.Module, disp: AnchorDispatch): | |
| super().__init__() | |
| self.block = block | |
| self.disp = disp | |
| def forward(self, *args, **kwargs): | |
| out = self.block(*args, **kwargs) | |
| if isinstance(out, tuple): | |
| return (self.disp(out[0]),) + out[1:] | |
| return self.disp(out) | |
| def prefill(self, *args, **kwargs): | |
| out, cache = self.block.prefill(*args, **kwargs) | |
| self.disp._call = getattr(self.disp, "_call", "forward") | |
| if isinstance(self.disp, GainDispatch): | |
| self.disp._call = "prefill" | |
| try: | |
| return self.disp(out), cache | |
| finally: | |
| self.disp._call = "forward" | |
| return self.disp(out), cache | |
| def step(self, *args, **kwargs): | |
| out = self.block.step(*args, **kwargs) | |
| if isinstance(self.disp, GainDispatch): | |
| self.disp._call = "step" | |
| try: | |
| return self.disp(out[0]) if isinstance(out, tuple) else self.disp(out) | |
| finally: | |
| self.disp._call = "forward" | |
| if isinstance(out, tuple): | |
| return (self.disp(out[0]),) + out[1:] | |
| return self.disp(out) | |
| # --------------------------------------------------------------- handles | |
| def _probe(aleph) -> torch.Tensor: | |
| """Deterministic logits fingerprint on a tiny fixed input.""" | |
| dev = next(aleph.parameters()).device | |
| ids = torch.arange(8, dtype=torch.long, device=dev).unsqueeze(0) % 7 + 1 | |
| was_training = aleph.training | |
| aleph.eval() | |
| with torch.no_grad(): | |
| out = aleph(ids) | |
| if was_training: | |
| aleph.train() | |
| logits = out[0] if isinstance(out, (tuple, list)) else out.logits | |
| return logits.detach().float().cpu() | |
| class ArmHandle: | |
| """One mounted recipe. `detach(verify=True)` restores the core and | |
| proves the restoration is bit-exact.""" | |
| def __init__(self, aleph, original, names, wrapped, fingerprint, card=None): | |
| self.aleph = aleph | |
| self._original = original | |
| self.names = list(names) | |
| self._blocks = wrapped | |
| self._fingerprint = fingerprint | |
| self.card = card or {} | |
| # -- masking (the toggle law) -------------------------------------- | |
| def set_mask(self, mask: dict) -> None: | |
| enabled = [bool(mask.get(n, True)) for n in self.names] | |
| for b in self._blocks: | |
| if isinstance(b, BlockWithDispatch): | |
| b.disp.enabled = list(enabled) | |
| else: | |
| b.enabled = enabled[0] | |
| def enable(self, *names): | |
| self.set_mask({n: (n in names or not names) for n in self.names}) | |
| def disable(self, *names): | |
| self.set_mask({n: n not in names for n in self.names}) | |
| def _snapshot(self): | |
| return [list(b.disp.enabled) if isinstance(b, BlockWithDispatch) else b.enabled | |
| for b in self._blocks] | |
| def _restore(self, snap): | |
| for b, s in zip(self._blocks, snap): | |
| if isinstance(b, BlockWithDispatch): | |
| b.disp.enabled = list(s) | |
| else: | |
| b.enabled = s | |
| def only(self, *names): | |
| prev = self._snapshot() | |
| self.set_mask({n: n in names for n in self.names}) | |
| try: | |
| yield self | |
| finally: | |
| self._restore(prev) | |
| def all_off(self): | |
| prev = self._snapshot() | |
| self.set_mask({n: False for n in self.names}) | |
| try: | |
| yield self | |
| finally: | |
| self._restore(prev) | |
| # -- decode state --------------------------------------------------- | |
| def reset_state(self): | |
| """Forget the carried gain state (a fresh context starts over).""" | |
| for b in self._blocks: | |
| if isinstance(b, BlockWithDispatch): | |
| getattr(b.disp, "reset_state", lambda: None)() | |
| # -- detach --------------------------------------------------------- | |
| def detach(self, *, verify: bool = True): | |
| self.aleph.blocks = self._original | |
| if verify: | |
| post = _probe(self.aleph) | |
| if not torch.equal(post, self._fingerprint): | |
| raise RuntimeError( | |
| "detach verification FAILED: the restored core's logits " | |
| "are not bit-exact with the pre-mount fingerprint") | |
| return self.aleph | |
| class StackHandle: | |
| """Nested always-on handles, detached outermost first.""" | |
| def __init__(self, handles, names=None, card=None): | |
| self.handles = list(handles) | |
| self.names = list(names or [h.names[0] for h in self.handles]) | |
| self.card = card or {} | |
| def reset_state(self): | |
| for h in self.handles: | |
| h.reset_state() | |
| def set_mask(self, mask: dict) -> None: | |
| for h, n in zip(self.handles, self.names): | |
| h.set_mask({h.names[0]: bool(mask.get(n, True))}) | |
| def all_off(self): | |
| prev = [h._snapshot() for h in self.handles] | |
| for h in self.handles: | |
| h.set_mask({h.names[0]: False}) | |
| try: | |
| yield self | |
| finally: | |
| for h, p in zip(self.handles, prev): | |
| h._restore(p) | |
| def detach(self, *, verify: bool = True): | |
| for h in reversed(self.handles): | |
| h.detach(verify=verify) | |
| return self.handles[0].aleph if self.handles else None | |
| # --------------------------------------------------------------- loading | |
| def load_arm_file(path: str): | |
| """-> (tensors, meta). Reads the packaged safetensors arm; also accepts | |
| an amoe .pt anchor/dispatch checkpoint for local experiments.""" | |
| if str(path).endswith(".safetensors"): | |
| from safetensors.torch import load_file | |
| from safetensors import safe_open | |
| with safe_open(path, framework="pt") as f: | |
| meta = dict(f.metadata() or {}) | |
| card = json.loads(meta.get("arm", "{}")) | |
| return load_file(path), card | |
| blob = torch.load(path, map_location="cpu", weights_only=True) | |
| if "adapters" in blob: | |
| return blob["adapters"], dict(blob.get("meta") or {}) | |
| if "dispatch" in blob: # flatten to "<block>.<key>" | |
| flat = {f"{i}.{k}": v for i, st in enumerate(blob["dispatch"]) | |
| for k, v in st.items()} | |
| return flat, dict(blob.get("meta") or {}) | |
| raise ValueError(f"unrecognized arm file: {path}") | |
| def _per_block(state: dict, i: int) -> dict: | |
| pre = f"{i}." | |
| return {k[len(pre):]: v for k, v in state.items() if k.startswith(pre)} | |
| def _target_blocks(aleph): | |
| return list(aleph.blocks) | |
| # --------------------------------------------------------------- mounting | |
| def mount_anchor(aleph, state: dict, name: str = "arm", card=None) -> ArmHandle: | |
| layers = _target_blocks(aleph) | |
| d = int(aleph.cfg.d_model) | |
| fingerprint = _probe(aleph) | |
| spec = spec_from_state(state) | |
| blocks, wrapped = [], [] | |
| for i, layer in enumerate(layers): | |
| st = _per_block(state, i) | |
| if not st: | |
| blocks.append(layer) | |
| continue | |
| p = next(layer.parameters()) | |
| a = RelayPatchwork(d, spec) | |
| a.load_state_dict(st, strict=True) | |
| a.to(device=p.device, dtype=p.dtype) | |
| w = BlockWithAdapter(layer, a) | |
| blocks.append(w) | |
| wrapped.append(w) | |
| original = aleph.blocks | |
| aleph.blocks = nn.ModuleList(blocks) | |
| return ArmHandle(aleph, original, [name], wrapped, fingerprint, card) | |
| def mount_routed(aleph, state: dict, names, gain=None, mode="growth", | |
| coherent=None, card=None) -> ArmHandle: | |
| layers = _target_blocks(aleph) | |
| d = int(aleph.cfg.d_model) | |
| fingerprint = _probe(aleph) | |
| A = len(names) | |
| member0 = {k[len("anchors.0."):]: v for k, v in _per_block(state, 0).items() | |
| if k.startswith("anchors.0.")} | |
| spec = spec_from_state(member0, prefix="") | |
| holder, blocks, wrapped = {}, [], [] | |
| for i, layer in enumerate(layers): | |
| st = _per_block(state, i) | |
| p = next(layer.parameters()) | |
| members = nn.ModuleList([RelayPatchwork(d, spec) for _ in range(A)]) | |
| emb = int(st["key_proj"].shape[0]) | |
| dp = AnchorDispatch(members, d, emb=emb) | |
| dp.load_state_dict(st, strict=True) | |
| dp.to(device=p.device, dtype=p.dtype) | |
| if gain: | |
| dp.__class__ = GainDispatch | |
| dp.configure(mode=mode, rho=gain["rho"], gamma=gain["gamma"], | |
| a_min=gain["a_min"], a_max=gain["a_max"], | |
| coherent=coherent, role="ref" if i == 0 else "follow", | |
| holder=holder) | |
| w = BlockWithDispatch(layer, dp) | |
| blocks.append(w) | |
| wrapped.append(w) | |
| original = aleph.blocks | |
| aleph.blocks = nn.ModuleList(blocks) | |
| return ArmHandle(aleph, original, list(names), wrapped, fingerprint, card) | |
| def mount(aleph, card: dict, files) -> "ArmHandle | StackHandle": | |
| """Attach one recipe. `card` is an arms/index.json row; `files` maps the | |
| row's weight keys to local paths (a stack passes one file per member, | |
| inner first).""" | |
| kind = (card.get("mount") or {}).get("type", "anchor") | |
| if kind == "anchor": | |
| state, _ = load_arm_file(files[0]) | |
| return mount_anchor(aleph, state, card.get("name", "arm"), card) | |
| if kind == "stack": | |
| members = (card.get("mount") or {}).get("members") or [] | |
| handles = [] | |
| try: | |
| for f, n in zip(files, members or [f"member{i}" for i in range(len(files))]): | |
| state, _ = load_arm_file(f) | |
| handles.append(mount_anchor(aleph, state, n, card)) | |
| except Exception: | |
| StackHandle(handles).detach(verify=False) # never leave half a stack on | |
| raise | |
| return StackHandle(handles, members, card) | |
| if kind == "routed": | |
| state, _ = load_arm_file(files[0]) | |
| m = card.get("mount") or {} | |
| return mount_routed(aleph, state, m.get("members") or [], | |
| gain=m.get("gain"), mode=m.get("gain_mode", "growth"), | |
| coherent=m.get("coherent"), card=card) | |
| raise ValueError(f"unknown mount type {kind!r} (packaged types: {ARM_TYPES})") | |
| # --------------------------------------------------------------- frames | |
| CHAT_HEADER = ("A conversation with Beatrix, a small byte-level " | |
| "language model still in pretraining.\n") | |
| USER_TAG = "User: " | |
| ASSISTANT_TAG = "Beatrix:" | |
| TEXT_STOPS = ("\nUser:", "\nBeatrix:", "User:", "Beatrix:", "\n\n") | |
| SPECIALS = {"DOC": 0xFF, "USER": 0xFE, "MODEL": 0xFD, "END": 0xFC, "SYS": 0xFB} | |
| def render_text_frame(prompt: str, history=None) -> bytes: | |
| """The plain text frame every packaged arm was trained in: | |
| header + 'User: ...\\n' + 'Beatrix:' — the reply follows the tag.""" | |
| out = CHAT_HEADER | |
| for turn in (history or []): | |
| tag = USER_TAG if turn.get("role") == "user" else ASSISTANT_TAG + " " | |
| out += f"{tag}{turn.get('content', '')}\n" | |
| out += f"{USER_TAG}{prompt}\n{ASSISTANT_TAG}" | |
| return out.encode("utf-8") | |
| def render_specials_frame(prompt: str, system: str | None = None, | |
| history=None) -> bytes: | |
| """The core's own annealed frame: [SYS] .. [END] [USER] .. [END] [MODEL]. | |
| Encoded text can never contain a special, so the frame is unforgeable.""" | |
| S = SPECIALS | |
| ids = [] | |
| if system: | |
| ids += [S["SYS"]] + list(system.encode("utf-8")) + [S["END"]] | |
| for turn in (history or []): | |
| tag = S["USER"] if turn.get("role") == "user" else S["MODEL"] | |
| ids += [tag] + list(str(turn.get("content", "")).encode("utf-8")) + [S["END"]] | |
| ids += [S["USER"]] + list(prompt.encode("utf-8")) + [S["END"], S["MODEL"]] | |
| return bytes(ids) | |
| def cut_text_reply(raw: bytes) -> bytes: | |
| """Cut a text-frame continuation at the first stop string.""" | |
| best = len(raw) | |
| for s in TEXT_STOPS: | |
| j = raw.find(s.encode("utf-8")) | |
| if j != -1: | |
| best = min(best, j) | |
| return raw[:best] | |
| def cut_specials_reply(raw: bytes) -> bytes: | |
| j = raw.find(bytes([SPECIALS["END"]])) | |
| return raw if j == -1 else raw[:j] | |