| from transformers import PretrainedConfig | |
| class HaloSConfig(PretrainedConfig): | |
| model_type = "halo_s" | |
| def __init__( | |
| self, | |
| vocab_size=50257, | |
| hidden_size=512, | |
| num_layers=6, | |
| num_heads=8, | |
| num_kv_heads=2, | |
| num_globals=2, | |
| local_window=64, | |
| dilated_offsets=None, | |
| num_random=2, | |
| dropout=0.1, | |
| max_seq_len=4096, | |
| **kwargs | |
| ): | |
| super().__init__(**kwargs) | |
| if dilated_offsets is None: | |
| dilated_offsets = [1, 2, 4, 8, 16, 32, 64, 128] | |
| self.vocab_size = vocab_size | |
| self.hidden_size = hidden_size | |
| self.num_layers = num_layers | |
| self.num_heads = num_heads | |
| self.num_kv_heads = num_kv_heads | |
| self.num_globals = num_globals | |
| self.local_window = local_window | |
| self.dilated_offsets = dilated_offsets | |
| self.num_random = num_random | |
| self.dropout = dropout | |
| self.max_seq_len = max_seq_len | |
| def head_dim(self): | |
| return self.hidden_size // self.num_heads |