from transformers import PretrainedConfig

class Sol2Config(PretrainedConfig):
    model_type = "sol2"
    def __init__(self, vocab_size=4096, width=256, heads=8, kv_heads=4,
                 blocks=10, ffn_width=1552, recurrent_start=1, recurrent_blocks=4,
                 passes=2, loop_conditioning=True, rope_theta=20000.0,
                 max_context=2048, backend="sdpa", **kwargs):
        kwargs.pop("tie_word_embeddings", None)
        kwargs.pop("is_decoder", None)
        super().__init__(tie_word_embeddings=True, is_decoder=True, **kwargs)
        for key,value in locals().copy().items():
            if key not in ("self", "kwargs", "__class__"): setattr(self,key,value)
        self.hidden_size=width
        self.num_attention_heads=heads
        self.num_key_value_heads=kv_heads
        self.num_hidden_layers=blocks
        self.max_position_embeddings=max_context
        self.use_cache=False