2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license file for loading the pop2piano model weights from the official repository and to show how tokenizer vocab was constructed import json import torch from transformers import pop2pianoconfig pop2pianoforconditionalgeneration model weights this weights were downloaded from the official pop2piano repository https huggingface cosweetcocoapop2pianoblobmainmodel1999val0 67311615 ckpt officialweights torch load model1999val0 67311615 ckpt statedict load the config and init the model cfg pop2pianoconfig frompretrainedsweetcocoapop2piano model pop2pianoforconditionalgenerationcfg load relative attention bias statedictencoder block 0 layer 0 selfattention relativeattentionbias weight officialweightsstatedict transformer encoder block 0 layer 0 selfattention relativeattentionbias weight statedictdecoder block 0 layer 0 selfattention relativeattentionbias weight officialweightsstatedict transformer decoder block 0 layer 0 selfattention relativeattentionbias weight load embed tokens and final layer norm for both encoder and decoder statedictencoder embedtokens weight officialweightsstatedicttransformer encoder embedtokens weight statedictdecoder embedtokens weight officialweightsstatedicttransformer decoder embedtokens weight statedictencoder finallayernorm weight officialweightsstatedict transformer encoder finallayernorm weight statedictdecoder finallayernorm weight officialweightsstatedict transformer decoder finallayernorm weight load lmhead melconditioner emb and shared statedictlmhead weight officialweightsstatedicttransformer lmhead weight statedictmelconditioner embedding weight officialweightsstatedictmelconditioner embedding weight statedictshared weight officialweightsstatedicttransformer shared weight load each encoder blocks for i in rangecfg numlayers layer 0 statedictfencoder block i layer 0 selfattention q weight officialweightsstatedict ftransformer encoder block i layer 0 selfattention q weight statedictfencoder block i layer 0 selfattention k weight officialweightsstatedict ftransformer encoder block i layer 0 selfattention k weight statedictfencoder block i layer 0 selfattention v weight officialweightsstatedict ftransformer encoder block i layer 0 selfattention v weight statedictfencoder block i layer 0 selfattention o weight officialweightsstatedict ftransformer encoder block i layer 0 selfattention o weight statedictfencoder block i layer 0 layernorm weight officialweightsstatedict ftransformer encoder block i layer 0 layernorm weight layer 1 statedictfencoder block i layer 1 densereludense wi0 weight officialweightsstatedict ftransformer encoder block i layer 1 densereludense wi0 weight statedictfencoder block i layer 1 densereludense wi1 weight officialweightsstatedict ftransformer encoder block i layer 1 densereludense wi1 weight statedictfencoder block i layer 1 densereludense wo weight officialweightsstatedict ftransformer encoder block i layer 1 densereludense wo weight statedictfencoder block i layer 1 layernorm weight officialweightsstatedict ftransformer encoder block i layer 1 layernorm weight load each decoder blocks for i in range6 layer 0 statedictfdecoder block i layer 0 selfattention q weight officialweightsstatedict ftransformer decoder block i layer 0 selfattention q weight statedictfdecoder block i layer 0 selfattention k weight officialweightsstatedict ftransformer decoder block i layer 0 selfattention k weight statedictfdecoder block i layer 0 selfattention v weight officialweightsstatedict ftransformer decoder block i layer 0 selfattention v weight statedictfdecoder block i layer 0 selfattention o weight officialweightsstatedict ftransformer decoder block i layer 0 selfattention o weight statedictfdecoder block i layer 0 layernorm weight officialweightsstatedict ftransformer decoder block i layer 0 layernorm weight layer 1 statedictfdecoder block i layer 1 encdecattention q weight officialweightsstatedict ftransformer decoder block i layer 1 encdecattention q weight statedictfdecoder block i layer 1 encdecattention k weight officialweightsstatedict ftransformer decoder block i layer 1 encdecattention k weight statedictfdecoder block i layer 1 encdecattention v weight officialweightsstatedict ftransformer decoder block i layer 1 encdecattention v weight statedictfdecoder block i layer 1 encdecattention o weight officialweightsstatedict ftransformer decoder block i layer 1 encdecattention o weight statedictfdecoder block i layer 1 layernorm weight officialweightsstatedict ftransformer decoder block i layer 1 layernorm weight layer 2 statedictfdecoder block i layer 2 densereludense wi0 weight officialweightsstatedict ftransformer decoder block i layer 2 densereludense wi0 weight statedictfdecoder block i layer 2 densereludense wi1 weight officialweightsstatedict ftransformer decoder block i layer 2 densereludense wi1 weight statedictfdecoder block i layer 2 densereludense wo weight officialweightsstatedict ftransformer decoder block i layer 2 densereludense wo weight statedictfdecoder block i layer 2 layernorm weight officialweightsstatedict ftransformer decoder block i layer 2 layernorm weight model loadstatedictstatedict stricttrue save the weights torch savestatedict pytorchmodel bin tokenizer the tokenize and detokenize methods are taken from the official implementation link https github comsweetcocoapop2pianoblobfac11e8dcfc73487513f4588e8d0c22a22f2fdc5miditokenizer pyl34 def tokenizeidx tokentype nspecial4 nnote128 nvelocity2 if tokentype tokentime return nspecial nnote nvelocity idx elif tokentype tokenvelocity return nspecial nnote idx elif tokentype tokennote return nspecial idx elif tokentype tokenspecial return idx else return 1 link https github comsweetcocoapop2pianoblobfac11e8dcfc73487513f4588e8d0c22a22f2fdc5miditokenizer pyl48 def detokenizeidx nspecial4 nnote128 nvelocity2 timeidxoffset0 if idx nspecial nnote nvelocity return tokentime idx nspecial nnote nvelocity timeidxoffset elif idx nspecial nnote return tokenvelocity idx nspecial nnote elif idx nspecial return tokennote idx nspecial else return tokenspecial idx crate the decoder and then the encoder of the tokenizer decoder for i in rangecfg vocabsize decoder updatei fdetokenizei1detokenizei0 encoder v k for k v in decoder items save the vocab with open vocab json w as file file writejson dumpsencoder 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license file for loading the pop2piano model weights from the official repository and to show how tokenizer vocab was constructed model weights this weights were downloaded from the official pop2piano repository https huggingface co sweetcocoa pop2piano blob main model 1999 val_0 67311615 ckpt load the config and init the model load relative attention bias load embed tokens and final layer norm for both encoder and decoder load lm_head mel_conditioner emb and shared load each encoder blocks layer 0 layer 1 load each decoder blocks layer 0 layer 1 layer 2 save the weights tokenizer the tokenize and detokenize methods are taken from the official implementation link https github com sweetcocoa pop2piano blob fac11e8dcfc73487513f4588e8d0c22a22f2fdc5 midi_tokenizer py l34 link https github com sweetcocoa pop2piano blob fac11e8dcfc73487513f4588e8d0c22a22f2fdc5 midi_tokenizer py l48 crate the decoder and then the encoder of the tokenizer save the vocab
import json

