File size: 2,614 Bytes
c4609e6 | 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 | """HuggingFace config for RoST.
Ships inside the published model repository and runs on the downloader's
machine, so it must not import anything from `nanochat`.
Field names mirror `nanochat.gpt.GPTConfig` exactly rather than being renamed
to Llama's vocabulary. A rename would need a mapping table that nothing checks,
and a silently wrong mapping produces a model that loads and computes the wrong
thing -- the one failure mode this whole export has to avoid.
"""
from transformers.configuration_utils import PretrainedConfig
class RostConfig(PretrainedConfig):
model_type = "rost"
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(
self,
vocab_size=32768,
n_layer=24,
n_head=12,
n_kv_head=12,
n_embd=1536,
sequence_len=4096,
rope_base=100000,
window_pattern="SSSL",
pad_vocab_size_to=64,
logit_softcap=15.0,
attention_scale=1.2,
ve_gate_channels=12,
smear_gate_channels=24,
bos_token_id=None,
eos_token_id=None,
**kwargs,
):
self.vocab_size = vocab_size
self.n_layer = n_layer
self.n_head = n_head
self.n_kv_head = n_kv_head
self.n_embd = n_embd
self.sequence_len = sequence_len
self.rope_base = rope_base
self.window_pattern = window_pattern
self.pad_vocab_size_to = pad_vocab_size_to
# Constants in the training code, carried as config so a checkpoint
# trained under different ones cannot be served under these.
self.logit_softcap = logit_softcap
self.attention_scale = attention_scale
self.ve_gate_channels = ve_gate_channels
self.smear_gate_channels = smear_gate_channels
super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
@property
def padded_vocab_size(self):
pad = self.pad_vocab_size_to
return ((self.vocab_size + pad - 1) // pad) * pad
@property
def head_dim(self):
return self.n_embd // self.n_head
# Aliases so generic HuggingFace code (generation, device maps, pipelines)
# finds what it expects without the weights being renamed.
@property
def hidden_size(self):
return self.n_embd
@property
def num_attention_heads(self):
return self.n_head
@property
def num_key_value_heads(self):
return self.n_kv_head
@property
def num_hidden_layers(self):
return self.n_layer
@property
def max_position_embeddings(self):
return self.sequence_len
|