# Copyright 2026. Released under the MIT license.
"""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,
        )