source dump

app.py

@@ -0,0 +1,62 @@
+from utils.musicautobot.numpy_encode import *
+from utils.musicautobot.utils.file_processing import process_all, process_file
+from utils.musicautobot.config import *
+from utils.musicautobot.music_transformer import *
+
+import gradio as gr
+from midi2audio import FluidSynth
+import tempfile
+import os
+
+# Bootloading model
+data_path = Path('./')
+data = MusicDataBunch.empty(data_path)
+vocab = data.vocab
+pretrained_path='./music_transformer.pth'
+learn = music_model_learner(data, pretrained_path=pretrained_path, config=default_config())
+
+
+def predict(seed_midi, n_words=400, temperature1=1.1, temperature2=0.4, min_bars=12, top_k=24, top_p=0.7):
+    # Load input MIDI file as MusicItem
+    cutoff_beat = 10
+    item = MusicItem.from_file(seed_midi.name, data.vocab)
+    seed_item = item.trim_to_beat(cutoff_beat)
+
+    # Generate prediction
+    pred, full = learn.predict(seed_item, n_words=n_words, temperatures=(temperature1, temperature2), min_bars=min_bars, top_k=top_k, top_p=top_p)
+
+    # Convert input MIDI to audio
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as seed_audio_temp:
+        FluidSynth("sound_font.sf2").midi_to_audio(seed_midi.name, seed_audio_temp.name)
+
+    # Save generated MIDI as temporary file
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.midi') as pred_midi_temp:
+        pred.stream.write('midi', fp=pred_midi_temp.name)
+
+    # Convert generated MIDI to audio
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as pred_audio_temp:
+        FluidSynth("sound_font.sf2").midi_to_audio(pred_midi_temp.name, pred_audio_temp.name)
+
+    # Cleanup temporary MIDI file
+    os.remove(pred_midi_temp.name)
+
+    return seed_audio_temp.name, pred_audio_temp.name
+
+
+
+iface = gr.Interface(fn=predict,
+                     inputs=[
+                         gr.inputs.File(label="Seed MIDI"),
+                         gr.inputs.Slider(50, 1000, step=10, default=400, label="Number of Words"),
+                         gr.inputs.Slider(0.0, 2.0, step=0.1, default=1.1, label="Temperature 1"),
+                         gr.inputs.Slider(0.0, 2.0, step=0.1, default=0.4, label="Temperature 2"),
+                         gr.inputs.Slider(1, 32, step=1, default=12, label="Min Bars"),
+                         gr.inputs.Slider(1, 50, step=1, default=24, label="Top K"),
+                         gr.inputs.Slider(0.0, 1.0, step=0.1, default=0.7, label="Top P")
+                     ],
+                     outputs=[
+                         gr.outputs.Audio(type='filepath', label="Seed Audio"),
+                         gr.outputs.Audio(type='filepath', label="Generated Audio")
+                     ],)
+
+iface.launch()

music_transformer.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9856190b46abee88440104c661349f577eca4754ae485b63cf77030772b0c8cf
+size 657241884

requirements.txt

@@ -0,0 +1,4 @@
+gradio
+midi2audio
+music21
+fastai

utils/musicautobot/__init__.py

@@ -0,0 +1,3 @@
+from .utils.setup_musescore import setup_musescore
+
+setup_musescore()

utils/musicautobot/config.py

@@ -0,0 +1,47 @@
+from fastai.text.models.transformer import tfmerXL_lm_config, Activation
+# from .vocab import MusicVocab
+
+def default_config():
+    config = tfmerXL_lm_config.copy()
+    config['act'] = Activation.GeLU
+
+    config['mem_len'] = 512
+    config['d_model'] = 512
+    config['d_inner'] = 2048
+    config['n_layers'] = 16
+    
+    config['n_heads'] = 8
+    config['d_head'] = 64
+
+    return config
+
+def music_config():
+    config = default_config()
+    config['encode_position'] = True
+    return config
+
+def musicm_config():
+    config = music_config()
+    config['d_model'] = 768
+    config['d_inner'] = 3072
+    config['n_heads'] = 12
+    config['d_head'] = 64
+    config['n_layers'] = 12
+    return config
+    
+def multitask_config():
+    config = default_config()
+    config['bias'] = True
+    config['enc_layers'] = 8
+    config['dec_layers'] = 8
+    del config['n_layers']
+    return config
+
+def multitaskm_config():
+    config = musicm_config()
+    config['bias'] = True
+    config['enc_layers'] = 12
+    config['dec_layers'] = 12
+    del config['n_layers']
+    return config
+

utils/musicautobot/multitask_transformer/__init__.py

@@ -0,0 +1,3 @@
+from .dataloader import *
+from .model import *
+from .learner import *

utils/musicautobot/multitask_transformer/dataloader.py

@@ -0,0 +1,146 @@
+from fastai.basics import *
+from .transform import *
+from ..music_transformer.dataloader import MusicDataBunch, MusicItemList
+# Sequence 2 Sequence Translate
+
+class S2SFileProcessor(PreProcessor):
+    "`PreProcessor` that opens the filenames and read the texts."
+    def process_one(self,item):
+        out = np.load(item, allow_pickle=True)
+        if out.shape != (2,): return None
+        if not 16 < len(out[0]) < 2048: return None
+        if not 16 < len(out[1]) < 2048: return None
+        return out
+    
+    def process(self, ds:Collection):
+        ds.items = [self.process_one(item) for item in ds.items]
+        ds.items = [i for i in ds.items if i is not None] # filter out None
+
+class S2SPartsProcessor(PreProcessor):
+    "Encodes midi file into 2 separate parts - melody and chords."
+    
+    def process_one(self, item):
+        m, c = item
+        mtrack = MultitrackItem.from_npenc_parts(m, c, vocab=self.vocab)
+        return mtrack.to_idx()
+    
+    def process(self, ds):
+        self.vocab = ds.vocab
+        ds.items = [self.process_one(item) for item in ds.items]
+        
+class Midi2MultitrackProcessor(PreProcessor):
+    "Converts midi files to multitrack items"
+    def process_one(self, midi_file):
+        try:
+            item = MultitrackItem.from_file(midi_file, vocab=self.vocab)
+        except Exception as e:
+            print(e)
+            return None
+        return item.to_idx()
+        
+    def process(self, ds):
+        self.vocab = ds.vocab
+        ds.items = [self.process_one(item) for item in ds.items]
+        ds.items = [i for i in ds.items if i is not None]
+    
+class S2SPreloader(Callback):
+    def __init__(self, dataset:LabelList, bptt:int=512, 
+                 transpose_range=None, **kwargs):
+        self.dataset,self.bptt = dataset,bptt
+        self.vocab = self.dataset.vocab
+        self.transpose_range = transpose_range
+        self.rand_transpose = partial(rand_transpose_value, rand_range=transpose_range) if transpose_range is not None else None
+        
+    def __getitem__(self, k:int):
+        item,empty_label = self.dataset[k]
+        
+        if self.rand_transpose is not None:
+            val = self.rand_transpose()
+            item = item.transpose(val)
+        item = item.pad_to(self.bptt+1)
+        ((m_x, m_pos), (c_x, c_pos)) = item.to_idx()
+        return m_x, m_pos, c_x, c_pos
+    
+    def __len__(self):
+        return len(self.dataset)
+
+def rand_transpose_value(rand_range=(0,24), p=0.5):
+    if np.random.rand() < p: return np.random.randint(*rand_range)-rand_range[1]//2
+    return 0
+
+class S2SItemList(MusicItemList):
+    _bunch = MusicDataBunch
+    def get(self, i):
+        return MultitrackItem.from_idx(self.items[i], self.vocab)
+
+# DATALOADING AND TRANSFORMATIONS
+# These transforms happen on batch
+
+def mask_tfm(b, mask_range, mask_idx, pad_idx, p=0.3):
+    # mask range (min, max)
+    # replacement vals - [x_replace, y_replace]. Usually [mask_idx, pad_idx]
+    # p = replacement probability
+    x,y = b
+    x,y = x.clone(),y.clone()
+    rand = torch.rand(x.shape, device=x.device)
+    rand[x < mask_range[0]] = 1.0
+    rand[x >= mask_range[1]] = 1.0
+    
+    # p(15%) of words are replaced. Of those p(15%) - 80% are masked. 10% wrong word. 10% unchanged
+    y[rand > p] = pad_idx # pad unchanged 80%. Remove these from loss/acc metrics
+    x[rand <= (p*.8)] = mask_idx # 80% = mask
+    wrong_word = (rand > (p*.8)) & (rand <= (p*.9)) # 10% = wrong word
+    x[wrong_word] = torch.randint(*mask_range, [wrong_word.sum().item()], device=x.device)
+    return x, y
+
+def mask_lm_tfm_default(b, vocab, mask_p=0.3):
+    return mask_lm_tfm(b, mask_range=vocab.npenc_range, mask_idx=vocab.mask_idx, pad_idx=vocab.pad_idx, mask_p=mask_p)
+
+def mask_lm_tfm_pitchdur(b, vocab, mask_p=0.9):
+    mask_range = vocab.dur_range if np.random.rand() < 0.5 else vocab.note_range
+    return mask_lm_tfm(b, mask_range=mask_range, mask_idx=vocab.mask_idx, pad_idx=vocab.pad_idx, mask_p=mask_p)
+
+def mask_lm_tfm(b, mask_range, mask_idx, pad_idx, mask_p):
+    x,y = b
+    x_lm,x_pos = x[...,0], x[...,1]
+    y_lm,y_pos = y[...,0], y[...,1]
+    
+    # Note: masking y_lm instead of x_lm. Just in case we ever do sequential s2s training
+    x_msk, y_msk = mask_tfm((y_lm, y_lm), mask_range=mask_range, mask_idx=mask_idx, pad_idx=pad_idx, p=mask_p)
+    msk_pos = y_pos
+    
+    x_dict = { 
+        'msk': { 'x': x_msk, 'pos': msk_pos },
+        'lm': { 'x': x_lm, 'pos': msk_pos }
+    }
+    y_dict = { 'msk': y_msk, 'lm': y_lm }
+    return x_dict, y_dict
+
+def melody_chord_tfm(b):
+    m,m_pos,c,c_pos = b
+    
+    # offset x and y for next word prediction
+    y_m = m[:,1:]
+    x_m, m_pos = m[:,:-1], m_pos[:,:-1]
+    
+    y_c = c[:,1:]
+    x_c, c_pos = c[:,:-1], c_pos[:,:-1]
+    
+    x_dict = { 
+        'c2m': {
+            'enc': x_c,
+            'enc_pos': c_pos,
+            'dec': x_m,
+            'dec_pos': m_pos
+        },
+        'm2c': {
+            'enc': x_m,
+            'enc_pos': m_pos,
+            'dec': x_c,
+            'dec_pos': c_pos
+        }
+    }
+    y_dict = {
+        'c2m': y_m, 'm2c': y_c
+    }
+    return x_dict, y_dict

utils/musicautobot/multitask_transformer/learner.py

