midi_model.py

import json 
from typing import Union, Dict, Any 
 
import numpy as np 
import torch 
import torch.nn as nn 
import torch.nn.functional as F 
import tqdm 
from peft import PeftConfig, LoraModel, load_peft_weights, set_peft_model_state_dict 
from transformers import LlamaModel, LlamaConfig, DynamicCache, PretrainedConfig, PreTrainedModel 
 
from midi_tokenizer import MIDITokenizerV1, MIDITokenizerV2, MIDITokenizer 
 
config_name_list = ["tv1-medium", "tv2-medium", "tv2o-medium", "tv2-large", "tv2o-large"] 
 
 
class MIDIModelConfig(PretrainedConfig): 
 model_type = "midi_model" 
 
 def __init__(self, 
 tokenizer: Union[MIDITokenizerV1, MIDITokenizerV2, Dict]=None, 
 net_config: Union[LlamaConfig, Dict]=None, 
 net_token_config: Union[LlamaConfig, Dict]=None, 
 **kwargs): 
 super().__init__(**kwargs) 
 if tokenizer: 
 if isinstance(tokenizer, dict): 
 self.tokenizer = MIDITokenizer(tokenizer["version"]) 
 self.tokenizer.set_optimise_midi(tokenizer["optimise_midi"]) 
 else: 
 self.tokenizer = tokenizer 
 else: 
 self.tokenizer = MIDITokenizer() 
 if net_config: 
 if isinstance(net_config, dict): 
 self.net_config = LlamaConfig(**net_config) 
 else: 
 self.net_config = net_config 
 else: 
 self.net_config = LlamaConfig() 
 if net_token_config: 
 if isinstance(net_token_config, dict): 
 self.net_token_config = LlamaConfig(**net_token_config) 
 else: 
 self.net_token_config = net_token_config 
 else: 
 self.net_token_config = LlamaConfig() 
 self.n_embd = self.net_token_config.hidden_size 
 
 def to_dict(self) -> Dict[str, Any]: 
 d = super().to_dict() 
 d["tokenizer"] = self.tokenizer.to_dict() 
 return d 
 
 def __str__(self): 
 d = { 
 "net": self.net_config.to_json_string(use_diff=False), 
 "net_token": self.net_token_config.to_json_string(use_diff=False) 
 } 
 return json.dumps(d, indent=4) 
 
 @staticmethod 
 def get_config(tokenizer_ver="v2", optimise_midi=True, n_layer=12, n_head=16, n_embd=1024, n_inner=4096): 
 tokenizer = MIDITokenizer(tokenizer_ver) 
 tokenizer.set_optimise_midi(optimise_midi) 
 net_config = LlamaConfig(vocab_size=tokenizer.vocab_size, 
 hidden_size=n_embd, num_attention_heads=n_head, 
 num_hidden_layers=n_layer, intermediate_size=n_inner, 
 pad_token_id=tokenizer.pad_id, max_position_embeddings=4096, 
 use_cache=False) 
 net_token_config = LlamaConfig(vocab_size=tokenizer.vocab_size, 
 hidden_size=n_embd, num_attention_heads=n_head // 4, 
 num_hidden_layers=n_layer // 4, intermediate_size=n_inner // 4, 
 pad_token_id=tokenizer.pad_id, max_position_embeddings=4096, 
 use_cache=False) 
 return MIDIModelConfig(tokenizer, net_config, net_token_config) 
 
 @staticmethod 
 def from_name(name="tv2o-medium"): 
 tv, size = name.split("-") 
 tv = tv[1:] 
 if tv[-1] == "o": 
 o = True 
 tv = tv[:-1] 
 else: 
 o = False 
 if tv not in ["v1", "v2"]: 
 raise ValueError(f"Unknown tokenizer version {tv}") 
 if size == "medium": 
 return MIDIModelConfig.get_config(tokenizer_ver=tv, optimise_midi=o, 
 n_layer=12, n_head=16, n_embd=1024, n_inner=4096) 
 elif size == "large": 
 return MIDIModelConfig.get_config(tokenizer_ver=tv, optimise_midi=o, 
 n_layer=24, n_head=16, n_embd=1024, n_inner=4096) 
 else: 
 raise ValueError(f"Unknown model size {size}") 
 
 
