from transformers import PretrainedConfig


class EdgeFaceConfig(PretrainedConfig):
    """
    Configuration for EdgeFace face-recognition models.

    EdgeFace is a `timm` edgenext backbone with the classifier reset to output a
    `featdim`-dimensional embedding. Some variants additionally replace their
    nn.Linear layers with a static low-rank factorization (two smaller linears)
    to cut parameters -- this is EdgeFace's "gamma" trick and is baked into the
    weights. It is NOT PEFT/LoRA adapters; you can still train real LoRA on top
    of the resulting model.

    The four published variants map to:
        edgeface_base        -> timm_model="edgenext_base",    use_low_rank=False
        edgeface_s_gamma_05  -> timm_model="edgenext_small",   use_low_rank=True, low_rank_ratio=0.5
        edgeface_xs_gamma_06 -> timm_model="edgenext_x_small", use_low_rank=True, low_rank_ratio=0.6
        edgeface_xxs         -> timm_model="edgenext_xx_small", use_low_rank=False
    """

    model_type = "edgeface"

    def __init__(
        self,
        timm_model: str = "edgenext_x_small",
        featdim: int = 512,
        use_low_rank: bool = False,
        low_rank_ratio: float = 0.6,
        **kwargs,
    ):
        self.timm_model = timm_model
        self.featdim = featdim
        self.use_low_rank = use_low_rank
        self.low_rank_ratio = low_rank_ratio
        super().__init__(**kwargs)