@@ -0,0 +1,340 @@
+from fastai.basics import *
+from ..vocab import *
+from ..utils.top_k_top_p import top_k_top_p
+from ..utils.midifile import is_empty_midi
+from ..music_transformer.transform import *
+from ..music_transformer.learner import filter_invalid_indexes
+from .model import get_multitask_model
+from .dataloader import *
+
+def multitask_model_learner(data:DataBunch, config:dict=None, drop_mult:float=1., 
+                            pretrained_path:PathOrStr=None, **learn_kwargs) -> 'LanguageLearner':
+    "Create a `Learner` with a language model from `data` and `arch`."
+    vocab = data.vocab
+    vocab_size = len(vocab)
+
+    if pretrained_path: 
+        state = torch.load(pretrained_path, map_location='cpu')
+        if config is None: config = state['config']
+
+    model = get_multitask_model(vocab_size, config=config, drop_mult=drop_mult, pad_idx=vocab.pad_idx)
+    metrics = [AverageMultiMetric(partial(m, pad_idx=vocab.pad_idx)) for m in [mask_acc, lm_acc, c2m_acc, m2c_acc]]
+    loss_func = MultiLoss(ignore_index=data.vocab.pad_idx)
+    learn = MultitaskLearner(data, model, loss_func=loss_func, metrics=metrics, **learn_kwargs)
+    
+    if pretrained_path: 
+        get_model(model).load_state_dict(state['model'], strict=False)
+        if not hasattr(learn, 'opt'): learn.create_opt(defaults.lr, learn.wd)
+        try:    learn.opt.load_state_dict(state['opt'])
+        except: pass
+        del state
+        gc.collect()
+        
+    return learn
+
+class MultitaskLearner(Learner):
+    def save(self, file:PathLikeOrBinaryStream=None, with_opt:bool=True, config=None):
+        "Save model and optimizer state (if `with_opt`) with `file` to `self.model_dir`. `file` can be file-like (file or buffer)"
+        out_path = super().save(file, return_path=True, with_opt=with_opt)
+        if config and out_path:
+            state = torch.load(out_path)
+            state['config'] = config
+            torch.save(state, out_path)
+            del state
+            gc.collect()
+        return out_path
+
+    def predict_nw(self, item:MusicItem, n_words:int=128,
+                     temperatures:float=(1.0,1.0), min_bars=4,
+                     top_k=30, top_p=0.6):
+        "Return the `n_words` that come after `text`."
+        self.model.reset()
+        new_idx = []
+        vocab = self.data.vocab
+        x, pos = item.to_tensor(), item.get_pos_tensor()
+        last_pos = pos[-1] if len(pos) else 0
+        y = torch.tensor([0])
+
+        start_pos = last_pos
+
+        sep_count = 0
+        bar_len = SAMPLE_FREQ * 4 # assuming 4/4 time
+        vocab = self.data.vocab
+
+        repeat_count = 0
+
+        for i in progress_bar(range(n_words), leave=True):
+            batch = { 'lm': { 'x': x[None], 'pos': pos[None] } }, y
+            logits = self.pred_batch(batch=batch)['lm'][-1][-1]
+
+            prev_idx = new_idx[-1] if len(new_idx) else vocab.pad_idx
+
+            # Temperature
+            # Use first temperatures value if last prediction was duration
+            temperature = temperatures[0] if vocab.is_duration_or_pad(prev_idx) else temperatures[1]
+            repeat_penalty = max(0, np.log((repeat_count+1)/4)/5) * temperature
+            temperature += repeat_penalty
+            if temperature != 1.: logits = logits / temperature
+                
+
+            # Filter
+            # bar = 16 beats
+            filter_value = -float('Inf')
+            if ((last_pos - start_pos) // 16) <= min_bars: logits[vocab.bos_idx] = filter_value
+
+            logits = filter_invalid_indexes(logits, prev_idx, vocab, filter_value=filter_value)
+            logits = top_k_top_p(logits, top_k=top_k, top_p=top_p, filter_value=filter_value)
+            
+            # Sample
+            probs = F.softmax(logits, dim=-1)
+            idx = torch.multinomial(probs, 1).item()
+
+            # Update repeat count
+            num_choices = len(probs.nonzero().view(-1))
+            if num_choices <= 2: repeat_count += 1
+            else: repeat_count = repeat_count // 2
+
+            if prev_idx==vocab.sep_idx: 
+                duration = idx - vocab.dur_range[0]
+                last_pos = last_pos + duration
+
+                bars_pred = (last_pos - start_pos) // 16
+                abs_bar = last_pos // 16
+                # if (bars % 8 == 0) and (bars_pred > min_bars): break
+                if (i / n_words > 0.80) and (abs_bar % 4 == 0): break
+
+
+            if idx==vocab.bos_idx: 
+                print('Predicted BOS token. Returning prediction...')
+                break
+
+            new_idx.append(idx)
+            x = x.new_tensor([idx])
+            pos = pos.new_tensor([last_pos])
+
+        pred = vocab.to_music_item(np.array(new_idx))
+        full = item.append(pred)
+        return pred, full
+
+    def predict_mask(self, masked_item:MusicItem,
+                    temperatures:float=(1.0,1.0),
+                    top_k=20, top_p=0.8):
+        x = masked_item.to_tensor()
+        pos = masked_item.get_pos_tensor()
+        y = torch.tensor([0])
+        vocab = self.data.vocab
+        self.model.reset()
+        mask_idxs = (x == vocab.mask_idx).nonzero().view(-1)
+
+        repeat_count = 0
+
+        for midx in progress_bar(mask_idxs, leave=True):
+            prev_idx = x[midx-1]
+
+            # Using original positions, otherwise model gets too off track
+            # pos = torch.tensor(-position_enc(xb[0].cpu().numpy()), device=xb.device)[None]
+    
+            # Next Word
+            logits = self.pred_batch(batch=({ 'msk': { 'x': x[None], 'pos': pos[None] } }, y) )['msk'][0][midx]
+
+            # Temperature
+            # Use first temperatures value if last prediction was duration
+            temperature = temperatures[0] if vocab.is_duration_or_pad(prev_idx) else temperatures[1]
+            repeat_penalty = max(0, np.log((repeat_count+1)/4)/5) * temperature
+            temperature += repeat_penalty
+            if temperature != 1.: logits = logits / temperature
+
+            # Filter
+            filter_value = -float('Inf')
+            special_idxs = [vocab.bos_idx, vocab.sep_idx, vocab.stoi[EOS]]
+            logits[special_idxs] = filter_value # Don't allow any special tokens (as we are only removing notes and durations)
+            logits = filter_invalid_indexes(logits, prev_idx, vocab, filter_value=filter_value)
+            logits = top_k_top_p(logits, top_k=top_k, top_p=top_p, filter_value=filter_value)
+
+            # Sampling
+            probs = F.softmax(logits, dim=-1)
+            idx = torch.multinomial(probs, 1).item()
+
+            # Update repeat count
+            num_choices = len(probs.nonzero().view(-1))
+            if num_choices <= 2: repeat_count += 1
+            else: repeat_count = repeat_count // 2
+
+            x[midx] = idx
+
+        return vocab.to_music_item(x.cpu().numpy())
+
+    def predict_s2s(self, input_item:MusicItem, target_item:MusicItem, n_words:int=256,
+                        temperatures:float=(1.0,1.0), top_k=30, top_p=0.8,
+                        use_memory=True):
+        vocab = self.data.vocab
+        
+        # Input doesn't change. We can reuse the encoder output on each prediction
+        with torch.no_grad():
+            inp, inp_pos = input_item.to_tensor(), input_item.get_pos_tensor()
+            x_enc = self.model.encoder(inp[None], inp_pos[None])
+        
+        # target
+        targ = target_item.data.tolist()
+        targ_pos = target_item.position.tolist()
+        last_pos = targ_pos[-1]
+        self.model.reset()
+
+        repeat_count = 0
+
+        max_pos = input_item.position[-1] + SAMPLE_FREQ * 4 # Only predict until both tracks/parts have the same length
+        x, pos = inp.new_tensor(targ), inp_pos.new_tensor(targ_pos)
+        
+        for i in progress_bar(range(n_words), leave=True):
+            # Predict
+            with torch.no_grad():
+                dec = self.model.decoder(x[None], pos[None], x_enc)
+                logits = self.model.head(dec)[-1, -1]
+
+            # Temperature
+            # Use first temperatures value if last prediction was duration
+            prev_idx = targ[-1] if len(targ) else vocab.pad_idx
+            temperature = temperatures[0] if vocab.is_duration_or_pad(prev_idx) else temperatures[1]
+            repeat_penalty = max(0, np.log((repeat_count+1)/4)/5) * temperature
+            temperature += repeat_penalty
+            if temperature != 1.: logits = logits / temperature
+                
+            # Filter
+            filter_value = -float('Inf')
+            logits = filter_invalid_indexes(logits, prev_idx, vocab, filter_value=filter_value)
+            logits = top_k_top_p(logits, top_k=top_k, top_p=top_p, filter_value=filter_value)
+
+            # Sample
+            probs = F.softmax(logits, dim=-1)
+            idx = torch.multinomial(probs, 1).item()
+
+            # Update repeat count
+            num_choices = len(probs.nonzero().view(-1))
+            if num_choices <= 2: repeat_count += 1
+            else: repeat_count = repeat_count // 2
+
+            if idx == vocab.bos_idx | idx == vocab.stoi[EOS]: 
+                print('Predicting BOS/EOS')
+                break
+
+            if prev_idx == vocab.sep_idx: 
+                duration = idx - vocab.dur_range[0]
+                last_pos = last_pos + duration
+                if last_pos > max_pos:
+                    print('Predicted past counter-part length. Returning early')
+                    break
+
+            targ_pos.append(last_pos)
+            targ.append(idx)
+            
+            if use_memory:
+                # Relying on memory for kv. Only need last prediction index
+                x, pos = inp.new_tensor([targ[-1]]), inp_pos.new_tensor([targ_pos[-1]])
+            else:
+                # Reset memory after each prediction, since we feeding the whole sequence every time
+                self.model.reset()
+                x, pos = inp.new_tensor(targ), inp_pos.new_tensor(targ_pos)
+
+        return vocab.to_music_item(np.array(targ))
+    
+# High level prediction functions from midi file
+def nw_predict_from_midi(learn, midi=None, n_words=400, 
+                      temperatures=(1.0,1.0), top_k=30, top_p=0.6, seed_len=None, **kwargs):
+    vocab = learn.data.vocab
+    seed = MusicItem.from_file(midi, vocab) if not is_empty_midi(midi) else MusicItem.empty(vocab)
+    if seed_len is not None: seed = seed.trim_to_beat(seed_len)
+        
+    pred, full = learn.predict_nw(seed, n_words=n_words, temperatures=temperatures, top_k=top_k, top_p=top_p, **kwargs)
+    return full
+
+def s2s_predict_from_midi(learn, midi=None, n_words=200, 
+                      temperatures=(1.0,1.0), top_k=24, top_p=0.7, seed_len=None, pred_melody=True, **kwargs):
+    multitrack_item = MultitrackItem.from_file(midi, learn.data.vocab)
+    melody, chords = multitrack_item.melody, multitrack_item.chords
+    inp, targ = (chords, melody) if pred_melody else (melody, chords)
+    
+    # if seed_len is passed, cutoff sequence so we can predict the rest
+    if seed_len is not None: targ = targ.trim_to_beat(seed_len)
+    targ = targ.remove_eos()
+        
+    pred = learn.predict_s2s(inp, targ, n_words=n_words, temperatures=temperatures, top_k=top_k, top_p=top_p, **kwargs)
+    
+    part_order = (pred, inp) if pred_melody else (inp, pred)
+    return MultitrackItem(*part_order)
+
+def mask_predict_from_midi(learn, midi=None, predict_notes=True,
+                           temperatures=(1.0,1.0), top_k=30, top_p=0.7, section=None, **kwargs):
+    item = MusicItem.from_file(midi, learn.data.vocab)
+    masked_item = item.mask_pitch(section) if predict_notes else item.mask_duration(section)
+    pred = learn.predict_mask(masked_item, temperatures=temperatures, top_k=top_k, top_p=top_p, **kwargs)
+    return pred
+
+# LOSS AND METRICS
+
+class MultiLoss():
+    def __init__(self, ignore_index=None):
+        "Loss mult - Mask, NextWord, Seq2Seq"
+        self.loss = CrossEntropyFlat(ignore_index=ignore_index)
+        
+    def __call__(self, inputs:Dict[str,Tensor], targets:Dict[str,Tensor])->Rank0Tensor:
+        losses = [self.loss(inputs[key], target) for key,target in targets.items()]
+        return sum(losses)
+    
+def acc_ignore_pad(input:Tensor, targ:Tensor, pad_idx)->Rank0Tensor:
+    if input is None or targ is None: return None
+    n = targ.shape[0]
+    input = input.argmax(dim=-1).view(n,-1)
+    targ = targ.view(n,-1)
+    mask = targ != pad_idx
+    return (input[mask]==targ[mask]).float().mean()
+
+def acc_index(inputs, targets, key, pad_idx):
+    return acc_ignore_pad(inputs.get(key), targets.get(key), pad_idx)
+    
+def mask_acc(inputs, targets, pad_idx): return acc_index(inputs, targets, 'msk', pad_idx)
+def lm_acc(inputs, targets, pad_idx): return acc_index(inputs, targets, 'lm', pad_idx)
+def c2m_acc(inputs, targets, pad_idx): return acc_index(inputs, targets, 'c2m', pad_idx)
+def m2c_acc(inputs, targets, pad_idx): return acc_index(inputs, targets, 'm2c', pad_idx)
+
+
+class AverageMultiMetric(AverageMetric):
+    "Updated fastai.AverageMetric to support multi task metrics."
+    def on_batch_end(self, last_output, last_target, **kwargs):
+        "Update metric computation with `last_output` and `last_target`."
+        if not is_listy(last_target): last_target=[last_target]
+        val = self.func(last_output, *last_target)
+        if val is None: return
+        self.count += first_el(last_target).size(0)
+        if self.world:
+            val = val.clone()
+            dist.all_reduce(val, op=dist.ReduceOp.SUM)
+            val /= self.world
+        self.val += first_el(last_target).size(0) * val.detach().cpu()
+
+    def on_epoch_end(self, last_metrics, **kwargs):
+        "Set the final result in `last_metrics`."
+        if self.count == 0: return add_metrics(last_metrics, 0)
+        return add_metrics(last_metrics, self.val/self.count)
+    
+
+# MODEL LOADING
+class MTTrainer(LearnerCallback):
+    "`Callback` that regroups lr adjustment to seq_len, AR and TAR."
+    def __init__(self, learn:Learner, dataloaders=None, starting_mask_window=1):
+        super().__init__(learn)
+        self.count = 1
+        self.mw_start = starting_mask_window
+        self.dataloaders = dataloaders
+
+    def on_epoch_begin(self, **kwargs):
+        "Reset the hidden state of the model."
+        model = get_model(self.learn.model)
+        model.reset()
+        model.encoder.mask_steps = max(self.count+self.mw_start, 100)
+        
+    def on_epoch_end(self, last_metrics, **kwargs):
+        "Finish the computation and sends the result to the Recorder."
+        if self.dataloaders is not None: 
+            self.learn.data = self.dataloaders[self.count % len(self.dataloaders)]
+        self.count += 1
+

utils/musicautobot/multitask_transformer/model.py

@@ -0,0 +1,258 @@
+from fastai.basics import *
+from fastai.text.models.transformer import Activation, PositionalEncoding, feed_forward, init_transformer, _line_shift
+from fastai.text.models.awd_lstm import RNNDropout
+from ..utils.attention_mask import *
+
+def get_multitask_model(vocab_size:int, config:dict=None, drop_mult:float=1., pad_idx=None):
+    "Create a language model from `arch` and its `config`, maybe `pretrained`."
+    for k in config.keys(): 
+        if k.endswith('_p'): config[k] *= drop_mult
+    n_hid = config['d_model']
+    mem_len = config.pop('mem_len')
+    embed = TransformerEmbedding(vocab_size, n_hid, embed_p=config['embed_p'], mem_len=mem_len, pad_idx=pad_idx)
+    encoder = MTEncoder(embed, n_hid, n_layers=config['enc_layers'], mem_len=0, **config) # encoder doesn't need memory
+    decoder = MTEncoder(embed, n_hid, is_decoder=True, n_layers=config['dec_layers'], mem_len=mem_len, **config)
+    head = MTLinearDecoder(n_hid, vocab_size, tie_encoder=embed.embed, **config)
+    model = MultiTransformer(encoder, decoder, head, mem_len=mem_len)
+    return model.apply(init_transformer)
+
+class MultiTransformer(nn.Module):
+    "Multitask Transformer for training mask, next word, and sequence 2 sequence"
+    def __init__(self, encoder, decoder, head, mem_len):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.head = head
+        self.default_mem_len = mem_len
+        self.current_mem_len = None
+    
+    def forward(self, inp):
+        # data order: mask, next word, melody, chord
+        outputs = {}
+        msk, lm, c2m, m2c = [inp.get(key) for key in ['msk', 'lm', 'c2m', 'm2c']]
+        
+        if msk is not None:
+            outputs['msk'] = self.head(self.encoder(msk['x'], msk['pos']))
+        if lm is not None:
+            outputs['lm'] = self.head(self.decoder(lm['x'], lm['pos']))
+        
+        if c2m is not None:
+            self.reset()
+            c2m_enc = self.encoder(c2m['enc'], c2m['enc_pos'])
+            c2m_dec = self.decoder(c2m['dec'], c2m['dec_pos'], c2m_enc)
+            outputs['c2m'] = self.head(c2m_dec)
+            
+        if m2c is not None:
+            self.reset()
+            m2c_enc = self.encoder(m2c['enc'], m2c['enc_pos'])
+            m2c_dec = self.decoder(m2c['dec'], m2c['dec_pos'], m2c_enc)
+            outputs['m2c'] = self.head(m2c_dec)
+            
+        return outputs
+    
+    "A sequential module that passes the reset call to its children."
+    def reset(self):
+        for module in self.children(): 
+            reset_children(module)
+        
+def reset_children(mod):
+    if hasattr(mod, 'reset'): mod.reset()
+    for module in mod.children(): 
+        reset_children(module)
+
+ # COMPONENTS
+class TransformerEmbedding(nn.Module):
+    "Embedding + positional encoding + dropout"
+    def __init__(self, vocab_size:int, emb_sz:int, embed_p:float=0., mem_len=512, beat_len=32, max_bar_len=1024, pad_idx=None):
+        super().__init__()
+        self.emb_sz = emb_sz
+        self.pad_idx = pad_idx
+        
+        self.embed = nn.Embedding(vocab_size, emb_sz, padding_idx=pad_idx)
+        self.pos_enc = PositionalEncoding(emb_sz)
+        self.beat_len, self.max_bar_len = beat_len, max_bar_len
+        self.beat_enc = nn.Embedding(beat_len, emb_sz, padding_idx=0)
+        self.bar_enc = nn.Embedding(max_bar_len, emb_sz, padding_idx=0)
+        
+        self.drop = nn.Dropout(embed_p)
+        self.mem_len = mem_len
+    
+    def forward(self, inp, pos):
+        beat_enc = self.beat_enc(pos % self.beat_len)
+        bar_pos = pos // self.beat_len % self.max_bar_len
+        bar_pos[bar_pos >= self.max_bar_len] = self.max_bar_len - 1
+        bar_enc = self.bar_enc((bar_pos))
+        emb = self.drop(self.embed(inp) + beat_enc + bar_enc)
+        return emb
+    
+    def relative_pos_enc(self, emb):
+#         return torch.arange(640-1, -1, -1).float().cuda()
+        seq_len = emb.shape[1] + self.mem_len
+        pos = torch.arange(seq_len-1, -1, -1, device=emb.device, dtype=emb.dtype) # backwards (txl pos encoding)
+        return self.pos_enc(pos)
+
+class MTLinearDecoder(nn.Module):
+    "To go on top of a RNNCore module and create a Language Model."
+    initrange=0.1
+
+    def __init__(self, n_hid:int, n_out:int, output_p:float, tie_encoder:nn.Module=None, out_bias:bool=True, **kwargs):
+        super().__init__()
+        self.decoder = nn.Linear(n_hid, n_out, bias=out_bias)
+        self.decoder.weight.data.uniform_(-self.initrange, self.initrange)
+        self.output_dp = RNNDropout(output_p)
+        if out_bias: self.decoder.bias.data.zero_()
+        if tie_encoder: self.decoder.weight = tie_encoder.weight
+
+    def forward(self, input:Tuple[Tensor,Tensor])->Tuple[Tensor,Tensor,Tensor]:
+        output = self.output_dp(input)
+        decoded = self.decoder(output)
+        return decoded
+
+    
+# DECODER TRANSLATE BLOCK
+class MTEncoder(nn.Module):
+    def __init__(self, embed:nn.Module, n_hid:int, n_layers:int, n_heads:int, d_model:int, d_head:int, d_inner:int, 
+                 resid_p:float=0., attn_p:float=0., ff_p:float=0., bias:bool=True, scale:bool=True,
+                 act:Activation=Activation.ReLU, double_drop:bool=True, mem_len:int=512, is_decoder=False,
+                 mask_steps=1, mask_p=0.3, **kwargs):
+        super().__init__()
+        self.embed = embed
+        self.u = nn.Parameter(torch.Tensor(n_heads, 1, d_head)) #Remove 1 for einsum implementation of attention
+        self.v = nn.Parameter(torch.Tensor(n_heads, 1, d_head)) #Remove 1 for einsum implementation of attention
+        self.n_layers,self.d_model = n_layers,d_model
+        self.layers = nn.ModuleList([MTEncoderBlock(n_heads, d_model, d_head, d_inner, resid_p=resid_p, attn_p=attn_p,
+                      ff_p=ff_p, bias=bias, scale=scale, act=act, double_drop=double_drop, mem_len=mem_len,
+                      ) for k in range(n_layers)])
+
+        self.mask_steps, self.mask_p = mask_steps, mask_p
+        self.is_decoder = is_decoder
+    
+        nn.init.normal_(self.u, 0., 0.02)
+        nn.init.normal_(self.v, 0., 0.02)
+        
+    def forward(self, x_lm, lm_pos, msk_emb=None):
+        bs,lm_len = x_lm.size()
+        
+        lm_emb = self.embed(x_lm, lm_pos)
+        if msk_emb is not None and msk_emb.shape[1] > lm_emb.shape[1]:
+            pos_enc = self.embed.relative_pos_enc(msk_emb)
+        else:
+            pos_enc = self.embed.relative_pos_enc(lm_emb)
+    
+        # Masks
+        if self.is_decoder:
+            lm_mask = rand_window_mask(lm_len, self.embed.mem_len, x_lm.device,
+                                       max_size=self.mask_steps, p=self.mask_p, is_eval=not self.training)
+        else:
+            lm_mask = None
+        
+        for i, layer in enumerate(self.layers):
+            lm_emb = layer(lm_emb, msk_emb, lm_mask=lm_mask,
+                        r=pos_enc, g_u=self.u, g_v=self.v)
+        return lm_emb
+
+class MTEncoderBlock(nn.Module):
+    "Decoder block of a Transformer model."
+    #Can't use Sequential directly cause more than one input...
+    def __init__(self, n_heads:int, d_model:int, d_head:int, d_inner:int, resid_p:float=0., attn_p:float=0., ff_p:float=0.,
+                 bias:bool=True, scale:bool=True, double_drop:bool=True, mem_len:int=512, mha2_mem_len=0, **kwargs):
+        super().__init__()
+        attn_cls = MemMultiHeadRelativeAttentionKV
+        self.mha1 = attn_cls(n_heads, d_model, d_head, resid_p=resid_p, attn_p=attn_p, bias=bias, scale=scale, mem_len=mem_len, r_mask=False)
+        self.mha2 = attn_cls(n_heads, d_model, d_head, resid_p=resid_p, attn_p=attn_p, bias=bias, scale=scale, mem_len=mha2_mem_len, r_mask=True)
+        self.ff   = feed_forward(d_model, d_inner, ff_p=ff_p, double_drop=double_drop)
+    
+    def forward(self, enc_lm:Tensor, enc_msk:Tensor,
+                r=None, g_u=None, g_v=None,
+                msk_mask:Tensor=None, lm_mask:Tensor=None): 
+        
+        y_lm = self.mha1(enc_lm, enc_lm, enc_lm, r, g_u, g_v, mask=lm_mask)
+        if enc_msk is None: return y_lm
+        return self.ff(self.mha2(y_lm, enc_msk, enc_msk, r, g_u, g_v, mask=msk_mask))
+    
+
+    # Attention Layer
+
+
+# Attn
+
+class MemMultiHeadRelativeAttentionKV(nn.Module):
+    "Attention Layer monster - relative positioning, keeps track of own memory, separate kv weights to support sequence2sequence decoding."
+    def __init__(self, n_heads:int, d_model:int, d_head:int=None, resid_p:float=0., attn_p:float=0., bias:bool=True,
+                 scale:bool=True, mem_len:int=512, r_mask=True):
+        super().__init__()
+        d_head = ifnone(d_head, d_model//n_heads)
+        self.n_heads,self.d_head,self.scale = n_heads,d_head,scale
+        
+        assert(d_model == d_head * n_heads)
+        self.q_wgt = nn.Linear(d_model, n_heads * d_head, bias=bias)
+        self.k_wgt = nn.Linear(d_model, n_heads * d_head, bias=bias)
+        self.v_wgt = nn.Linear(d_model, n_heads * d_head, bias=bias)
+        
+        self.drop_att,self.drop_res = nn.Dropout(attn_p),nn.Dropout(resid_p)
+        self.ln = nn.LayerNorm(d_model)
+        self.r_attn = nn.Linear(d_model, n_heads * d_head, bias=bias)
+        self.r_mask = r_mask
+
+        self.mem_len = mem_len
+        self.prev_k = None
+        self.prev_v = None
+        
+    def forward(self, q:Tensor, k:Tensor=None, v:Tensor=None, 
+                r:Tensor=None, g_u:Tensor=None, g_v:Tensor=None, 
+                mask:Tensor=None, **kwargs):
+        if k is None: k = q
+        if v is None: v = q
+        return self.ln(q + self.drop_res(self._apply_attention(q, k, v, r, g_u, g_v, mask=mask, **kwargs)))
+
+    def mem_k(self, k):
+        if self.mem_len == 0: return k
+        if self.prev_k is None or (self.prev_k.shape[0] != k.shape[0]): # reset if wrong batch size
+            self.prev_k = k[:, -self.mem_len:]
+            return k
+        with torch.no_grad():
+            k_ext = torch.cat([self.prev_k, k], dim=1)
+            self.prev_k = k_ext[:, -self.mem_len:]
+        return k_ext.detach()
+    
+    def mem_v(self, v):
+        if self.mem_len == 0: return v
+        if self.prev_v is None or (self.prev_v.shape[0] != v.shape[0]): # reset if wrong batch size
+            self.prev_v = v[:, -self.mem_len:]
+            return v
+        with torch.no_grad():
+            v_ext = torch.cat([self.prev_v, v], dim=1)
+            self.prev_v = v_ext[:, -self.mem_len:]
+        return v_ext.detach()
+        
+    def reset(self):
+        self.prev_v = None
+        self.prev_k = None
+        
+    def _apply_attention(self, q:Tensor, k:Tensor, v:Tensor, 
+                         r:Tensor=None, g_u:Tensor=None, g_v:Tensor=None, 
+                         mask:Tensor=None, **kwargs):
+        #Notations from the paper: x input, r vector of relative distance between two elements, u et v learnable
+        #parameters of the model common between all layers, mask to avoid cheating and mem the previous hidden states.
+#         bs,x_len,seq_len = q.size(0),q.size(1),r.size(0)
+        k = self.mem_k(k)
+        v = self.mem_v(v)
+        bs,x_len,seq_len = q.size(0),q.size(1),k.size(1)
+        wq,wk,wv = self.q_wgt(q),self.k_wgt(k),self.v_wgt(v)
+        wq = wq[:,-x_len:]
+        wq,wk,wv = map(lambda x:x.view(bs, x.size(1), self.n_heads, self.d_head), (wq,wk,wv))
+        wq,wk,wv = wq.permute(0, 2, 1, 3),wk.permute(0, 2, 3, 1),wv.permute(0, 2, 1, 3)
+        wkr = self.r_attn(r[-seq_len:])
+        wkr = wkr.view(seq_len, self.n_heads, self.d_head)
+        wkr = wkr.permute(1,2,0)
+        #### compute attention score (AC is (a) + (c) and BS is (b) + (d) in the paper)
+        AC = torch.matmul(wq+g_u,wk)
+        BD = _line_shift(torch.matmul(wq+g_v, wkr), mask=self.r_mask)
+        if self.scale: attn_score = (AC + BD).mul_(1/(self.d_head ** 0.5))
+        if mask is not None: 
+            mask = mask[...,-seq_len:]
+            if hasattr(mask, 'bool'): mask = mask.bool()
+            attn_score = attn_score.float().masked_fill(mask, -float('inf')).type_as(attn_score)
+        attn_prob = self.drop_att(F.softmax(attn_score, dim=-1))
+        attn_vec = torch.matmul(attn_prob, wv)
+        return attn_vec.permute(0, 2, 1, 3).contiguous().view(bs, x_len, -1)

utils/musicautobot/multitask_transformer/transform.py

@@ -0,0 +1,68 @@
+from ..music_transformer.transform import *
+
+class MultitrackItem():
+    def __init__(self, melody:MusicItem, chords:MusicItem, stream=None):
+        self.melody,self.chords = melody, chords
+        self.vocab = melody.vocab
+        self._stream = stream
+        
+    @classmethod
+    def from_file(cls, midi_file, vocab):
+        return cls.from_stream(file2stream(midi_file), vocab)
+
+    @classmethod
+    def from_stream(cls, stream, vocab):
+        if not isinstance(stream, music21.stream.Score): stream = stream.voicesToParts()
+        num_parts = len(stream.parts)
+        sort_pitch = False
+        if num_parts > 2: 
+            raise ValueError('Could not extract melody and chords from midi file. Please make sure file contains exactly 2 tracks')
+        elif num_parts == 1: 
+            print('Warning: only 1 track found. Inferring melody/chords')
+            stream = separate_melody_chord(stream)
+            sort_pitch = False
+            
+        mpart, cpart = stream2npenc_parts(stream, sort_pitch=sort_pitch)
+        return cls.from_npenc_parts(mpart, cpart, vocab, stream)
+        
+    @classmethod
+    def from_npenc_parts(cls, mpart, cpart, vocab, stream=None):
+        mpart = npenc2idxenc(mpart, seq_type=SEQType.Melody, vocab=vocab, add_eos=False)
+        cpart = npenc2idxenc(cpart, seq_type=SEQType.Chords, vocab=vocab, add_eos=False)
+        return MultitrackItem(MusicItem(mpart, vocab), MusicItem(cpart, vocab), stream)
+        
+    @classmethod
+    def from_idx(cls, item, vocab):
+        m, c = item
+        return MultitrackItem(MusicItem.from_idx(m, vocab), MusicItem.from_idx(c, vocab))
+    def to_idx(self): return np.array((self.melody.to_idx(), self.chords.to_idx()))
+    
+    @property
+    def stream(self):
+        self._stream = self.to_stream() if self._stream is None else self._stream
+        return self._stream
+    
+    def to_stream(self, bpm=120):
+        ps = self.melody.to_npenc(), self.chords.to_npenc()
+        ps = [npenc2chordarr(p) for p in ps]
+        chordarr = chordarr_combine_parts(ps)
+        return chordarr2stream(chordarr, bpm=bpm)
+
+    
+    def show(self, format:str=None):
+        return self.stream.show(format)
+    def play(self): self.stream.show('midi')
+        
+    def transpose(self, val):
+        return MultitrackItem(self.melody.transpose(val), self.chords.transpose(val))
+    def pad_to(self, val):
+        return MultitrackItem(self.melody.pad_to(val), self.chords.pad_to(val))
+    def trim_to_beat(self, beat):
+        return MultitrackItem(self.melody.trim_to_beat(beat), self.chords.trim_to_beat(beat))
+    
+def combine2chordarr(np1, np2, vocab):
+    if len(np1.shape) == 1: np1 = idxenc2npenc(np1, vocab)
+    if len(np2.shape) == 1: np2 = idxenc2npenc(np2, vocab)
+    p1 = npenc2chordarr(np1)
+    p2 = npenc2chordarr(np2)
+    return chordarr_combine_parts((p1, p2))

utils/musicautobot/music_transformer/__init__.py

@@ -0,0 +1,3 @@
+from .dataloader import *
+from .model import *
+from .learner import *

utils/musicautobot/music_transformer/dataloader.py

@@ -0,0 +1,229 @@
+"Fastai Language Model Databunch modified to work with music"
+from fastai.basics import *
+# from fastai.basic_data import DataBunch
+from fastai.text.data import LMLabelList
+from .transform import *
+from ..vocab import MusicVocab
+
+
+class MusicDataBunch(DataBunch):
+    "Create a `TextDataBunch` suitable for training a language model."
+    @classmethod
+    def create(cls, train_ds, valid_ds, test_ds=None, path:PathOrStr='.', no_check:bool=False, bs=64, val_bs:int=None, 
+               num_workers:int=0, device:torch.device=None, collate_fn:Callable=data_collate, 
+               dl_tfms:Optional[Collection[Callable]]=None, bptt:int=70,
+               preloader_cls=None, shuffle_dl=False, transpose_range=(0,12), **kwargs) -> DataBunch:
+        "Create a `TextDataBunch` in `path` from the `datasets` for language modelling."
+        datasets = cls._init_ds(train_ds, valid_ds, test_ds)
+        preloader_cls = MusicPreloader if preloader_cls is None else preloader_cls
+        val_bs = ifnone(val_bs, bs)
+        datasets = [preloader_cls(ds, shuffle=(i==0), bs=(bs if i==0 else val_bs), bptt=bptt, transpose_range=transpose_range, **kwargs) 
+                    for i,ds in enumerate(datasets)]
+        val_bs = bs
+        dl_tfms = [partially_apply_vocab(tfm, train_ds.vocab) for tfm in listify(dl_tfms)]
+        dls = [DataLoader(d, b, shuffle=shuffle_dl) for d,b in zip(datasets, (bs,val_bs,val_bs,val_bs)) if d is not None]
+        return cls(*dls, path=path, device=device, dl_tfms=dl_tfms, collate_fn=collate_fn, no_check=no_check)
+    
+    @classmethod    
+    def from_folder(cls, path:PathOrStr, extensions='.npy', **kwargs):
+        files = get_files(path, extensions=extensions, recurse=True);
+        return cls.from_files(files, path, **kwargs)
+    
+    @classmethod
+    def from_files(cls, files, path, processors=None, split_pct=0.1, 
+                   vocab=None, list_cls=None, **kwargs):
+        if vocab is None: vocab = MusicVocab.create()
+        if list_cls is None: list_cls = MusicItemList
+        src = (list_cls(items=files, path=path, processor=processors, vocab=vocab)
+                .split_by_rand_pct(split_pct, seed=6)
+                .label_const(label_cls=LMLabelList))
+        return src.databunch(**kwargs)
+
+    @classmethod
+    def empty(cls, path, **kwargs):
+        vocab = MusicVocab.create()
+        src = MusicItemList([], path=path, vocab=vocab, ignore_empty=True).split_none()
+        return src.label_const(label_cls=LMLabelList).databunch()
+        
+def partially_apply_vocab(tfm, vocab):
+    if 'vocab' in inspect.getfullargspec(tfm).args:
+        return partial(tfm, vocab=vocab)
+    return tfm
+    
+class MusicItemList(ItemList):
+    _bunch = MusicDataBunch
+    
+    def __init__(self, items:Iterator, vocab:MusicVocab=None, **kwargs):
+        super().__init__(items, **kwargs)
+        self.vocab = vocab
+        self.copy_new += ['vocab']
+    
+    def get(self, i):
+        o = super().get(i)
+        if is_pos_enc(o): 
+            return MusicItem.from_idx(o, self.vocab)
+        return MusicItem(o, self.vocab)
+
+def is_pos_enc(idxenc):
+    if len(idxenc.shape) == 2 and idxenc.shape[0] == 2: return True
+    return idxenc.dtype == np.object and idxenc.shape == (2,)
+
+class MusicItemProcessor(PreProcessor):
+    "`PreProcessor` that transforms numpy files to indexes for training"
+    def process_one(self,item):
+        item = MusicItem.from_npenc(item, vocab=self.vocab)
+        return item.to_idx()
+    
+    def process(self, ds):
+        self.vocab = ds.vocab
+        super().process(ds)
+        
+class OpenNPFileProcessor(PreProcessor):
+    "`PreProcessor` that opens the filenames and read the texts."
+    def process_one(self,item):
+        return np.load(item, allow_pickle=True) if isinstance(item, Path) else item
+
+class Midi2ItemProcessor(PreProcessor):
+    "Skips midi preprocessing step. And encodes midi files to MusicItems"
+    def process_one(self,item):
+        item = MusicItem.from_file(item, vocab=self.vocab)
+        return item.to_idx()
+    
+    def process(self, ds):
+        self.vocab = ds.vocab
+        super().process(ds)
+    
+## For npenc dataset
+class MusicPreloader(Callback):
+    "Transforms the tokens in `dataset` to a stream of contiguous batches for language modelling."
+    
+    class CircularIndex():
+        "Handles shuffle, direction of indexing, wraps around to head tail in the ragged array as needed"
+        def __init__(self, length:int, forward:bool): self.idx, self.forward = np.arange(length), forward
+        def __getitem__(self, i): 
+            return self.idx[ i%len(self.idx) if self.forward else len(self.idx)-1-i%len(self.idx)]
+        def __len__(self) -> int: return len(self.idx)
+        def shuffle(self): np.random.shuffle(self.idx)
+
+    def __init__(self, dataset:LabelList, lengths:Collection[int]=None, bs:int=32, bptt:int=70, backwards:bool=False, 
+                 shuffle:bool=False, y_offset:int=1, 
+                 transpose_range=None, transpose_p=0.5,
+                 encode_position=True,
+                 **kwargs):
+        self.dataset,self.bs,self.bptt,self.shuffle,self.backwards,self.lengths = dataset,bs,bptt,shuffle,backwards,lengths
+        self.vocab = self.dataset.vocab
+        self.bs *= num_distrib() or 1
+        self.totalToks,self.ite_len,self.idx = int(0),None,None
+        self.y_offset = y_offset
+        
+        self.transpose_range,self.transpose_p = transpose_range,transpose_p
+        self.encode_position = encode_position
+        self.bptt_len = self.bptt
+        
+        self.allocate_buffers() # needed for valid_dl on distributed training - otherwise doesn't get initialized on first epoch
+
+    def __len__(self): 
+        if self.ite_len is None:
+            if self.lengths is None: self.lengths = np.array([len(item) for item in self.dataset.x])
+            self.totalToks = self.lengths.sum()
+            self.ite_len   = self.bs*int( math.ceil( self.totalToks/(self.bptt*self.bs) )) if self.item is None else 1
+        return self.ite_len
+
+    def __getattr__(self,k:str)->Any: return getattr(self.dataset, k)
+   
+    def allocate_buffers(self):
+        "Create the ragged array that will be filled when we ask for items."
+        if self.ite_len is None: len(self)
+        self.idx   = MusicPreloader.CircularIndex(len(self.dataset.x), not self.backwards)
+        
+        # batch shape = (bs, bptt, 2 - [index, pos]) if encode_position. Else - (bs, bptt)
+        buffer_len = (2,) if self.encode_position else ()
+        self.batch = np.zeros((self.bs, self.bptt+self.y_offset) + buffer_len, dtype=np.int64)
+        self.batch_x, self.batch_y = self.batch[:,0:self.bptt], self.batch[:,self.y_offset:self.bptt+self.y_offset] 
+        #ro: index of the text we're at inside our datasets for the various batches
+        self.ro    = np.zeros(self.bs, dtype=np.int64)
+        #ri: index of the token we're at inside our current text for the various batches
+        self.ri    = np.zeros(self.bs, dtype=np.int)
+        
+        # allocate random transpose values. Need to allocate this before hand.
+        self.transpose_values = self.get_random_transpose_values()
+        
+    def get_random_transpose_values(self):
+        if self.transpose_range is None: return None
+        n = len(self.dataset)
+        rt_arr = torch.randint(*self.transpose_range, (n,))-self.transpose_range[1]//2
+        mask = torch.rand(rt_arr.shape) > self.transpose_p
+        rt_arr[mask] = 0
+        return rt_arr
+
+    def on_epoch_begin(self, **kwargs):
+        if self.idx is None: self.allocate_buffers()
+        elif self.shuffle:   
+            self.ite_len = None
+            self.idx.shuffle()
+            self.transpose_values = self.get_random_transpose_values()
+            self.bptt_len = self.bptt
+        self.idx.forward = not self.backwards 
+
+        step = self.totalToks / self.bs
+        ln_rag, countTokens, i_rag = 0, 0, -1
+        for i in range(0,self.bs):
+            #Compute the initial values for ro and ri 
+            while ln_rag + countTokens <= int(step * i):
+                countTokens += ln_rag
+                i_rag       += 1
+                ln_rag       = self.lengths[self.idx[i_rag]]
+            self.ro[i] = i_rag
+            self.ri[i] = ( ln_rag - int(step * i - countTokens) ) if self.backwards else int(step * i - countTokens)
+        
+    #Training dl gets on_epoch_begin called, val_dl, on_epoch_end
+    def on_epoch_end(self, **kwargs): self.on_epoch_begin()
+
+    def __getitem__(self, k:int):
+        j = k % self.bs
+        if j==0:
+            if self.item is not None: return self.dataset[0]
+            if self.idx is None: self.on_epoch_begin()
+                
+        self.ro[j],self.ri[j] = self.fill_row(not self.backwards, self.dataset.x, self.idx, self.batch[j][:self.bptt_len+self.y_offset], 
+                                              self.ro[j], self.ri[j], overlap=1, lengths=self.lengths)
+        return self.batch_x[j][:self.bptt_len], self.batch_y[j][:self.bptt_len]
+
+    def fill_row(self, forward, items, idx, row, ro, ri, overlap, lengths):
+        "Fill the row with tokens from the ragged array. --OBS-- overlap != 1 has not been implemented"
+        ibuf = n = 0 
+        ro  -= 1
+        while ibuf < row.shape[0]:  
+            ro   += 1 
+            ix    = idx[ro]
+            
+            item = items[ix]
+            if self.transpose_values is not None: 
+                item = item.transpose(self.transpose_values[ix].item())
+                
+            if self.encode_position:
+                # Positions are colomn stacked with indexes. This makes it easier to keep in sync
+                rag = np.stack([item.data, item.position], axis=1)
+            else:
+                rag = item.data
+                
+            if forward:
+                ri = 0 if ibuf else ri
+                n  = min(lengths[ix] - ri, row.shape[0] - ibuf)
+                row[ibuf:ibuf+n] = rag[ri:ri+n]
+            else:    
+                ri = lengths[ix] if ibuf else ri
+                n  = min(ri, row.size - ibuf) 
+                row[ibuf:ibuf+n] = rag[ri-n:ri][::-1]
+            ibuf += n
+        return ro, ri + ((n-overlap) if forward else -(n-overlap))
+
+def batch_position_tfm(b):
+    "Batch transform for training with positional encoding"
+    x,y = b
+    x = {
+        'x': x[...,0],
+        'pos': x[...,1]
+    }
+    return x, y[...,0]
+    

utils/musicautobot/music_transformer/learner.py

@@ -0,0 +1,171 @@
+from fastai.basics import *
+from fastai.text.learner import LanguageLearner, get_language_model, _model_meta
+from .model import *
+from .transform import MusicItem
+from ..numpy_encode import SAMPLE_FREQ
+from ..utils.top_k_top_p import top_k_top_p
+from ..utils.midifile import is_empty_midi
+
+_model_meta[MusicTransformerXL] = _model_meta[TransformerXL] # copy over fastai's model metadata
+
+def music_model_learner(data:DataBunch, arch=MusicTransformerXL, config:dict=None, drop_mult:float=1.,
+                        pretrained_path:PathOrStr=None, **learn_kwargs) -> 'LanguageLearner':
+    "Create a `Learner` with a language model from `data` and `arch`."
+    meta = _model_meta[arch]
+
+    if pretrained_path: 
+        state = torch.load(pretrained_path, map_location='cpu')
+        if config is None: config = state['config']
+        
+    model = get_language_model(arch, len(data.vocab.itos), config=config, drop_mult=drop_mult)
+    learn = MusicLearner(data, model, split_func=meta['split_lm'], **learn_kwargs)
+
+    if pretrained_path: 
+        get_model(model).load_state_dict(state['model'], strict=False)
+        if not hasattr(learn, 'opt'): learn.create_opt(defaults.lr, learn.wd)
+        try:    learn.opt.load_state_dict(state['opt'])
+        except: pass
+        del state
+        gc.collect()
+
+    return learn
+
+# Predictions
+from fastai import basic_train # for predictions
+class MusicLearner(LanguageLearner):
+    def save(self, file:PathLikeOrBinaryStream=None, with_opt:bool=True, config=None):
+        "Save model and optimizer state (if `with_opt`) with `file` to `self.model_dir`. `file` can be file-like (file or buffer)"
+        out_path = super().save(file, return_path=True, with_opt=with_opt)
+        if config and out_path:
+            state = torch.load(out_path)
+            state['config'] = config
+            torch.save(state, out_path)
+            del state
+            gc.collect()
+        return out_path
+
+    def beam_search(self, xb:Tensor, n_words:int, top_k:int=10, beam_sz:int=10, temperature:float=1.,
+                    ):
+        "Return the `n_words` that come after `text` using beam search."
+        self.model.reset()
+        self.model.eval()
+        xb_length = xb.shape[-1]
+        if xb.shape[0] > 1: xb = xb[0][None]
+        yb = torch.ones_like(xb)
+
+        nodes = None
+        xb = xb.repeat(top_k, 1)
+        nodes = xb.clone()
+        scores = xb.new_zeros(1).float()
+        with torch.no_grad():
+            for k in progress_bar(range(n_words), leave=False):
+                out = F.log_softmax(self.model(xb)[0][:,-1], dim=-1)
+                values, indices = out.topk(top_k, dim=-1)
+                scores = (-values + scores[:,None]).view(-1)
+                indices_idx = torch.arange(0,nodes.size(0))[:,None].expand(nodes.size(0), top_k).contiguous().view(-1)
+                sort_idx = scores.argsort()[:beam_sz]
+                scores = scores[sort_idx]
+                nodes = torch.cat([nodes[:,None].expand(nodes.size(0),top_k,nodes.size(1)),
+                                indices[:,:,None].expand(nodes.size(0),top_k,1),], dim=2)
+                nodes = nodes.view(-1, nodes.size(2))[sort_idx]
+                self.model[0].select_hidden(indices_idx[sort_idx])
+                xb = nodes[:,-1][:,None]
+        if temperature != 1.: scores.div_(temperature)
+        node_idx = torch.multinomial(torch.exp(-scores), 1).item()
+        return [i.item() for i in nodes[node_idx][xb_length:] ]
+
+    def predict(self, item:MusicItem, n_words:int=128,
+                     temperatures:float=(1.0,1.0), min_bars=4,
+                     top_k=30, top_p=0.6):
+        "Return the `n_words` that come after `text`."
+        self.model.reset()
+        new_idx = []
+        vocab = self.data.vocab
+        x, pos = item.to_tensor(), item.get_pos_tensor()
+        last_pos = pos[-1] if len(pos) else 0
+        y = torch.tensor([0])
+
+        start_pos = last_pos
+
+        sep_count = 0
+        bar_len = SAMPLE_FREQ * 4 # assuming 4/4 time
+        vocab = self.data.vocab
+
+        repeat_count = 0
+        if hasattr(self.model[0], 'encode_position'):
+            encode_position = self.model[0].encode_position
+        else: encode_position = False
+
+        for i in progress_bar(range(n_words), leave=True):
+            with torch.no_grad():
+                if encode_position:
+                    batch = { 'x': x[None], 'pos': pos[None] }
+                    logits = self.model(batch)[0][-1][-1]
+                else:
+                    logits = self.model(x[None])[0][-1][-1]
+
+            prev_idx = new_idx[-1] if len(new_idx) else vocab.pad_idx
+
+            # Temperature
+            # Use first temperatures value if last prediction was duration
+            temperature = temperatures[0] if vocab.is_duration_or_pad(prev_idx) else temperatures[1]
+            repeat_penalty = max(0, np.log((repeat_count+1)/4)/5) * temperature
+            temperature += repeat_penalty
+            if temperature != 1.: logits = logits / temperature
+                
+
+            # Filter
+            # bar = 16 beats
+            filter_value = -float('Inf')
+            if ((last_pos - start_pos) // 16) <= min_bars: logits[vocab.bos_idx] = filter_value
+
+            logits = filter_invalid_indexes(logits, prev_idx, vocab, filter_value=filter_value)
+            logits = top_k_top_p(logits, top_k=top_k, top_p=top_p, filter_value=filter_value)
+            
+            # Sample
+            probs = F.softmax(logits, dim=-1)
+            idx = torch.multinomial(probs, 1).item()
+
+            # Update repeat count
+            num_choices = len(probs.nonzero().view(-1))
+            if num_choices <= 2: repeat_count += 1
+            else: repeat_count = repeat_count // 2
+
+            if prev_idx==vocab.sep_idx: 
+                duration = idx - vocab.dur_range[0]
+                last_pos = last_pos + duration
+
+                bars_pred = (last_pos - start_pos) // 16
+                abs_bar = last_pos // 16
+                # if (bars % 8 == 0) and (bars_pred > min_bars): break
+                if (i / n_words > 0.80) and (abs_bar % 4 == 0): break
+
+
+            if idx==vocab.bos_idx: 
+                print('Predicted BOS token. Returning prediction...')
+                break
+
+            new_idx.append(idx)
+            x = x.new_tensor([idx])
+            pos = pos.new_tensor([last_pos])
+
+        pred = vocab.to_music_item(np.array(new_idx))
+        full = item.append(pred)
+        return pred, full
+    
+# High level prediction functions from midi file
+def predict_from_midi(learn, midi=None, n_words=400, 
+                      temperatures=(1.0,1.0), top_k=30, top_p=0.6, seed_len=None, **kwargs):
+    vocab = learn.data.vocab
+    seed = MusicItem.from_file(midi, vocab) if not is_empty_midi(midi) else MusicItem.empty(vocab)
+    if seed_len is not None: seed = seed.trim_to_beat(seed_len)
+
+    pred, full = learn.predict(seed, n_words=n_words, temperatures=temperatures, top_k=top_k, top_p=top_p, **kwargs)
+    return full
+
+def filter_invalid_indexes(res, prev_idx, vocab, filter_value=-float('Inf')):
+    if vocab.is_duration_or_pad(prev_idx):
+        res[list(range(*vocab.dur_range))] = filter_value
+    else:
+        res[list(range(*vocab.note_range))] = filter_value
+    return res

utils/musicautobot/music_transformer/model.py

@@ -0,0 +1,66 @@
+from fastai.basics import *
+from fastai.text.models.transformer import TransformerXL
+from ..utils.attention_mask import rand_window_mask
+
+class MusicTransformerXL(TransformerXL):
+    "Exactly like fastai's TransformerXL, but with more aggressive attention mask: see `rand_window_mask`"
+    def __init__(self, *args, encode_position=True, mask_steps=1, **kwargs):
+        import inspect
+        sig = inspect.signature(TransformerXL)
+        arg_params = { k:kwargs[k] for k in sig.parameters if k in kwargs }
+        super().__init__(*args, **arg_params)
+
+        self.encode_position = encode_position
+        if self.encode_position: self.beat_enc = BeatPositionEncoder(kwargs['d_model'])
+            
+        self.mask_steps=mask_steps
+        
+        
+    def forward(self, x):
+        #The hidden state has to be initiliazed in the forward pass for nn.DataParallel
+        if self.mem_len > 0 and not self.init: 
+            self.reset()
+            self.init = True
+
+        benc = 0
+        if self.encode_position:
+            x,pos = x['x'], x['pos']
+            benc = self.beat_enc(pos)
+
+        bs,x_len = x.size()
+        inp = self.drop_emb(self.encoder(x) + benc) #.mul_(self.d_model ** 0.5)
+        m_len = self.hidden[0].size(1) if hasattr(self, 'hidden') and len(self.hidden[0].size()) > 1 else 0
+        seq_len = m_len + x_len
+        
+        mask = rand_window_mask(x_len, m_len, inp.device, max_size=self.mask_steps, is_eval=not self.training) if self.mask else None
+        if m_len == 0: mask[...,0,0] = 0
+        #[None,:,:None] for einsum implementation of attention
+        hids = []
+        pos = torch.arange(seq_len-1, -1, -1, device=inp.device, dtype=inp.dtype)
+        pos_enc = self.pos_enc(pos)
+        hids.append(inp)
+        for i, layer in enumerate(self.layers):
+            mem = self.hidden[i] if self.mem_len > 0 else None
+            inp = layer(inp, r=pos_enc, u=self.u, v=self.v, mask=mask, mem=mem)
+            hids.append(inp)
+        core_out = inp[:,-x_len:]
+        if self.mem_len > 0 : self._update_mems(hids)
+        return (self.hidden if self.mem_len > 0 else [core_out]),[core_out]
+
+
+ # Beat encoder
+class BeatPositionEncoder(nn.Module):
+    "Embedding + positional encoding + dropout"
+    def __init__(self, emb_sz:int, beat_len=32, max_bar_len=1024):
+        super().__init__()
+
+        self.beat_len, self.max_bar_len = beat_len, max_bar_len
+        self.beat_enc = nn.Embedding(beat_len, emb_sz, padding_idx=0)
+        self.bar_enc = nn.Embedding(max_bar_len, emb_sz, padding_idx=0)
+    
+    def forward(self, pos):
+        beat_enc = self.beat_enc(pos % self.beat_len)
+        bar_pos = pos // self.beat_len % self.max_bar_len
+        bar_pos[bar_pos >= self.max_bar_len] = self.max_bar_len - 1
+        bar_enc = self.bar_enc((bar_pos))
+        return beat_enc + bar_enc

utils/musicautobot/music_transformer/transform.py

@@ -0,0 +1,235 @@
+from ..numpy_encode import *
+import numpy as np
+from enum import Enum
+import torch
+from ..vocab import *
+from functools import partial
+
+SEQType = Enum('SEQType', 'Mask, Sentence, Melody, Chords, Empty')
+
+class MusicItem():
+    def __init__(self, data, vocab, stream=None, position=None):
+        self.data = data
+        self.vocab = vocab
+        self._stream = stream
+        self._position = position
+    def __repr__(self): return '\n'.join([
+        f'\n{self.__class__.__name__} - {self.data.shape}',
+        f'{self.vocab.textify(self.data[:10])}...'])
+    def __len__(self): return len(self.data)
+
+    @classmethod
+    def from_file(cls, midi_file, vocab):
+        return cls.from_stream(file2stream(midi_file), vocab)
+    @classmethod
+    def from_stream(cls, stream, vocab):
+        if not isinstance(stream, music21.stream.Score): stream = stream.voicesToParts()
+        chordarr = stream2chordarr(stream) # 2.
+        npenc = chordarr2npenc(chordarr) # 3.
+        return cls.from_npenc(npenc, vocab, stream)
+    @classmethod
+    def from_npenc(cls, npenc, vocab, stream=None): return MusicItem(npenc2idxenc(npenc, vocab), vocab, stream)
+
+    @classmethod
+    def from_idx(cls, item, vocab):
+        idx,pos = item
+        return MusicItem(idx, vocab=vocab, position=pos)
+    def to_idx(self): return self.data, self.position
+
+    @classmethod
+    def empty(cls, vocab, seq_type=SEQType.Sentence):
+        return MusicItem(seq_prefix(seq_type, vocab), vocab)
+
+    @property
+    def stream(self):
+        self._stream = self.to_stream() if self._stream is None else self._stream
+        return self._stream
+
+    def to_stream(self, bpm=120):
+        return idxenc2stream(self.data, self.vocab, bpm=bpm)
+
+    def to_tensor(self, device=None):
+        return to_tensor(self.data, device)
+
+    def to_text(self, sep=' '): return self.vocab.textify(self.data, sep)
+
+    @property
+    def position(self):
+        self._position = position_enc(self.data, self.vocab) if self._position is None else self._position
+        return self._position
+
+    def get_pos_tensor(self, device=None): return to_tensor(self.position, device)
+
+    def to_npenc(self):
+        return idxenc2npenc(self.data, self.vocab)
+
+    def show(self, format:str=None):
+        return self.stream.show(format)
+    def play(self): self.stream.show('midi')
+
+    #Added by caslabs
+    def download(self, filename:str=None, ext:str=None):
+        return self.stream.write('midi', fp=filename)
+
+    @property
+    def new(self):
+        return partial(type(self), vocab=self.vocab)
+
+    def trim_to_beat(self, beat, include_last_sep=False):
+        return self.new(trim_to_beat(self.data, self.position, self.vocab, beat, include_last_sep))
+
+    def transpose(self, interval):
+        return self.new(tfm_transpose(self.data, interval, self.vocab), position=self._position)
+
+    def append(self, item):
+        return self.new(np.concatenate((self.data, item.data), axis=0))
+
+    def mask_pitch(self, section=None):
+        return self.new(self.mask(self.vocab.note_range, section), position=self.position)
+
+    def mask_duration(self, section=None, keep_position_enc=True):
+        masked_data = self.mask(self.vocab.dur_range, section)
+        if keep_position_enc: return self.new(masked_data, position=self.position)
+        return self.new(masked_data)
+
+    def mask(self, token_range, section_range=None):
+        return mask_section(self.data, self.position, token_range, self.vocab.mask_idx, section_range=section_range)
+
+    def pad_to(self, bptt):
+        data = pad_seq(self.data, bptt, self.vocab.pad_idx)
+        pos = pad_seq(self.position, bptt, 0)
+        return self.new(data, stream=self._stream, position=pos)
+
+    def split_stream_parts(self):
+        self._stream = separate_melody_chord(self.stream)
+        return self.stream
+
+    def remove_eos(self):
+        if self.data[-1] == self.vocab.stoi[EOS]: return self.new(self.data, stream=self.stream)
+        return self
+
+    def split_parts(self):
+        return self.new(self.data, stream=separate_melody_chord(self.stream), position=self.position)
+
+def pad_seq(seq, bptt, value):
+    pad_len = max(bptt-seq.shape[0], 0)
+    return np.pad(seq, (0, pad_len), 'constant', constant_values=value)[:bptt]
+
+def to_tensor(t, device=None):
+    t = t if isinstance(t, torch.Tensor) else torch.tensor(t)
+    if device is None and torch.cuda.is_available(): t = t.cuda()
+    else: t.to(device)
+    return t.long()
+
+def midi2idxenc(midi_file, vocab):
+    "Converts midi file to index encoding for training"
+    npenc = midi2npenc(midi_file) # 3.
+    return npenc2idxenc(npenc, vocab)
+
+def idxenc2stream(arr, vocab, bpm=120):
+    "Converts index encoding to music21 stream"
+    npenc = idxenc2npenc(arr, vocab)
+    return npenc2stream(npenc, bpm=bpm)
+
+# single stream instead of note,dur
+def npenc2idxenc(t, vocab, seq_type=SEQType.Sentence, add_eos=False):
+    "Transforms numpy array from 2 column (note, duration) matrix to a single column"
+    "[[n1, d1], [n2, d2], ...] -> [n1, d1, n2, d2]"
+    if isinstance(t, (list, tuple)) and len(t) == 2:
+        return [npenc2idxenc(x, vocab, start_seq) for x in t]
+    t = t.copy()
+
+    t[:, 0] = t[:, 0] + vocab.note_range[0]
+    t[:, 1] = t[:, 1] + vocab.dur_range[0]
+
+    prefix = seq_prefix(seq_type, vocab)
+    suffix = np.array([vocab.stoi[EOS]]) if add_eos else np.empty(0, dtype=int)
+    return np.concatenate([prefix, t.reshape(-1), suffix])
+
+def seq_prefix(seq_type, vocab):
+    if seq_type == SEQType.Empty: return np.empty(0, dtype=int)
+    start_token = vocab.bos_idx
+    if seq_type == SEQType.Chords: start_token = vocab.stoi[CSEQ]
+    if seq_type == SEQType.Melody: start_token = vocab.stoi[MSEQ]
+    return np.array([start_token, vocab.pad_idx])
+
+def idxenc2npenc(t, vocab, validate=True):
+    if validate: t = to_valid_idxenc(t, vocab.npenc_range)
+    t = t.copy().reshape(-1, 2)
+    if t.shape[0] == 0: return t
+
+    t[:, 0] = t[:, 0] - vocab.note_range[0]
+    t[:, 1] = t[:, 1] - vocab.dur_range[0]
+
+    if validate: return to_valid_npenc(t)
+    return t
+
+def to_valid_idxenc(t, valid_range):
+    r = valid_range
+    t = t[np.where((t >= r[0]) & (t < r[1]))]
+    if t.shape[-1] % 2 == 1: t = t[..., :-1]
+    return t
+
+def to_valid_npenc(t):
+    is_note = (t[:, 0] < VALTSEP) | (t[:, 0] >= NOTE_SIZE)
+    invalid_note_idx = is_note.argmax()
+    invalid_dur_idx = (t[:, 1] < 0).argmax()
+
+    invalid_idx = max(invalid_dur_idx, invalid_note_idx)
+    if invalid_idx > 0:
+        if invalid_note_idx > 0 and invalid_dur_idx > 0: invalid_idx = min(invalid_dur_idx, invalid_note_idx)
+        print('Non midi note detected. Only returning valid portion. Index, seed', invalid_idx, t.shape)
+        return t[:invalid_idx]
+    return t
+
+def position_enc(idxenc, vocab):
+    "Calculates positional beat encoding."
+    sep_idxs = (idxenc == vocab.sep_idx).nonzero()[0]
+    sep_idxs = sep_idxs[sep_idxs+2 < idxenc.shape[0]] # remove any indexes right before out of bounds (sep_idx+2)
+    dur_vals = idxenc[sep_idxs+1]
+    dur_vals[dur_vals == vocab.mask_idx] = vocab.dur_range[0] # make sure masked durations are 0
+    dur_vals -= vocab.dur_range[0]
+
+    posenc = np.zeros_like(idxenc)
+    posenc[sep_idxs+2] = dur_vals
+    return posenc.cumsum()
+
+def beat2index(idxenc, pos, vocab, beat, include_last_sep=False):
+    cutoff = find_beat(pos, beat)
+    if cutoff < 2: return 2 # always leave starter tokens
+    if len(idxenc) < 2 or include_last_sep: return cutoff
+    if idxenc[cutoff - 2] == vocab.sep_idx: return cutoff - 2
+    return cutoff
+
+def find_beat(pos, beat, sample_freq=SAMPLE_FREQ, side='left'):
+    return np.searchsorted(pos, beat * sample_freq, side=side)
+
+# TRANSFORMS
+
+def tfm_transpose(x, value, vocab):
+    x = x.copy()
+    x[(x >= vocab.note_range[0]) & (x < vocab.note_range[1])] += value
+    return x
+
+def trim_to_beat(idxenc, pos, vocab, to_beat=None, include_last_sep=True):
+    if to_beat is None: return idxenc
+    cutoff = beat2index(idxenc, pos, vocab, to_beat, include_last_sep=include_last_sep)
+    return idxenc[:cutoff]
+
+def mask_input(xb, mask_range, replacement_idx):
+    xb = xb.copy()
+    xb[(xb >= mask_range[0]) & (xb < mask_range[1])] = replacement_idx
+    return xb
+
+def mask_section(xb, pos, token_range, replacement_idx, section_range=None):
+    xb = xb.copy()
+    token_mask = (xb >= token_range[0]) & (xb < token_range[1])
+
+    if section_range is None: section_range = (None, None)
+    section_mask = np.zeros_like(xb, dtype=bool)
+    start_idx = find_beat(pos, section_range[0]) if section_range[0] is not None else 0
+    end_idx = find_beat(pos, section_range[1]) if section_range[1] is not None else xb.shape[0]
+    section_mask[start_idx:end_idx] = True
+
+    xb[token_mask & section_mask] = replacement_idx
+    return xb

utils/musicautobot/numpy_encode.py

@@ -0,0 +1,302 @@
+"Encoding music21 streams -> numpy array -> text"
+
+# import re
+import music21
+import numpy as np
+# from pathlib import Path
+
+BPB = 4 # beats per bar
+TIMESIG = f'{BPB}/4' # default time signature
+PIANO_RANGE = (21, 108)
+VALTSEP = -1 # separator value for numpy encoding
+VALTCONT = -2 # numpy value for TCONT - needed for compressing chord array
+
+SAMPLE_FREQ = 4
+NOTE_SIZE = 128
+DUR_SIZE = (10*BPB*SAMPLE_FREQ)+1 # Max length - 8 bars. Or 16 beats/quarternotes
+MAX_NOTE_DUR = (8*BPB*SAMPLE_FREQ)
+
+# Encoding process
+# 1. midi -> music21.Stream
+# 2. Stream -> numpy chord array (timestep X instrument X noterange)
+# 3. numpy array -> List[Timestep][NoteEnc]
+def midi2npenc(midi_file, skip_last_rest=True):
+    "Converts midi file to numpy encoding for language model"
+    stream = file2stream(midi_file) # 1.
+    chordarr = stream2chordarr(stream) # 2.
+    return chordarr2npenc(chordarr, skip_last_rest=skip_last_rest) # 3.
+
+# Decoding process
+# 1. NoteEnc -> numpy chord array
+# 2. numpy array -> music21.Stream
+def npenc2stream(arr, bpm=120):
+    "Converts numpy encoding to music21 stream"
+    chordarr = npenc2chordarr(np.array(arr)) # 1.
+    return chordarr2stream(chordarr, bpm=bpm) # 2.
+
+##### ENCODING ######
+
+# 1. File To STream
+
+def file2stream(fp):
+    if isinstance(fp, music21.midi.MidiFile): return music21.midi.translate.midiFileToStream(fp)
+    return music21.converter.parse(fp)
+
+# 2.
+def stream2chordarr(s, note_size=NOTE_SIZE, sample_freq=SAMPLE_FREQ, max_note_dur=MAX_NOTE_DUR):
+    "Converts music21.Stream to 1-hot numpy array"
+    # assuming 4/4 time
+    # note x instrument x pitch
+    # FYI: midi middle C value=60
+    
+    # (AS) TODO: need to order by instruments most played and filter out percussion or include the channel
+    highest_time = max(s.flat.getElementsByClass('Note').highestTime, s.flat.getElementsByClass('Chord').highestTime)
+    maxTimeStep = round(highest_time * sample_freq)+1
+    score_arr = np.zeros((maxTimeStep, len(s.parts), NOTE_SIZE))
+
+    def note_data(pitch, note):
+        return (pitch.midi, int(round(note.offset*sample_freq)), int(round(note.duration.quarterLength*sample_freq)))
+
+    for idx,part in enumerate(s.parts):
+        notes=[]
+        for elem in part.flat:
+            if isinstance(elem, music21.note.Note):
+                notes.append(note_data(elem.pitch, elem))
+            if isinstance(elem, music21.chord.Chord):
+                for p in elem.pitches:
+                    notes.append(note_data(p, elem))
+                
+        # sort notes by offset (1), duration (2) so that hits are not overwritten and longer notes have priority
+        notes_sorted = sorted(notes, key=lambda x: (x[1], x[2])) 
+        for n in notes_sorted:
+            if n is None: continue
+            pitch,offset,duration = n
+            if max_note_dur is not None and duration > max_note_dur: duration = max_note_dur
+            score_arr[offset, idx, pitch] = duration
+            score_arr[offset+1:offset+duration, idx, pitch] = VALTCONT      # Continue holding note
+    return score_arr
+
+def chordarr2npenc(chordarr, skip_last_rest=True):
+    # combine instruments
+    result = []
+    wait_count = 0
+    for idx,timestep in enumerate(chordarr):
+        flat_time = timestep2npenc(timestep)
+        if len(flat_time) == 0:
+            wait_count += 1
+        else:
+            # pitch, octave, duration, instrument
+            if wait_count > 0: result.append([VALTSEP, wait_count])
+            result.extend(flat_time)
+            wait_count = 1
+    if wait_count > 0 and not skip_last_rest: result.append([VALTSEP, wait_count])
+    return np.array(result, dtype=int).reshape(-1, 2) # reshaping. Just in case result is empty
+
+# Note: not worrying about overlaps - as notes will still play. just look tied
+# http://web.mit.edu/music21/doc/moduleReference/moduleStream.html#music21.stream.Stream.getOverlaps
+def timestep2npenc(timestep, note_range=PIANO_RANGE, enc_type=None):
+    # inst x pitch
+    notes = []
+    for i,n in zip(*timestep.nonzero()):
+        d = timestep[i,n]
+        if d < 0: continue # only supporting short duration encoding for now
+        if n < note_range[0] or n >= note_range[1]: continue # must be within midi range
+        notes.append([n,d,i])
+        
+    notes = sorted(notes, key=lambda x: x[0], reverse=True) # sort by note (highest to lowest)
+    
+    if enc_type is None: 
+        # note, duration
+        return [n[:2] for n in notes] 
+    if enc_type == 'parts':
+        # note, duration, part
+        return [n for n in notes]
+    if enc_type == 'full':
+        # note_class, duration, octave, instrument
+        return [[n%12, d, n//12, i] for n,d,i in notes] 
+
+##### DECODING #####
+
+# 1.
+def npenc2chordarr(npenc, note_size=NOTE_SIZE):
+    num_instruments = 1 if len(npenc.shape) <= 2 else npenc.max(axis=0)[-1]
+    
+    max_len = npenc_len(npenc)
+    # score_arr = (steps, inst, note)
+    score_arr = np.zeros((max_len, num_instruments, note_size))
+    
+    idx = 0
+    for step in npenc:
+        n,d,i = (step.tolist()+[0])[:3] # or n,d,i
+        if n < VALTSEP: continue # special token
+        if n == VALTSEP:
+            idx += d
+            continue
+        score_arr[idx,i,n] = d
+    return score_arr
+
+def npenc_len(npenc):
+    duration = 0
+    for t in npenc:
+        if t[0] == VALTSEP: duration += t[1]
+    return duration + 1
+
+
+# 2.
+def chordarr2stream(arr, sample_freq=SAMPLE_FREQ, bpm=120):
+    duration = music21.duration.Duration(1. / sample_freq)
+    stream = music21.stream.Score()
+    stream.append(music21.meter.TimeSignature(TIMESIG))
+    stream.append(music21.tempo.MetronomeMark(number=bpm))
+    stream.append(music21.key.KeySignature(0))
+    for inst in range(arr.shape[1]):
+        p = partarr2stream(arr[:,inst,:], duration)
+        stream.append(p)
+    stream = stream.transpose(0)
+    return stream
+
+# 2b.
+def partarr2stream(partarr, duration):
+    "convert instrument part to music21 chords"
+    part = music21.stream.Part()
+    part.append(music21.instrument.Piano())
+    part_append_duration_notes(partarr, duration, part) # notes already have duration calculated
+
+    return part
+
+def part_append_duration_notes(partarr, duration, stream):
+    "convert instrument part to music21 chords"
+    for tidx,t in enumerate(partarr):
+        note_idxs = np.where(t > 0)[0] # filter out any negative values (continuous mode)
+        if len(note_idxs) == 0: continue
+        notes = []
+        for nidx in note_idxs:
+            note = music21.note.Note(nidx)
+            note.duration = music21.duration.Duration(partarr[tidx,nidx]*duration.quarterLength)
+            notes.append(note)
+        for g in group_notes_by_duration(notes):
+            if len(g) == 1:
+                stream.insert(tidx*duration.quarterLength, g[0])
+            else:
+                chord = music21.chord.Chord(g)
+                stream.insert(tidx*duration.quarterLength, chord)
+    return stream
+
+from itertools import groupby
+#  combining notes with different durations into a single chord may overwrite conflicting durations. Example: aylictal/still-waters-run-deep
+def group_notes_by_duration(notes):
+    "separate notes into chord groups"
+    keyfunc = lambda n: n.duration.quarterLength
+    notes = sorted(notes, key=keyfunc)
+    return [list(g) for k,g in groupby(notes, keyfunc)]
+
+
+# Midi -> npenc Conversion helpers
+def is_valid_npenc(npenc, note_range=PIANO_RANGE, max_dur=DUR_SIZE, 
+                   min_notes=32, input_path=None, verbose=True):
+    if len(npenc) < min_notes:
+        if verbose: print('Sequence too short:', len(npenc), input_path)
+        return False
+    if (npenc[:,1] >= max_dur).any(): 
+        if verbose: print(f'npenc exceeds max {max_dur} duration:', npenc[:,1].max(), input_path)
+        return False
+    # https://en.wikipedia.org/wiki/Scientific_pitch_notation - 88 key range - 21 = A0, 108 = C8
+    if ((npenc[...,0] > VALTSEP) & ((npenc[...,0] < note_range[0]) | (npenc[...,0] >= note_range[1]))).any(): 
+        print(f'npenc out of piano note range {note_range}:', input_path)
+        return False
+    return True
+
+# seperates overlapping notes to different tracks
+def remove_overlaps(stream, separate_chords=True):
+    if not separate_chords:
+        return stream.flat.makeVoices().voicesToParts()
+    return separate_melody_chord(stream)
+
+# seperates notes and chords to different tracks
+def separate_melody_chord(stream):
+    new_stream = music21.stream.Score()
+    if stream.timeSignature: new_stream.append(stream.timeSignature)
+    new_stream.append(stream.metronomeMarkBoundaries()[0][-1])
+    if stream.keySignature: new_stream.append(stream.keySignature)
+    
+    melody_part = music21.stream.Part(stream.flat.getElementsByClass('Note'))
+    melody_part.insert(0, stream.getInstrument())
+    chord_part = music21.stream.Part(stream.flat.getElementsByClass('Chord'))
+    chord_part.insert(0, stream.getInstrument())
+    new_stream.append(melody_part)
+    new_stream.append(chord_part)
+    return new_stream
+
+# processing functions for sanitizing data
+
+def compress_chordarr(chordarr):
+    return shorten_chordarr_rests(trim_chordarr_rests(chordarr))
+
+def trim_chordarr_rests(arr, max_rests=4, sample_freq=SAMPLE_FREQ):
+    # max rests is in quarter notes
+    # max 1 bar between song start and end
+    start_idx = 0
+    max_sample = max_rests*sample_freq
+    for idx,t in enumerate(arr):
+        if (t != 0).any(): break
+        start_idx = idx+1
+        
+    end_idx = 0
+    for idx,t in enumerate(reversed(arr)):
+        if (t != 0).any(): break
+        end_idx = idx+1
+    start_idx = start_idx - start_idx % max_sample
+    end_idx = end_idx - end_idx % max_sample
+#     if start_idx > 0 or end_idx > 0: print('Trimming rests. Start, end:', start_idx, len(arr)-end_idx, end_idx)
+    return arr[start_idx:(len(arr)-end_idx)]
+
+def shorten_chordarr_rests(arr, max_rests=8, sample_freq=SAMPLE_FREQ):
+    # max rests is in quarter notes
+    # max 2 bar pause
+    rest_count = 0
+    result = []
+    max_sample = max_rests*sample_freq
+    for timestep in arr:
+        if (timestep==0).all(): 
+            rest_count += 1
+        else:
+            if rest_count > max_sample:
+#                 old_count = rest_count
+                rest_count = (rest_count % sample_freq) + max_sample
+#                 print(f'Compressing rests: {old_count} -> {rest_count}')
+            for i in range(rest_count): result.append(np.zeros(timestep.shape))
+            rest_count = 0
+            result.append(timestep)
+    for i in range(rest_count): result.append(np.zeros(timestep.shape))
+    return np.array(result)
+
+# sequence 2 sequence convenience functions
+
+def stream2npenc_parts(stream, sort_pitch=True):
+    chordarr = stream2chordarr(stream)
+    _,num_parts,_ = chordarr.shape
+    parts = [part_enc(chordarr, i) for i in range(num_parts)]
+    return sorted(parts, key=avg_pitch, reverse=True) if sort_pitch else parts
+
+def chordarr_combine_parts(parts):
+    max_ts = max([p.shape[0] for p in parts])
+    parts_padded = [pad_part_to(p, max_ts) for p in parts]
+    chordarr_comb = np.concatenate(parts_padded, axis=1)
+    return chordarr_comb
+
+def pad_part_to(p, target_size):
+    pad_width = ((0,target_size-p.shape[0]),(0,0),(0,0))
+    return np.pad(p, pad_width, 'constant')
+
+def part_enc(chordarr, part):
+    partarr = chordarr[:,part:part+1,:]
+    npenc = chordarr2npenc(partarr)
+    return npenc
+
+def avg_tempo(t, sep_idx=VALTSEP):
+    avg = t[t[:, 0] == sep_idx][:, 1].sum()/t.shape[0]
+    avg = int(round(avg/SAMPLE_FREQ))
+    return 'mt'+str(min(avg, MTEMPO_SIZE-1))
+
+def avg_pitch(t, sep_idx=VALTSEP):
+    return t[t[:, 0] > sep_idx][:, 0].mean()

utils/musicautobot/utils/attention_mask.py

@@ -0,0 +1,21 @@
+import numpy as np
+import torch
+
+def window_mask(x_len, device, m_len=0, size=(1,1)):
+    win_size,k = size
+    mem_mask = torch.zeros((x_len,m_len), device=device)
+    tri_mask = torch.triu(torch.ones((x_len//win_size+1,x_len//win_size+1), device=device),diagonal=k)
+    window_mask = tri_mask.repeat_interleave(win_size,dim=0).repeat_interleave(win_size,dim=1)[:x_len,:x_len]
+    if x_len: window_mask[...,0] = 0 # Always allowing first index to see. Otherwise you'll get NaN loss
+    mask = torch.cat((mem_mask, window_mask), dim=1)[None,None]
+    return mask.bool() if hasattr(mask, 'bool') else mask.byte()
+    
+def rand_window_mask(x_len,m_len,device,max_size:int=None,p:float=0.2,is_eval:bool=False):
+    if is_eval or np.random.rand() >= p or max_size is None: 
+        win_size,k = (1,1)
+    else: win_size,k = (np.random.randint(0,max_size)+1,0)
+    return window_mask(x_len, device, m_len, size=(win_size,k))
+
+def lm_mask(x_len, device):
+    mask = torch.triu(torch.ones((x_len, x_len), device=device), diagonal=1)[None,None]
+    return mask.bool() if hasattr(mask, 'bool') else mask.byte()

utils/musicautobot/utils/file_processing.py

@@ -0,0 +1,52 @@
+"Parallel processing for midi files"
+import csv
+from fastprogress.fastprogress import master_bar, progress_bar
+from pathlib import Path
+from pebble import ProcessPool
+from concurrent.futures import TimeoutError
+import numpy as np
+
+# https://stackoverflow.com/questions/20991968/asynchronous-multiprocessing-with-a-worker-pool-in-python-how-to-keep-going-aft
+def process_all(func, arr, timeout_func=None, total=None, max_workers=None, timeout=None):
+    with ProcessPool() as pool:
+        future = pool.map(func, arr, timeout=timeout)
+
+        iterator = future.result()
+        results = []
+        for i in progress_bar(range(len(arr)), total=len(arr)):
+            try:
+                result = next(iterator)
+                if result: results.append(result)
+            except StopIteration:
+                break  
+            except TimeoutError as error:
+                if timeout_func: timeout_func(arr[i], error.args[1])
+    return results
+
+def process_file(file_path, tfm_func=None, src_path=None, dest_path=None):
+    "Utility function that transforms midi file to numpy array."
+    output_file = Path(str(file_path).replace(str(src_path), str(dest_path))).with_suffix('.npy')
+    if output_file.exists(): return output_file
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+    
+    # Call tfm_func and save file
+    npenc = tfm_func(file_path)
+    if npenc is not None: 
+        np.save(output_file, npenc)
+        return output_file
+
+def arr2csv(arr, out_file):
+    "Convert metadata array to csv"
+    all_keys = {k for d in arr for k in d.keys()}
+    arr = [format_values(x) for x in arr]
+    with open(out_file, 'w') as f:
+        dict_writer = csv.DictWriter(f, list(all_keys))
+        dict_writer.writeheader()
+        dict_writer.writerows(arr)
+        
+def format_values(d):
+    "Format array values for csv encoding"
+    def format_value(v):
+        if isinstance(v, list): return ','.join(v)
+        return v
+    return {k:format_value(v) for k,v in d.items()}

utils/musicautobot/utils/lamb.py

@@ -0,0 +1,106 @@
+# SOURCE: https://github.com/cybertronai/pytorch-lamb/
+
+import collections
+import math
+
+import torch
+from torch.optim import Optimizer
+
+
+class Lamb(Optimizer):
+    r"""Implements Lamb algorithm.
+
+    It has been proposed in `Reducing BERT Pre-Training Time from 3 Days to 76 Minutes`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        adam (bool, optional): always use trust ratio = 1, which turns this into
+            Adam. Useful for comparison purposes.
+
+    .. _Reducing BERT Pre-Training Time from 3 Days to 76 Minutes:
+        https://arxiv.org/abs/1904.00962
+    """
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-4,
+                 weight_decay=0, adam=False):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(lr=lr, betas=betas, eps=eps,
+                        weight_decay=weight_decay)
+        self.adam = adam
+        super(Lamb, self).__init__(params, defaults)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('Lamb does not support sparse gradients, consider SparseAdam instad.')
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+
+                if group['weight_decay'] != 0:
+                    grad.add_(group['weight_decay'], p.data)
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+                denom = exp_avg_sq.sqrt().add_(group['eps'])
+
+                bias_correction1 = 1 - beta1 ** state['step']
+                bias_correction2 = 1 - beta2 ** state['step']
+                # Apply bias to lr to avoid broadcast.
+                step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1
+
+                adam_step = exp_avg / denom
+                # L2 norm uses sum, but here since we're dividing, use mean to avoid overflow.
+                r1 = p.data.pow(2).mean().sqrt()
+                r2 = adam_step.pow(2).mean().sqrt()
+                r = 1 if r1 == 0 or r2 == 0 else  min(r1/r2, 10)
+                state['r1'] = r1
+                state['r2'] = r2
+                state['r'] = r
+                if self.adam:
+                    r = 1
+
+                p.data.add_(-step_size * r, adam_step)
+
+        return loss

utils/musicautobot/utils/midifile.py

@@ -0,0 +1,107 @@
+"Transform functions for raw midi files"
+from enum import Enum
+import music21
+
+PIANO_TYPES = list(range(24)) + list(range(80, 96)) # Piano, Synths
+PLUCK_TYPES = list(range(24, 40)) + list(range(104, 112)) # Guitar, Bass, Ethnic
+BRIGHT_TYPES = list(range(40, 56)) + list(range(56, 80))
+
+PIANO_RANGE = (21, 109) # https://en.wikipedia.org/wiki/Scientific_pitch_notation
+
+class Track(Enum):
+    PIANO = 0 # discrete instruments - keyboard, woodwinds
+    PLUCK = 1 # continuous instruments with pitch bend: violin, trombone, synths
+    BRIGHT = 2
+    PERC = 3
+    UNDEF = 4
+    
+type2inst = {
+    # use print_music21_instruments() to see supported types
+    Track.PIANO: 0, # Piano
+    Track.PLUCK: 24, # Guitar
+    Track.BRIGHT: 40, # Violin
+    Track.PERC: 114, # Steel Drum
+}
+
+# INFO_TYPES = set(['TIME_SIGNATURE', 'KEY_SIGNATURE'])
+INFO_TYPES = set(['TIME_SIGNATURE', 'KEY_SIGNATURE', 'SET_TEMPO'])
+
+def file2mf(fp):
+    mf = music21.midi.MidiFile()
+    if isinstance(fp, bytes):
+        mf.readstr(fp)
+    else:
+        mf.open(fp)
+        mf.read()
+        mf.close()
+    return mf
+
+def mf2stream(mf): return music21.midi.translate.midiFileToStream(mf)
+
+def is_empty_midi(fp):
+    if fp is None: return False
+    mf = file2mf(fp)
+    return not any([t.hasNotes() for t in mf.tracks])
+
+def num_piano_tracks(fp):
+    music_file = file2mf(fp)
+    note_tracks = [t for t in music_file.tracks if t.hasNotes() and get_track_type(t) == Track.PIANO]
+    return len(note_tracks)
+
+def is_channel(t, c_val):
+    return any([c == c_val for c in t.getChannels()])
+
+def track_sort(t): # sort by 1. variation of pitch, 2. number of notes
+    return len(unique_track_notes(t)), len(t.events)
+
+def is_piano_note(pitch):
+    return (pitch >= PIANO_RANGE[0]) and (pitch < PIANO_RANGE[1])
+
+def unique_track_notes(t):
+    return { e.pitch for e in t.events if e.pitch is not None }
+
+def compress_midi_file(fp, cutoff=6, min_variation=3, supported_types=set([Track.PIANO, Track.PLUCK, Track.BRIGHT])):
+    music_file = file2mf(fp)
+    
+    info_tracks = [t for t in music_file.tracks if not t.hasNotes()]
+    note_tracks = [t for t in music_file.tracks if t.hasNotes()]
+    
+    if len(note_tracks) > cutoff:
+        note_tracks = sorted(note_tracks, key=track_sort, reverse=True)
+        
+    supported_tracks = []
+    for idx,t in enumerate(note_tracks):
+        if len(supported_tracks) >= cutoff: break
+        track_type = get_track_type(t)
+        if track_type not in supported_types: continue
+        pitch_set = unique_track_notes(t)
+        if (len(pitch_set) < min_variation): continue # must have more than x unique notes
+        if not all(map(is_piano_note, pitch_set)): continue # must not contain midi notes outside of piano range
+#         if track_type == Track.UNDEF: print('Could not designate track:', fp, t)
+        change_track_instrument(t, type2inst[track_type])
+        supported_tracks.append(t)
+    if not supported_tracks: return None
+    music_file.tracks = info_tracks + supported_tracks
+    return music_file
+
+def get_track_type(t):
+    if is_channel(t, 10): return Track.PERC
+    i = get_track_instrument(t)
+    if i in PIANO_TYPES: return Track.PIANO
+    if i in PLUCK_TYPES: return Track.PLUCK
+    if i in BRIGHT_TYPES: return Track.BRIGHT
+    return Track.UNDEF
+
+def get_track_instrument(t):
+    for idx,e in enumerate(t.events):
+        if e.type == 'PROGRAM_CHANGE': return e.data
+    return None
+
+def change_track_instrument(t, value):
+    for idx,e in enumerate(t.events):
+        if e.type == 'PROGRAM_CHANGE': e.data = value
+
+def print_music21_instruments():
+    for i in range(200):
+        try: print(i, music21.instrument.instrumentFromMidiProgram(i))
+        except: pass

utils/musicautobot/utils/setup_musescore.py

@@ -0,0 +1,46 @@
+def setup_musescore(musescore_path=None):
+    if not is_ipython(): return
+    
+    import platform
+    from music21 import environment
+    from pathlib import Path
+    
+    system = platform.system()
+    if system == 'Linux':
+        import os
+        os.environ['QT_QPA_PLATFORM']='offscreen' # https://musescore.org/en/node/29041
+        
+    existing_path = environment.get('musicxmlPath')
+    if existing_path: return
+    if musescore_path is None:
+        if system == 'Darwin':
+            app_paths = list(Path('/Applications').glob('MuseScore *.app'))
+            if len(app_paths): musescore_path = app_paths[-1]/'Contents/MacOS/mscore'
+        elif system == 'Linux':
+            musescore_path = '/usr/bin/musescore'
+    
+    if musescore_path is None or not Path(musescore_path).exists():
+        print('Warning: Could not find musescore installation. Please install musescore (see README) and/or update music21 environment paths')
+    else :
+        environment.set('musicxmlPath', musescore_path)
+        environment.set('musescoreDirectPNGPath', musescore_path)
+
+def is_ipython():
+    try: get_ipython
+    except: return False
+    return True
+
+def is_colab():
+    try: import google.colab
+    except: return False
+    return True
+
+def setup_fluidsynth():
+    from midi2audio import FluidSynth
+    from IPython.display import Audio
+
+def play_wav(stream):
+    out_midi = stream.write('midi')
+    out_wav = str(Path(out_midi).with_suffix('.wav'))
+    FluidSynth("font.sf2").midi_to_audio(out_midi, out_wav)
+    return Audio(out_wav)

utils/musicautobot/utils/stacked_dataloader.py

@@ -0,0 +1,70 @@
+"Dataloader wrapper that can combine and handle multiple dataloaders for multitask training"
+from fastai.callback import Callback
+from typing import Callable
+
+__all__ = ['StackedDataBunch']
+
+# DataLoading
+class StackedDataBunch():
+    def __init__(self, dbs, num_it=100):
+        self.dbs = dbs
+        self.train_dl = StackedDataloader([db.train_dl for db in self.dbs], num_it)
+        self.valid_dl = StackedDataloader([db.valid_dl for db in self.dbs], num_it)
+        self.train_ds = None
+        self.path = dbs[0].path
+        self.device = dbs[0].device
+        self.vocab = dbs[0].vocab
+        self.empty_val = False
+
+    def add_tfm(self,tfm:Callable)->None:
+        for dl in self.dbs: dl.add_tfm(tfm)
+
+    def remove_tfm(self,tfm:Callable)->None:
+        for dl in self.dbs: dl.remove_tfm(tfm)
+
+# Helper functions
+class StackedDataset(Callback):
+    def __init__(self, dss):
+        self.dss = dss
+    def __getattribute__(self, attr):
+        if attr == 'dss': return super().__getattribute__(attr)
+        def redirected(*args, **kwargs):
+            for ds in self.dss:
+                if hasattr(ds, attr): getattr(ds, attr)(*args, **kwargs)
+        return redirected
+    def __len__(self)->int: return sum([len(ds) for ds in self.dss])
+    def __repr__(self): return '\n'.join([self.__class__.__name__] + [repr(ds) for ds in self.dss])
+
+class StackedDataloader():
+    def __init__(self, dls, num_it=100):
+        self.dls = dls
+        self.dataset = StackedDataset([dl.dataset for dl in dls if hasattr(dl, 'dataset')])
+        self.num_it = num_it
+        self.dl_idx = -1
+        
+    def __len__(self)->int: return sum([len(dl) for dl in self.dls])
+    def __getattr__(self, attr):
+        def redirected(*args, **kwargs):
+            for dl in self.dls:
+                if hasattr(dl, attr):
+                    getattr(dl, attr)(*args, **kwargs)
+        return redirected
+        
+    def __iter__(self):
+        "Process and returns items from `DataLoader`."
+        iters = [iter(dl) for dl in self.dls]
+        self.dl_idx = -1
+        while len(iters):
+            self.dl_idx = (self.dl_idx+1) % len(iters)
+            for b in range(self.num_it):
+                try:
+                    yield next(iters[self.dl_idx])
+                except StopIteration as e:
+                    iters.remove(iters[self.dl_idx])
+                    break
+#         raise StopIteration
+
+    def new(self, **kwargs):
+        "Create a new copy of `self` with `kwargs` replacing current values."
+        new_dls = [dl.new(**kwargs) for dl in self.dls]
+        return StackedDataloader(new_dls, self.num_it)

utils/musicautobot/utils/top_k_top_p.py

@@ -0,0 +1,35 @@
+import torch
+import torch.nn.functional as F
+
+__all__ = ['top_k_top_p']
+
+# top_k + nucleus filter - https://twitter.com/thom_wolf/status/1124263861727760384?lang=en
+# https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+def top_k_top_p(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):
+    """ Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+        Args:
+            logits: logits distribution shape (vocabulary size)
+            top_k >0: keep only top k tokens with highest probability (top-k filtering).
+            top_p >0.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+    """
+    logits = logits.clone()
+    assert logits.dim() == 1  # batch size 1 for now - could be updated for more but the code would be less clear
+    top_k = min(top_k, logits.size(-1))  # Safety check
+    if top_k > 0:
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p > 0.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold
+        sorted_indices_to_remove = cumulative_probs > top_p
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        indices_to_remove = sorted_indices[sorted_indices_to_remove]
+        logits[indices_to_remove] = filter_value
+    return logits

utils/musicautobot/vocab.py

@@ -0,0 +1,93 @@
+from fastai.basics import *
+from .numpy_encode import *
+from .music_transformer import transform
+
+BOS = 'xxbos'
+PAD = 'xxpad'
+EOS = 'xxeos'
+MASK = 'xxmask' # Used for BERT masked language modeling. 
+CSEQ = 'xxcseq' # Used for Seq2Seq translation - denotes start of chord sequence
+MSEQ = 'xxmseq' # Used for Seq2Seq translation - denotes start of melody sequence
+
+# Deprecated tokens. Kept for compatibility 
+S2SCLS = 'xxs2scls' # deprecated
+NSCLS = 'xxnscls' # deprecated
+
+SEP = 'xxsep' # Used to denote end of timestep (required for polyphony). separator idx = -1 (part of notes)
+
+SPECIAL_TOKS = [BOS, PAD, EOS, S2SCLS, MASK, CSEQ, MSEQ, NSCLS, SEP] # Important: SEP token must be last
+
+NOTE_TOKS = [f'n{i}' for i in range(NOTE_SIZE)] 
+DUR_TOKS = [f'd{i}' for i in range(DUR_SIZE)]
+NOTE_START, NOTE_END = NOTE_TOKS[0], NOTE_TOKS[-1]
+DUR_START, DUR_END = DUR_TOKS[0], DUR_TOKS[-1]
+
+MTEMPO_SIZE = 10
+MTEMPO_OFF = 'mt0'
+MTEMPO_TOKS = [f'mt{i}' for i in range(MTEMPO_SIZE)]
+
+# Vocab - token to index mapping
+class MusicVocab():
+    "Contain the correspondence between numbers and tokens and numericalize."
+    def __init__(self, itos:Collection[str]):
+        self.itos = itos
+        self.stoi = {v:k for k,v in enumerate(self.itos)}
+
+    def numericalize(self, t:Collection[str]) -> List[int]:
+        "Convert a list of tokens `t` to their ids."
+        return [self.stoi[w] for w in t]
+
+    def textify(self, nums:Collection[int], sep=' ') -> List[str]:
+        "Convert a list of `nums` to their tokens."
+        items = [self.itos[i] for i in nums]
+        return sep.join(items) if sep is not None else items
+    
+    def to_music_item(self, idxenc):
+        return transform.MusicItem(idxenc, self)
+    
+    @property 
+    def mask_idx(self): return self.stoi[MASK]
+    @property 
+    def pad_idx(self): return self.stoi[PAD]
+    @property
+    def bos_idx(self): return self.stoi[BOS]
+    @property
+    def sep_idx(self): return self.stoi[SEP]
+    @property
+    def npenc_range(self): return (self.stoi[SEP], self.stoi[DUR_END]+1)
+    @property
+    def note_range(self): return self.stoi[NOTE_START], self.stoi[NOTE_END]+1
+    @property
+    def dur_range(self): return self.stoi[DUR_START], self.stoi[DUR_END]+1
+
+    def is_duration(self, idx): 
+        return idx >= self.dur_range[0] and idx < self.dur_range[1]
+    def is_duration_or_pad(self, idx):
+        return idx == self.pad_idx or self.is_duration(idx)
+        
+    def __getstate__(self):
+        return {'itos':self.itos}
+
+    def __setstate__(self, state:dict):
+        self.itos = state['itos']
+        self.stoi = {v:k for k,v in enumerate(self.itos)}
+        
+    def __len__(self): return len(self.itos)
+
+    def save(self, path):
+        "Save `self.itos` in `path`"
+        pickle.dump(self.itos, open(path, 'wb'))
+
+    @classmethod
+    def create(cls) -> 'Vocab':
+        "Create a vocabulary from a set of `tokens`."
+        itos = SPECIAL_TOKS + NOTE_TOKS + DUR_TOKS + MTEMPO_TOKS
+        if len(itos)%8 != 0:
+            itos = itos + [f'dummy{i}' for i in range(len(itos)%8)]
+        return cls(itos)
+    
+    @classmethod
+    def load(cls, path):
+        "Load the `Vocab` contained in `path`"
+        itos = pickle.load(open(path, 'rb'))
+        return cls(itos)