from dataclasses import dataclass

from typing import Tuple

import torch

from transformers import AutoModel, AutoConfig

from transformers import PreTrainedModel, PretrainedConfig

from transformers.utils import ModelOutput

@dataclass

class DummyModelOutput(ModelOutput):

    hidden_states: Tuple[torch.Tensor]

class DummyModelConfig(PretrainedConfig):

    model_type = "dummy_model"

    def __init__(self, hidden_size: int = 768, **kwargs) -> None:

        self.hidden_size = hidden_size

        super().__init__(**kwargs)

        return

class DummyModel(PreTrainedModel):

    config_class = DummyModelConfig

    def __init__(self, config: DummyModelConfig) -> None:

        super().__init__(config)

        self.param = torch.nn.Parameter(torch.zeros(1), requires_grad=False)

        return

    def forward(self, x: torch.Tensor, **kwargs) -> DummyModelOutput:

        assert x.shape[-1] == self.config.hidden_size

        dtype = self.param.dtype

        x = x.to(dtype)

        return DummyModelOutput(hidden_states=(x,))

AutoConfig.register("dummy_model", DummyModelConfig)

AutoModel.register(DummyModelConfig, DummyModel)

if __name__ == "__main__":

    config = DummyModelConfig(hidden_size=3)

    print(config)

    model = AutoModel.from_config(config)

    print(model)

    x = torch.randn(1, 2, 3)

    output = model(x)

    print(output)