from transformers import PretrainedConfig
from transformers.models.gemma3.configuration_gemma3 import Gemma3TextConfig

class Gemma3PXConfig(Gemma3TextConfig):
    model_type = "gemma_3_px"
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # Default config for 270m
        self.px_routing_mode = kwargs.get("px_routing_mode", "adaptive")
        self.px_gamma = kwargs.get("px_gamma", 0.10)