from transformers import PreTrainedModel, PretrainedConfig
from torchvision import models
import torch.nn as nn

class CustomEfficientNetConfig(PretrainedConfig):
    model_type = "custom_efficientnet"
    def __init__(self, num_classes=2, **kwargs):
        super().__init__(**kwargs)
        self.num_classes = num_classes

class CustomEfficientNetForImageClassification(PreTrainedModel):
    config_class = CustomEfficientNetConfig
    base_model_prefix = "efficientnet"

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_classes
        self.efficientnet = models.efficientnet_b0(num_classes=config.num_classes)
        
    def forward(self, pixel_values, labels=None):
        outputs = self.efficientnet(pixel_values)
        loss = None
        if labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(outputs, labels)
        return {"loss": loss, "logits": outputs}

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        config = CustomEfficientNetConfig.from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
        model = cls(config)
        state_dict = torch.load(pretrained_model_name_or_path + "/pytorch_model.bin")
        model.load_state_dict(state_dict)
        return model