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()