import torch

from transformers import Pop2PianoConfig, Pop2PianoForConditionalGeneration






official_weights = torch.load("./model-1999-val_0.67311615.ckpt")
state_dict = {}



cfg = Pop2PianoConfig.from_pretrained("sweetcocoa/pop2piano")
model = Pop2PianoForConditionalGeneration(cfg)



state_dict["encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"] = official_weights["state_dict"][
    "transformer.encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"
]
state_dict["decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"] = official_weights["state_dict"][
    "transformer.decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"
]


state_dict["encoder.embed_tokens.weight"] = official_weights["state_dict"]["transformer.encoder.embed_tokens.weight"]
state_dict["decoder.embed_tokens.weight"] = official_weights["state_dict"]["transformer.decoder.embed_tokens.weight"]

state_dict["encoder.final_layer_norm.weight"] = official_weights["state_dict"][
    "transformer.encoder.final_layer_norm.weight"
]
state_dict["decoder.final_layer_norm.weight"] = official_weights["state_dict"][
    "transformer.decoder.final_layer_norm.weight"
]


state_dict["lm_head.weight"] = official_weights["state_dict"]["transformer.lm_head.weight"]
state_dict["mel_conditioner.embedding.weight"] = official_weights["state_dict"]["mel_conditioner.embedding.weight"]
state_dict["shared.weight"] = official_weights["state_dict"]["transformer.shared.weight"]


for i in range(cfg.num_layers):
    
    state_dict[f"encoder.block.{i}.layer.0.SelfAttention.q.weight"] = official_weights["state_dict"][
        f"transformer.encoder.block.{i}.layer.0.SelfAttention.q.weight"
    ]
    state_dict[f"encoder.block.{i}.layer.0.SelfAttention.k.weight"] = official_weights["state_dict"][
        f"transformer.encoder.block.{i}.layer.0.SelfAttention.k.weight"
    ]
    state_dict[f"encoder.block.{i}.layer.0.SelfAttention.v.weight"] = official_weights["state_dict"][
        f"transformer.encoder.block.{i}.layer.0.SelfAttention.v.weight"
    ]
    state_dict[f"encoder.block.{i}.layer.0.SelfAttention.o.weight"] = official_weights["state_dict"][
        f"transformer.encoder.block.{i}.layer.0.SelfAttention.o.weight"
    ]
    state_dict[f"encoder.block.{i}.layer.0.layer_norm.weight"] = official_weights["state_dict"][
        f"transformer.encoder.block.{i}.layer.0.layer_norm.weight"
    ]

    
    state_dict[f"encoder.block.{i}.layer.1.DenseReluDense.wi_0.weight"] = official_weights["state_dict"][
        f"transformer.encoder.block.{i}.layer.1.DenseReluDense.wi_0.weight"
    ]
    state_dict[f"encoder.block.{i}.layer.1.DenseReluDense.wi_1.weight"] = official_weights["state_dict"][
        f"transformer.encoder.block.{i}.layer.1.DenseReluDense.wi_1.weight"
    ]
    state_dict[f"encoder.block.{i}.layer.1.DenseReluDense.wo.weight"] = official_weights["state_dict"][
        f"transformer.encoder.block.{i}.layer.1.DenseReluDense.wo.weight"
    ]
    state_dict[f"encoder.block.{i}.layer.1.layer_norm.weight"] = official_weights["state_dict"][
        f"transformer.encoder.block.{i}.layer.1.layer_norm.weight"
    ]