class MIDIModel(PreTrainedModel): 
 config_class = MIDIModelConfig 
 
 def __init__(self, config: MIDIModelConfig, *args, **kwargs): 
 super(MIDIModel, self).__init__(config, *args, **kwargs) 
 self.tokenizer = config.tokenizer 
 self.net = LlamaModel(config.net_config) 
 self.net_token = LlamaModel(config.net_token_config) 
 self.lm_head = nn.Linear(config.n_embd, self.tokenizer.vocab_size, bias=False) 
 
 def load_merge_lora(self, model_id): 
 peft_config = PeftConfig.from_pretrained(model_id) 
 model = LoraModel(self, peft_config, adapter_name="default") 
 adapter_state_dict = load_peft_weights(model_id, device=str(self.device)) 
 set_peft_model_state_dict(self, adapter_state_dict, "default") 
 return model.merge_and_unload() 
 
 def forward_token(self, hidden_state=None, x=None, cache=None): 
 """ 
 
 :param hidden_state: (batch_size, n_embd) 
 :param x: (batch_size, token_sequence_length) 
 :param cache: Cache 
 :return: (batch_size, 1 + token_sequence_length, vocab_size) 
 """ 
 if hidden_state is not None: 
 #if you use cache, you don't need to pass in hidden_state 
 hidden_state = hidden_state.unsqueeze(1) # (batch_size, 1, n_embd) 
 if x is not None: 
 x = self.net_token.embed_tokens(x) 
 if hidden_state is not None: 
 x = torch.cat([hidden_state, x], dim=1) 
 hidden_state = x 
 hidden_state = self.net_token.forward(inputs_embeds=hidden_state, 
 past_key_values=cache, 
 use_cache=cache is not None).last_hidden_state 
 return self.lm_head(hidden_state) 
 
 def forward(self, x, cache = None): 
 """ 
 :param x: (batch_size, midi_sequence_length, token_sequence_length) 
 :param cache: Cache 
 :return: hidden (batch_size, midi_sequence_length, n_embd) 
 """ 
 
 # merge token sequence 
 x = self.net.embed_tokens(x) 
 x = x.sum(dim=-2) 
 x = self.net.forward(inputs_embeds=x, 
 past_key_values=cache, 
 use_cache=cache is not None) 
 return x.last_hidden_state 
 
 def sample_top_p_k(self, probs, p, k, generator=None): 
 probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True) 
 probs_sum = torch.cumsum(probs_sort, dim=-1) 
 mask = probs_sum - probs_sort > p 
 probs_sort[mask] = 0.0 
 mask = torch.zeros(probs_sort.shape[-1], device=probs_sort.device) 
 mask[:k] = 1 
 probs_sort = probs_sort * mask 
 probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True)) 
 shape = probs_sort.shape 
 next_token = torch.multinomial(probs_sort.reshape(-1, shape[-1]), 
 num_samples=1, generator=generator).reshape(*shape[:-1], 1) 
 next_token = torch.gather(probs_idx, -1, next_token).reshape(*shape[:-1]) 
 return next_token 
 
 @torch.inference_mode() 
 def generate(self, prompt=None, batch_size=1, max_len=512, temp=1.0, top_p=0.98, top_k=20, generator=None): 
 tokenizer = self.tokenizer 
 max_token_seq = tokenizer.max_token_seq 
 if prompt is None: 
 input_tensor = torch.full((1, max_token_seq), tokenizer.pad_id, dtype=torch.long, device=self.device) 
 input_tensor[0, 0] = tokenizer.bos_id # bos 
 input_tensor = input_tensor.unsqueeze(0) 
 input_tensor = torch.cat([input_tensor] * batch_size, dim=0) 
 else: 
 if len(prompt.shape) == 2: 
 prompt = prompt[None, :] 
 prompt = np.repeat(prompt, repeats=batch_size, axis=0) 
 elif prompt.shape[0] == 1: 
 prompt = np.repeat(prompt, repeats=batch_size, axis=0) 
 elif len(prompt.shape) != 3 or prompt.shape[0] != batch_size: 
 raise ValueError(f"invalid shape for prompt, {prompt.shape}") 
 prompt = prompt[..., :max_token_seq] 
 if prompt.shape[-1] < max_token_seq: 
 prompt = np.pad(prompt, ((0, 0), (0, 0), (0, max_token_seq - prompt.shape[-1])), 
 mode="constant", constant_values=tokenizer.pad_id) 
 input_tensor = torch.from_numpy(prompt).to(dtype=torch.long, device=self.device) 
 
 cur_len = input_tensor.shape[1] 
 bar = tqdm.tqdm(desc="generating", total=max_len - cur_len) 
 cache1 = DynamicCache() 
 past_len = 0 
 with bar: 
 while cur_len < max_len: 
 end = [False] * batch_size 
 hidden = self.forward(input_tensor[:, past_len:], cache=cache1)[:, -1] 
 next_token_seq = None 
 event_names = [""] * batch_size 
 cache2 = DynamicCache() 
 for i in range(max_token_seq): 
 mask = torch.zeros((batch_size, tokenizer.vocab_size), dtype=torch.int64, device=self.device) 
 for b in range(batch_size): 
 if end[b]: 
 mask[b, tokenizer.pad_id] = 1 
 continue 
 if i == 0: 
 mask[b, list(tokenizer.event_ids.values()) + [tokenizer.eos_id]] = 1 
 else: 
 param_names = tokenizer.events[event_names[b]] 
 if i > len(param_names): 
 mask[b, tokenizer.pad_id] = 1 
 continue 
 mask[b, tokenizer.parameter_ids[param_names[i - 1]]] = 1 
 mask = mask.unsqueeze(1) 
 x = next_token_seq 
 if i != 0: 
 # cached 
 hidden = None 
 x = x[:, -1:] 
 logits = self.forward_token(hidden, x, cache=cache2)[:, -1:] 
 scores = torch.softmax(logits / temp, dim=-1) * mask 
 samples = self.sample_top_p_k(scores, top_p, top_k, generator=generator) 
 if i == 0: 
 next_token_seq = samples 
 for b in range(batch_size): 
 if end[b]: 
 continue 
 eid = samples[b].item() 
 if eid == tokenizer.eos_id: 
 end[b] = True 
 else: 
 event_names[b] = tokenizer.id_events[eid] 
 else: 
 next_token_seq = torch.cat([next_token_seq, samples], dim=1) 
 if all([len(tokenizer.events[event_names[b]]) == i for b in range(batch_size) if not end[b]]): 
 break 
 
 if next_token_seq.shape[1] < max_token_seq: 
 next_token_seq = F.pad(next_token_seq, (0, max_token_seq - next_token_seq.shape[1]), 
 "constant", value=tokenizer.pad_id) 
 next_token_seq = next_token_seq.unsqueeze(1) 
 input_tensor = torch.cat([input_tensor, next_token_seq], dim=1) 
 past_len = cur_len 
 cur_len += 1 
 bar.update(1) 
 
 if all(end): 
 break 
 return input_tensor.cpu().numpy()