from transformers import PretrainedConfig


class SolProHFConfig(PretrainedConfig):
    model_type = "sol_pro"

    def __init__(self, vocab_size=4096, width=768, heads=24, kv_heads=8,
                 stored_blocks=12, ffn_width=3968, recurrent_start=1,
                 recurrent_blocks=4, recurrent_passes=2, tn_gram_rank=299,
                 max_ngram_order=4, use_tn_gram=True, use_qk_norm=True,
                 use_loop_conditioning=True, max_position_embeddings=2048,
                 rope_theta=20000.0, **kwargs):
        kwargs.pop("tie_word_embeddings", None)
        kwargs.pop("is_decoder", None)
        super().__init__(tie_word_embeddings=True, is_decoder=True, **kwargs)
        settings = locals().copy()
        for name, value in settings.items():
            if name not in {"self", "kwargs", "__class__"}:
                setattr(self, name, value)
        self.hidden_size = width
        self.num_attention_heads = heads
        self.num_key_value_heads = kv_heads
        self.num_hidden_layers = stored_blocks
        self.use_cache = False