for i in range(6):
    
    state_dict[f"decoder.block.{i}.layer.0.SelfAttention.q.weight"] = official_weights["state_dict"][
        f"transformer.decoder.block.{i}.layer.0.SelfAttention.q.weight"
    ]
    state_dict[f"decoder.block.{i}.layer.0.SelfAttention.k.weight"] = official_weights["state_dict"][
        f"transformer.decoder.block.{i}.layer.0.SelfAttention.k.weight"
    ]
    state_dict[f"decoder.block.{i}.layer.0.SelfAttention.v.weight"] = official_weights["state_dict"][
        f"transformer.decoder.block.{i}.layer.0.SelfAttention.v.weight"
    ]
    state_dict[f"decoder.block.{i}.layer.0.SelfAttention.o.weight"] = official_weights["state_dict"][
        f"transformer.decoder.block.{i}.layer.0.SelfAttention.o.weight"
    ]
    state_dict[f"decoder.block.{i}.layer.0.layer_norm.weight"] = official_weights["state_dict"][
        f"transformer.decoder.block.{i}.layer.0.layer_norm.weight"
    ]

    
    state_dict[f"decoder.block.{i}.layer.1.EncDecAttention.q.weight"] = official_weights["state_dict"][
        f"transformer.decoder.block.{i}.layer.1.EncDecAttention.q.weight"
    ]
    state_dict[f"decoder.block.{i}.layer.1.EncDecAttention.k.weight"] = official_weights["state_dict"][
        f"transformer.decoder.block.{i}.layer.1.EncDecAttention.k.weight"
    ]
    state_dict[f"decoder.block.{i}.layer.1.EncDecAttention.v.weight"] = official_weights["state_dict"][
        f"transformer.decoder.block.{i}.layer.1.EncDecAttention.v.weight"
    ]
    state_dict[f"decoder.block.{i}.layer.1.EncDecAttention.o.weight"] = official_weights["state_dict"][
        f"transformer.decoder.block.{i}.layer.1.EncDecAttention.o.weight"
    ]
    state_dict[f"decoder.block.{i}.layer.1.layer_norm.weight"] = official_weights["state_dict"][
        f"transformer.decoder.block.{i}.layer.1.layer_norm.weight"
    ]

    
    state_dict[f"decoder.block.{i}.layer.2.DenseReluDense.wi_0.weight"] = official_weights["state_dict"][
        f"transformer.decoder.block.{i}.layer.2.DenseReluDense.wi_0.weight"
    ]
    state_dict[f"decoder.block.{i}.layer.2.DenseReluDense.wi_1.weight"] = official_weights["state_dict"][
        f"transformer.decoder.block.{i}.layer.2.DenseReluDense.wi_1.weight"
    ]
    state_dict[f"decoder.block.{i}.layer.2.DenseReluDense.wo.weight"] = official_weights["state_dict"][
        f"transformer.decoder.block.{i}.layer.2.DenseReluDense.wo.weight"
    ]
    state_dict[f"decoder.block.{i}.layer.2.layer_norm.weight"] = official_weights["state_dict"][
        f"transformer.decoder.block.{i}.layer.2.layer_norm.weight"
    ]

model.load_state_dict(state_dict, strict=True)


torch.save(state_dict, "./pytorch_model.bin")







def tokenize(idx, token_type, n_special=4, n_note=128, n_velocity=2):
    if token_type == "TOKEN_TIME":
        return n_special + n_note + n_velocity + idx
    elif token_type == "TOKEN_VELOCITY":
        return n_special + n_note + idx
    elif token_type == "TOKEN_NOTE":
        return n_special + idx
    elif token_type == "TOKEN_SPECIAL":
        return idx
    else:
        return -1



def detokenize(idx, n_special=4, n_note=128, n_velocity=2, time_idx_offset=0):
    if idx >= n_special + n_note + n_velocity:
        return "TOKEN_TIME", (idx - (n_special + n_note + n_velocity)) + time_idx_offset
    elif idx >= n_special + n_note:
        return "TOKEN_VELOCITY", idx - (n_special + n_note)
    elif idx >= n_special:
        return "TOKEN_NOTE", idx - n_special
    else:
        return "TOKEN_SPECIAL", idx



decoder = {}
for i in range(cfg.vocab_size):
    decoder.update({i: f"{detokenize(i)[1]}_{detokenize(i)[0]}"})

encoder = {v: k for k, v in decoder.items()}


with open("./vocab.json", "w") as file:
    file.write(json.dumps(encoder))