File size: 3,844 Bytes
b007aec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | """HF AutoModel wrapper for mini-beatrix-1 (AlephLM byte-trigram craft).
Loads with:
AutoModelForCausalLM.from_pretrained("AbstractPhil/mini-beatrix-1",
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-1", d_model=768, n_layers=16,
n_heads=12, context=2048, vocab_size=256,
tokenizer="byte-trigram", hub_layers=(4, 9, 14),
hub_K=512, hub_D=32, tau=0.1, bank_experts=3, bank_ff=None,
head_K=512, head_D=32, gate_init=-3.0,
tie_embeddings=False, hub_chunk=128, **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
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)
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:]}
|