"""Configuration classes for MoCo models compatible with transformers"""

from transformers import PretrainedConfig


class MoCoResNetConfig(PretrainedConfig):
    """Configuration for MoCo ResNet models"""
    
    model_type = "moco_resnet"
    
    def __init__(
        self,
        arch="resnet50",
        block="Bottleneck",
        layers=[3, 4, 6, 3],
        image_size=224,
        num_channels=3,
        num_labels=0,
        epoch=0,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.arch = arch
        self.block = block
        self.layers = layers
        self.image_size = image_size
        self.num_channels = num_channels
        self.num_labels = num_labels
        self.epoch = epoch