| from transformers import OPTForCausalLM, OPTConfig, AutoModel | |
| import torch | |
| from miditok import TokSequence | |
| # class OPTForMusicGenerationConfig(OPTConfig): | |
| class OPTForMusicGeneration(OPTForCausalLM): | |
| def generate_music(self, **kwargs): | |
| input = torch.tensor([[self.config.bos_token_id]], device=self.device) | |
| midi = self.generate(input, **kwargs) | |
| generated_ts = TokSequence(ids=midi.tolist()[0], ids_bpe_encoded=True) | |
| return generated_ts | |
| OPTForMusicGeneration.register_for_auto_class("AutoModel") | |
| OPTForMusicGeneration.register_for_auto_class("AutoModelForCausalLM") | |