Upload OPTForMusicGeneration

config.json

@@ -7,6 +7,9 @@
     "OPTForMusicGeneration"
   ],
   "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "model.OPTForMusicGeneration"
+  },
   "bos_token_id": 1,
   "do_layer_norm_before": true,
   "do_sample": true,

model.py

@@ -0,0 +1,16 @@
+from transformers import OPTForCausalLM, OPTConfig, AutoModel
+import torch
+
+from miditok import TokSequence
+
+
+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")