| |
| """Configuration for the 3-digit-basic-calc looped arithmetic transformer.""" |
|
|
| from transformers import PretrainedConfig |
|
|
|
|
| class ThreeDigitBasicCalcConfig(PretrainedConfig): |
| """Config for a looped ("universal") decoder-only calculator transformer. |
| |
| A small stack of ``num_hidden_layers`` decoder blocks is applied ``n_loops`` |
| times with shared weights, giving an effective depth of |
| ``num_hidden_layers * n_loops`` at the parameter cost of the stack alone. |
| """ |
|
|
| model_type = "3-digit-basic-calc" |
|
|
| def __init__( |
| self, |
| vocab_size: int = 34, |
| max_position_embeddings: int = 176, |
| hidden_size: int = 256, |
| num_attention_heads: int = 8, |
| num_hidden_layers: int = 2, |
| n_loops: int = 8, |
| dropout: float = 0.05, |
| rms_eps: float = 1e-6, |
| rope_base: float = 10000.0, |
| pad_token_id: int = 0, |
| bos_token_id: int = 1, |
| eos_token_id: int = 2, |
| unk_token_id: int = 3, |
| **kwargs, |
| ): |
| self.vocab_size = vocab_size |
| self.max_position_embeddings = max_position_embeddings |
| self.hidden_size = hidden_size |
| self.num_attention_heads = num_attention_heads |
| self.num_hidden_layers = num_hidden_layers |
| self.n_loops = n_loops |
| self.dropout = dropout |
| self.rms_eps = rms_eps |
| self.rope_base = rope_base |
| self.unk_token_id = unk_token_id |
| kwargs.setdefault("tie_word_embeddings", False) |
| super().__init__( |
| pad_token_id=pad_token_id, |
| bos_token_id=bos_token_id, |
| eos_token_id=eos_token_id, |
| **kwargs, |
| ) |
|
|
