from transformers.models.lfm2_moe.configuration_lfm2_moe import Lfm2MoeConfig


class Lfm2MoeCustomConfig(Lfm2MoeConfig):
    model_type = "lfm2_moe_custom"

    def __init__(
        self,
        budget: float = 1.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.budget = budget