from transformers import PreTrainedModel
from .mgpt_config import MGPTConfig
from .gpt_model import GPT

class MusicModel(PreTrainedModel):
    config_class = MGPTConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = GPT(config)

    def forward(self, inputs):
        return self.model(inputs)