| from transformers import PretrainedConfig | |
| class AttnExtConfig(PretrainedConfig): | |
| model_type = "attn_ext" | |
| keys_to_ignore_at_inference = ["past_key_values"] | |
| def __init__( | |
| self, | |
| vocab_size=49152, | |
| d_model=2048, | |
| n_layer=24, | |
| n_head=32, | |
| ffn_multiplier=4.0, | |
| multiple_of=256, | |
| block_size=2048, | |
| rope_theta=10000.0, | |
| dropout=0.0, | |
| rms_norm_eps=1e-5, | |
| initializer_range=0.02, | |
| attention_bias=False, | |
| mlp_bias=False, | |
| input_mode="learned", | |
| binary_dim=16, | |
| binary_encoding="zero_one", | |
| binary_scale=1.0, | |
| code_seed=12345, | |
| min_row_weight=4, | |
| min_col_weight=4, | |
| pad_token_id=None, | |
| bos_token_id=None, | |
| eos_token_id=None, | |
| tie_word_embeddings=False, | |
| use_cache=False, | |
| **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, | |
| **kwargs, | |
| ) | |
| if d_model % n_head != 0: | |
| raise ValueError("d_model must be divisible by n_head") | |
| head_dim = d_model // n_head | |
| if head_dim % 2 != 0: | |
| raise ValueError("RoPE requires an even head dimension") | |
| if input_mode not in {"learned", "binary16", "gf2"}: | |
| raise ValueError( | |
| "input_mode must be learned, binary16, or gf2" | |
| ) | |
| if input_mode != "learned": | |
| if binary_dim != 16: | |
| raise ValueError("Frozen-code models require binary_dim=16") | |
| if vocab_size > 2**binary_dim: | |
| raise ValueError("Vocabulary does not fit in 16 bits") | |
| if d_model % binary_dim != 0: | |
| raise ValueError( | |
| "d_model must be divisible by binary_dim" | |
| ) | |
| if tie_word_embeddings: | |
| raise ValueError( | |
| "Frozen input codes cannot be tied to lm_head" | |
| ) | |
| if binary_encoding not in {"zero_one", "bipolar"}: | |
| raise ValueError( | |
| "binary_encoding must be zero_one or bipolar" | |
| ) | |
| self.vocab_size = vocab_size | |
| self.d_model = d_model | |
| self.hidden_size = d_model | |
| self.n_layer = n_layer | |
| self.num_hidden_layers = n_layer | |
| self.n_head = n_head | |
| self.num_attention_heads = n_head | |
| self.head_dim = head_dim | |
| self.ffn_multiplier = ffn_multiplier | |
| self.multiple_of = multiple_of | |
| self.block_size = block_size | |
| self.max_position_embeddings = block_size | |
| self.rope_theta = rope_theta | |
| self.dropout = dropout | |
| self.rms_norm_eps = rms_norm_eps | |
| self.initializer_range = initializer_range | |
| self.attention_bias = attention_bias | |
| self.mlp_bias = mlp_bias | |
| self.input_mode = input_mode | |
| self.binary_dim = binary_dim | |
| self.binary_encoding = binary_encoding | |
| self.binary_scale = binary_scale | |
| self.binary_repeat = d_model // binary_dim | |
| self.code_seed = code_seed | |
| self.min_row_weight = min_row_weight | |
| self.min_col_weight = min_col_weight | |
| self.use_cache = use_cache | |
| self.is_decoder = True | |
| self.is_encoder_decoder = False | |