| 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, *args, **kwargs): | |
| return self.model(*args, **kwargs) | |
| def generate(self, *args, **kwargs): | |
| return self.model.generate(*args, **kwargs) | |