mini-beatrix-2s automodel: modeling_minibeatrix.py (mission final 16.101B, alephllm 0.8.6)

modeling_minibeatrix.py

@@ -0,0 +1,99 @@
+"""HF AutoModel wrapper for mini-beatrix-2s — the FULL-SPLAT governed craft
+(a CausalSplatHUB constellation in every block; 16.101B-token mission,
+completed 2026-08-31).
+
+Loads with:
+    AutoModelForCausalLM.from_pretrained("AbstractPhil/mini-beatrix-2s",
+                                         trust_remote_code=True)
+The model reads raw UTF-8 bytes: input_ids are byte values 0..255.
+"""
+from __future__ import annotations
+
+import torch
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .presets import AlephLMConfig
+from .alephlm import AlephLM
+
+
+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-2s", d_model=1024, n_layers=20,
+                 n_heads=16, context=4096, vocab_size=256,
+                 tokenizer="byte-trigram", hub_layers=tuple(range(20)),
+                 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, **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        # inference: always 0
+        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.post_init()
+
+    def _init_weights(self, module):
+        pass  # AlephLM initializes itself; loader must not re-init
+
+    def forward(self, input_ids=None, labels=None, attention_mask=None,
+                past_key_values=None, use_cache=None, **kwargs):
+        # attention_mask is safely ignored: causal model, right-padding
+        # with label masking is the trained convention. No KV cache in
+        # this wrapper; generate() recomputes the prefix each step.
+        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:]}