File size: 2,166 Bytes
da49047 | 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 | from transformers import PretrainedConfig
class BETConfig(PretrainedConfig):
model_type = "bet"
def __init__(
self,
vocab_size=259,
hidden_size=324,
intermediate_size=864,
prelude_layers=1,
body_blocks=6,
coda_layers=1,
num_attention_heads=6,
num_key_value_heads=2,
head_dim=54,
lora_rank=16,
hyper_lanes=2,
max_position_embeddings=1024,
max_loops=8,
rope_theta=10_000.0,
rms_norm_eps=1e-6,
ddl_beta_init=1.0,
ddl_k_eps=1e-2,
ddl_v_sigmoid_scale=4.0,
refinement_cycles=8,
use_cache=False,
tie_word_embeddings=True,
pad_token_id=256,
bos_token_id=257,
eos_token_id=258,
**kwargs,
):
super().__init__(
pad_token_id=pad_token_id,
bos_token_id=bos_token_id,
eos_token_id=eos_token_id,
tie_word_embeddings=tie_word_embeddings,
is_encoder_decoder=False,
**kwargs,
)
self.vocab_size=int(vocab_size)
self.hidden_size=int(hidden_size)
self.intermediate_size=int(intermediate_size)
self.prelude_layers=int(prelude_layers)
self.body_blocks=int(body_blocks)
self.coda_layers=int(coda_layers)
# Common HF tooling expects num_hidden_layers even though only the body loops.
self.num_hidden_layers=int(prelude_layers+body_blocks+coda_layers)
self.num_attention_heads=int(num_attention_heads)
self.num_key_value_heads=int(num_key_value_heads)
self.head_dim=int(head_dim)
self.lora_rank=int(lora_rank)
self.hyper_lanes=int(hyper_lanes)
self.max_position_embeddings=int(max_position_embeddings)
self.max_loops=int(max_loops)
self.rope_theta=float(rope_theta)
self.rms_norm_eps=float(rms_norm_eps)
self.ddl_beta_init=float(ddl_beta_init)
self.ddl_k_eps=float(ddl_k_eps)
self.ddl_v_sigmoid_scale=float(ddl_v_sigmoid_scale)
self.refinement_cycles=int(refinement_cycles)
self.use_cache=bool(use_cache)
|