import torch
from transformers import PreTrainedModel

from hg.config import MConfig


class MModel(PreTrainedModel):
    config_class = MConfig

    def __init__(self, config: MConfig):
        super().__init__(config)
        self.m = torch.nn.Sequential(
            torch.nn.Linear(1, config.hidden),
            torch.nn.ReLU(),
            torch.nn.Linear(config.hidden, 1),
        )
        self.post_init()

    def forward(self, input: torch.Tensor, label: torch.Tensor):

        return torch.nn.functional.mse_loss(self.m(input), label)