| """HF AutoModel wrapper for mini-beatrix-3 — the 32-block full-splat craft |
| (a CausalSplatHUB constellation and an anchored expert bank in every block; |
| 64.4B-byte mission, completed 2026-10-05), packaged with its detachable |
| stage arms. |
| |
| from transformers import AutoModelForCausalLM |
| m = AutoModelForCausalLM.from_pretrained("AbstractPhil/mini-beatrix-3", |
| trust_remote_code=True) |
| m.say("Who are you?") # arms off: the bare core |
| |
| m.mount_arm("s3_rules") # arms on: one stage arm ... |
| m.mount_arm("stages-1-4") # ... or a stack of them |
| m.detach_arm() # bit-exact core restored |
| |
| # arms on from the first call |
| m = AutoModelForCausalLM.from_pretrained("AbstractPhil/mini-beatrix-3", |
| trust_remote_code=True, |
| default_arm="stages-1-4") |
| |
| The model reads raw UTF-8 bytes: input_ids are byte values 0..255. The |
| default weight file is bf16; the fp32 master weights load with |
| variant="fp32". The core's weights are the same with arms off and on: an |
| arm is a guest on the core, never a change to it. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import os |
|
|
| import torch |
| from transformers import PretrainedConfig, PreTrainedModel |
| from transformers.generation import GenerationMixin |
| from transformers.modeling_outputs import CausalLMOutputWithPast |
|
|
| from .arms import (SPECIALS, TEXT_STOPS, cut_specials_reply, cut_text_reply, |
| mount, render_specials_frame, render_text_frame) |
| from .alephlm import AlephLM |
| from .presets import AlephLMConfig |
|
|
| ARMS_INDEX = "arms/index.json" |
|
|
|
|
| class MiniBeatrixConfig(PretrainedConfig): |
| model_type = "mini-beatrix" |
| attribute_map = {"num_hidden_layers": "n_layers", |
| "hidden_size": "d_model", |
| "num_attention_heads": "n_heads", |
| "max_position_embeddings": "context"} |
|
|
| def __init__(self, name="mini-beatrix-3", d_model=1024, n_layers=32, |
| n_heads=16, context=4096, vocab_size=256, |
| tokenizer="byte-trigram", hub_layers=tuple(range(32)), |
| hub_K=64, hub_D=128, tau=0.1, bank_experts=3, bank_ff=1024, |
| head_K=256, head_D=256, gate_init=-3.0, |
| tie_embeddings=False, hub_chunk=256, hub_const=4, |
| hub_ckpt=0, trunk=None, default_arm=None, **kwargs): |
| self.name = name |
| self.d_model = d_model |
| self.n_layers = n_layers |
| self.n_heads = n_heads |
| self.context = context |
| self.vocab_size = vocab_size |
| self.tokenizer = tokenizer |
| self.hub_layers = list(hub_layers) |
| self.hub_K = hub_K |
| self.hub_D = hub_D |
| self.tau = tau |
| self.bank_experts = bank_experts |
| self.bank_ff = bank_ff |
| self.head_K = head_K |
| self.head_D = head_D |
| self.gate_init = gate_init |
| self.tie_embeddings = tie_embeddings |
| self.hub_chunk = hub_chunk |
| self.hub_const = hub_const |
| self.hub_ckpt = hub_ckpt |
| |
| self.trunk = trunk or {"craft": "mini-beatrix-3", "step": 245674, |
| "base_model_id": "alephllm/mini-beatrix-3@step245674"} |
| |
| self.default_arm = default_arm |
| super().__init__(**kwargs) |
|
|
| def to_aleph(self) -> AlephLMConfig: |
| return AlephLMConfig( |
| name=self.name, d_model=self.d_model, n_layers=self.n_layers, |
| n_heads=self.n_heads, context=self.context, |
| vocab_size=self.vocab_size, tokenizer=self.tokenizer, |
| hub_layers=tuple(self.hub_layers), hub_K=self.hub_K, |
| hub_D=self.hub_D, tau=self.tau, bank_experts=self.bank_experts, |
| bank_ff=self.bank_ff, head_K=self.head_K, head_D=self.head_D, |
| gate_init=self.gate_init, tie_embeddings=self.tie_embeddings, |
| hub_chunk=self.hub_chunk, hub_const=self.hub_const, |
| hub_ckpt=0) |
|
|
|
|
| class MiniBeatrixForCausalLM(PreTrainedModel, GenerationMixin): |
| config_class = MiniBeatrixConfig |
| main_input_name = "input_ids" |
| supports_gradient_checkpointing = False |
| _tied_weights_keys = {} |
|
|
| def __init__(self, config: MiniBeatrixConfig): |
| super().__init__(config) |
| self.model = AlephLM(config.to_aleph()) |
| self._arm = None |
| self._index = None |
| self.post_init() |
|
|
| def _init_weights(self, module): |
| pass |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): |
| """Loads the bare core (arms off). With default_arm="<id>", passed |
| here or set in config.json, that arm is mounted once the core's |
| weights are in (arms on).""" |
| out = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) |
| model = out[0] if isinstance(out, tuple) else out |
| if getattr(model.config, "default_arm", None): |
| model.mount_arm(model.config.default_arm) |
| return out |
|
|
| |
| def forward(self, input_ids=None, labels=None, attention_mask=None, |
| past_key_values=None, use_cache=None, **kwargs): |
| |
| |
| |
| |
| if input_ids.shape[1] > self.config.context: |
| input_ids = input_ids[:, -self.config.context:] |
| if labels is not None: |
| labels = labels[:, -self.config.context:] |
| out = self.model(input_ids, labels=labels) |
| logits = out[0] |
| loss = out[1] if labels is not None else None |
| return CausalLMOutputWithPast(loss=loss, logits=logits) |
|
|
| def prepare_inputs_for_generation(self, input_ids, **kwargs): |
| return {"input_ids": input_ids[:, -self.config.context:]} |
|
|
| def save_pretrained(self, *args, **kwargs): |
| if self._arm is not None: |
| raise RuntimeError( |
| "an arm is mounted: detach_arm() before saving, or the arm's " |
| "parameters would be written into the core's weight file") |
| return super().save_pretrained(*args, **kwargs) |
|
|
| |
| def _repo_file(self, relpath: str) -> str: |
| """A file of this package: from the local snapshot if the model was |
| loaded from a directory, else fetched from the hub.""" |
| root = str(getattr(self.config, "_name_or_path", "") or "") |
| local = os.path.join(root, *relpath.split("/")) |
| if root and os.path.isfile(local): |
| return local |
| if root and os.path.isdir(root): |
| raise FileNotFoundError( |
| f"{relpath} is missing from {root}: this model was loaded from a " |
| "directory, so the arms are served from that directory too — " |
| "copy arms/ next to config.json, or load by repo id") |
| from huggingface_hub import hf_hub_download |
| return hf_hub_download(root, relpath) |
|
|
| |
| def arms(self) -> list: |
| """The packaged arm index: one row per arm or stack, with what it |
| does, the text it was trained on and what it measured.""" |
| if self._index is None: |
| with open(self._repo_file(ARMS_INDEX), encoding="utf-8") as f: |
| self._index = json.load(f) |
| return list(self._index.get("arms", [])) |
|
|
| def arm_card(self, name: str) -> dict: |
| for row in self.arms(): |
| if row.get("name") == name or row.get("id") == name: |
| return row |
| known = ", ".join(sorted(r.get("id", r.get("name", "?")) for r in self.arms())) |
| raise KeyError(f"no packaged arm {name!r}. Available: {known}") |
|
|
| @property |
| def arm(self): |
| """The live arm handle, or None when the bare core is serving.""" |
| return self._arm |
|
|
| def mount_arm(self, name: str, *, verify_base: bool = True): |
| """Mount one packaged arm (or stack) on the core and return its |
| handle. Mounting a second arm detaches the first.""" |
| card = self.arm_card(name) |
| if verify_base: |
| want = str((self.config.trunk or {}).get("base_model_id") or "") |
| got = str(card.get("base_model_id") or "") |
| if want and got and not (got.startswith(want) or want.startswith(got)): |
| raise ValueError( |
| f"arm {name!r} was trained on {got}, this core is {want} — " |
| "arms are trunk-bound (pass verify_base=False to override)") |
| if self._arm is not None: |
| self.detach_arm() |
| files = [self._repo_file(w) for w in card["weights"]] |
| self._arm = mount(self.model, card, files) |
| return self._arm |
|
|
| def detach_arm(self, *, verify: bool = True): |
| """Restore the bare core. With verify=True the restored logits are |
| asserted BIT-EXACT against the pre-mount fingerprint.""" |
| if self._arm is None: |
| return None |
| h, self._arm = self._arm, None |
| h.detach(verify=verify) |
| return h |
|
|
| |
| def render(self, prompt: str, *, frame: str | None = None, history=None, |
| system: str | None = None) -> torch.Tensor: |
| """Prompt bytes in a frame: 'specials' (the core's own chat frame), |
| 'raw' (the bytes as they are) or 'text'.""" |
| frame = frame or self.default_frame() |
| if frame == "specials": |
| ids = render_specials_frame(prompt, system=system, history=history) |
| elif frame == "text": |
| ids = render_text_frame(prompt, history=history) |
| elif frame == "raw": |
| ids = prompt.encode("utf-8") |
| else: |
| raise ValueError(f"unknown frame {frame!r} (raw | text | specials)") |
| return torch.tensor([list(ids)], dtype=torch.long, device=self.device) |
|
|
| def default_frame(self) -> str: |
| """The live arm's own frame (the stage arms continue raw text); the |
| core's chat frame when nothing is mounted.""" |
| if self._arm is not None: |
| return ((self._arm.card.get("template") or {}).get("frame") |
| or "raw") |
| return "specials" |
|
|
| @torch.no_grad() |
| def say(self, prompt: str, *, frame: str | None = None, history=None, |
| system: str | None = None, max_new: int = 120, |
| temperature: float = 0.0, top_p: float = 1.0, |
| use_cache: bool = True, raw: bool = False) -> str: |
| """Render, continue and cut at the frame's stop (greedy by default). |
| Returns the continuation text.""" |
| frame = frame or self.default_frame() |
| ids = self.render(prompt, frame=frame, history=history, system=system) |
| if getattr(self._arm, "reset_state", None): |
| self._arm.reset_state() |
| out = self._continue(ids, max_new=max_new, temperature=temperature, |
| top_p=top_p, use_cache=use_cache, frame=frame) |
| gen = bytes(int(b) for b in out[0, ids.shape[1]:].tolist()) |
| if raw or frame == "raw": |
| return gen.decode("utf-8", errors="replace") |
| cut = (cut_specials_reply(gen) if frame == "specials" |
| else cut_text_reply(gen)) |
| return cut.decode("utf-8", errors="replace").strip() |
|
|
| @torch.no_grad() |
| def _continue(self, idx: torch.Tensor, *, max_new: int, temperature: float, |
| top_p: float, use_cache: bool, |
| frame: str = "raw") -> torch.Tensor: |
| """The core's own decode: constant-size cached state while the |
| sequence fits the position table. Stops early on the frame's own |
| stop bytes: [END] or a document boundary in every frame, and the |
| text frame's stop strings in that frame only.""" |
| m = self.model |
| m.eval() |
| ctx = int(self.config.context) |
| stops = [bytes([SPECIALS["END"]]), bytes([SPECIALS["DOC"]])] |
| if frame == "text": |
| stops += [s.encode("utf-8") for s in TEXT_STOPS] |
| start = idx.shape[1] |
| if not (use_cache and idx.shape[1] < ctx): |
| for _ in range(max_new): |
| logits = self(input_ids=idx[:, -ctx:]).logits |
| nxt = m._sample(logits, temperature, top_p) |
| idx = torch.cat([idx, nxt], dim=1) |
| if _hit_stop(idx, start, stops): |
| break |
| return idx |
| logits, cache = m.prefill(idx) |
| for i in range(min(max_new, ctx - start)): |
| nxt = m._sample(logits, temperature, top_p) |
| idx = torch.cat([idx, nxt], dim=1) |
| if _hit_stop(idx, start, stops): |
| break |
| logits = m.decode_step(nxt, cache) |
| return idx |
|
|
|
|
| def _hit_stop(idx: torch.Tensor, start: int, stops) -> bool: |
| tail = bytes(int(b) for b in idx[0, start:].tolist())[-8:] |
| return any(tail.endswith(s) for s in stops) |
|
|
