from transformers import PretrainedConfig

class OmegaConfig(PretrainedConfig):
    """
    Fixed Configuration class for Omega-3060 quantized models.
    Removes property conflicts for seamless loading.
    """
    model_type = "omega_3060"
    
    def __init__(
        self,
        vocab_size=163840,
        hidden_size=7168,
        intermediate_size=2048,   
        expert_intermediate=2048, 
        num_hidden_layers=61,
        num_attention_heads=64,
        num_key_value_heads=64,
        num_experts_per_layer=10,  
        kv_lora_rank=512,
        q_lora_rank=1536,
        qk_rope_head_dim=64,
        v_head_dim=128,
        qk_nope_head_dim=128,
        hidden_act="silu",
        max_position_embeddings=262144,
        rms_norm_eps=1e-5,
        rope_theta=50000.0,
        tau=0.1,             
        expert_floor=0.05,    
        use_cache=False,      
        bits=4,               
        group_size=128,       
        quant_method="omega_custom",
        bos_token_id=163584,
        eos_token_id=163586,
        pad_token_id=163839,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.expert_intermediate = expert_intermediate
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        
        # Handle num_key_value_heads correctly as an attribute
        if num_key_value_heads is None:
            self.num_key_value_heads = num_attention_heads
        else:
            self.num_key_value_heads = num_key_value_heads
            
        self.num_experts_per_layer = num_experts_per_layer
        self.kv_lora_rank = kv_lora_rank
        self.q_lora_rank = q_lora_rank
        self.qk_rope_head_dim = qk_rope_head_dim
        self.v_head_dim = v_head_dim
        self.qk_nope_head_dim = qk_nope_head_dim
        self.hidden_act = hidden_act
        self.max_position_embeddings = max_position_embeddings
        self.rms_norm_eps = rms_norm_eps
        self.rope_theta = rope_theta
        self.tau = tau
        self.expert_floor = expert_floor
        self.use_cache = use_cache
        
        self.bits = bits
        self.group_size = group_size
        self.quant_method = quant_method
            
        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            **kwargs,
        )