Update app.py

app.py

@@ -2,10 +2,22 @@ import gradio as gr
 import torch
 import random
 from unidecode import unidecode
-from transformers import GPT2LMHeadModel
-from samplings import top_p_sampling, temperature_sampling
+import re
+from samplings import top_p_sampling, top_k_sampling, temperature_sampling
+from transformers import GPT2Model, GPT2LMHeadModel, PreTrainedModel
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+PATCH_LENGTH = 128      # Patch Length
+PATCH_SIZE = 32       # Patch Size
+
+PATCH_NUM_LAYERS = 9         # Number of layers in the encoder
+CHAR_NUM_LAYERS = 3          # Number of layers in the decoder
+
+NUM_EPOCHS = 32                 # Number of epochs to train for (if early stopping doesn't intervene)
+LEARNING_RATE = 5e-5            # Learning rate for the optimizer
+PATCH_SAMPLING_BATCH_SIZE = 0   # Batch size for patch during training, 0 for full context
+LOAD_FROM_CHECKPOINT = False     # Whether to load weights from a checkpoint
+SHARE_WEIGHTS = False            # Whether to share weights between the encoder and decoder
 
 description = """
 <div>
@@ -45,114 +57,379 @@ The first control code `[SECS_3]` specifies there are 3 sections in the tune, an
 
 """
 
-
-class ABCTokenizer():
+class Patchilizer:
+    """
+    A class for converting music bars to patches and vice versa. 
+    """
     def __init__(self):
+        self.delimiters = ["|:", "::", ":|", "[|", "||", "|]", "|"]
+        self.regexPattern = '(' + '|'.join(map(re.escape, self.delimiters)) + ')'
         self.pad_token_id = 0
-        self.bos_token_id = 2
-        self.eos_token_id = 3
-        self.merged_tokens = []
-        
-        for i in range(8):
-            self.merged_tokens.append('[SECS_'+str(i+1)+']')
-        for i in range(32):
-            self.merged_tokens.append('[BARS_'+str(i+1)+']')
-        for i in range(11):
-            self.merged_tokens.append('[SIM_'+str(i)+']')
-
-    def __len__(self):
-        return 128+len(self.merged_tokens)
+        self.bos_token_id = 1
+        self.eos_token_id = 2
+
+    def split_bars(self, body):
+        """
+        Split a body of music into individual bars.
+        """
+        bars = re.split(self.regexPattern, ''.join(body))
+        bars = list(filter(None, bars))  # remove empty strings
+        if bars[0] in self.delimiters:
+            bars[1] = bars[0] + bars[1]
+            bars = bars[1:]
+        bars = [bars[i * 2] + bars[i * 2 + 1] for i in range(len(bars) // 2)]
+        return bars
+    
+    def bar2patch(self, bar, patch_size=PATCH_SIZE):
+        """
+        Convert a bar into a patch of specified length.
+        """
+        patch = [self.bos_token_id] + [ord(c) for c in bar] + [self.eos_token_id]
+        patch = patch[:patch_size]
+        patch += [self.pad_token_id] * (patch_size - len(patch))
+        return patch
     
-    def encode(self, text):
-        encodings = {}
-        encodings['input_ids'] = torch.tensor(self.txt2ids(text, self.merged_tokens))
-        encodings['attention_mask'] = torch.tensor([1]*len(encodings['input_ids']))
-        return encodings
-
-    def decode(self, ids, skip_special_tokens=False):
-        txt = ""
-        for i in ids:
-            if i>=128:
-                if not skip_special_tokens:
-                    txt += self.merged_tokens[i-128]
-            elif i!=self.bos_token_id and i!=self.eos_token_id:
-                txt += chr(i)
-        return txt
-
-    def txt2ids(self, text, merged_tokens):
-        ids = ["\""+str(ord(c))+"\"" for c in text]
-        txt_ids = ' '.join(ids)
-        for t_idx, token in enumerate(merged_tokens):
-            token_ids = ["\""+str(ord(c))+"\"" for c in token]
-            token_txt_ids = ' '.join(token_ids)
-            txt_ids = txt_ids.replace(token_txt_ids, "\""+str(t_idx+128)+"\"")
+    def patch2bar(self, patch):
+        """
+        Convert a patch into a bar.
+        """
+        return ''.join(chr(idx) if idx > self.eos_token_id else '' for idx in patch if idx != self.eos_token_id)
+
+    def encode(self, abc_code, patch_length=PATCH_LENGTH, patch_size=PATCH_SIZE, add_special_patches=False):
+        """
+        Encode music into patches of specified length.
+        """
+        lines = unidecode(abc_code).split('\n')
+        lines = list(filter(None, lines))  # remove empty lines
+
+        body = ""
+        patches = []
+
+        for line in lines:
+            if len(line) > 1 and ((line[0].isalpha() and line[1] == ':') or line.startswith('%%score')):
+                if body:
+                    bars = self.split_bars(body)
+                    patches.extend(self.bar2patch(bar + '\n' if idx == len(bars) - 1 else bar, patch_size) 
+                                   for idx, bar in enumerate(bars))
+                    body = ""
+                patches.append(self.bar2patch(line + '\n', patch_size))
+            else:
+                body += line + '\n'
+
+        if body:
+            patches.extend(self.bar2patch(bar, patch_size) for bar in self.split_bars(body))
+
+        if add_special_patches:
+            bos_patch = [self.bos_token_id] * (patch_size-1) + [self.eos_token_id]
+            eos_patch = [self.bos_token_id] + [self.eos_token_id] * (patch_size-1)
+            patches = [bos_patch] + patches + [eos_patch]
+
+        return patches[:patch_length]
+
+    def decode(self, patches):
+        """
+        Decode patches into music.
+        """
+        return ''.join(self.patch2bar(patch) for patch in patches)
+
+class PatchLevelDecoder(PreTrainedModel):
+    """
+    An Patch-level Decoder model for generating patch features in an auto-regressive manner. 
+    It inherits PreTrainedModel from transformers.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.patch_embedding = torch.nn.Linear(PATCH_SIZE * 128, config.n_embd)
+        torch.nn.init.normal_(self.patch_embedding.weight, std=0.02)
+        self.base = GPT2Model(config)
+
+    def forward(self, patches: torch.Tensor) -> torch.Tensor:
+        """
+        The forward pass of the patch-level decoder model.
+        :param patches: the patches to be encoded
+        :return: the encoded patches
+        """
+        patches = torch.nn.functional.one_hot(patches, num_classes=128).float()
+        patches = patches.reshape(len(patches), -1, PATCH_SIZE * 128)
+        patches = self.patch_embedding(patches.to(self.device))
+
+        return self.base(inputs_embeds=patches)
+
+class CharLevelDecoder(PreTrainedModel):
+    """
+    A Char-level Decoder model for generating the characters within each bar patch sequentially. 
+    It inherits PreTrainedModel from transformers.
+    """
+    def __init__(self, config):
+        super().__init__(config)
+        self.pad_token_id = 0
+        self.bos_token_id = 1
+        self.eos_token_id = 2
+        self.base = GPT2LMHeadModel(config)
+
+    def forward(self, encoded_patches: torch.Tensor, target_patches: torch.Tensor, patch_sampling_batch_size: int):
+        """
+        The forward pass of the char-level decoder model.
+        :param encoded_patches: the encoded patches
+        :param target_patches: the target patches
+        :return: the decoded patches
+        """
+        # preparing the labels for model training
+        target_masks = target_patches == self.pad_token_id
+        labels = target_patches.clone().masked_fill_(target_masks, -100)
+
+        # masking the labels for model training
+        target_masks = torch.ones_like(labels)
+        target_masks = target_masks.masked_fill_(labels == -100, 0)
+
+        # select patches
+        if patch_sampling_batch_size!=0 and patch_sampling_batch_size<target_patches.shape[0]:
+            indices = list(range(len(target_patches)))
+            random.shuffle(indices)
+            selected_indices = sorted(indices[:patch_sampling_batch_size])
+
+            target_patches = target_patches[selected_indices,:]
+            target_masks = target_masks[selected_indices,:]
+            encoded_patches = encoded_patches[selected_indices,:]
+            labels = labels[selected_indices,:]
+
+        # get input embeddings
+        inputs_embeds = torch.nn.functional.embedding(target_patches, self.base.transformer.wte.weight)
+
+        # concatenate the encoded patches with the input embeddings
+        inputs_embeds = torch.cat((encoded_patches.unsqueeze(1), inputs_embeds[:,1:,:]), dim=1)
+
+        return self.base(inputs_embeds=inputs_embeds,
+                         attention_mask=target_masks,
+                         labels=labels)
+
+    def generate(self, encoded_patch: torch.Tensor, tokens: torch.Tensor):
+        """
+        The generate function for generating a patch based on the encoded patch and already generated tokens.
+        :param encoded_patch: the encoded patch
+        :param tokens: already generated tokens in the patch
+        :return: the probability distribution of next token
+        """
+        encoded_patch = encoded_patch.reshape(1, 1, -1)
+        tokens = tokens.reshape(1, -1)
+
+        # Get input embeddings
+        tokens = torch.nn.functional.embedding(tokens, self.base.transformer.wte.weight)
+
+        # Concatenate the encoded patch with the input embeddings
+        tokens = torch.cat((encoded_patch, tokens[:,1:,:]), dim=1)
         
-        txt_ids = txt_ids.split(' ')
-        txt_ids = [int(i[1:-1]) for i in txt_ids]
-        return [self.bos_token_id]+txt_ids+[self.eos_token_id]
+        # Get output from model
+        outputs = self.base(inputs_embeds=tokens)
+        
+        # Get probabilities of next token
+        probs = torch.nn.functional.softmax(outputs.logits.squeeze(0)[-1], dim=-1)
 
-def generate_abc(control_codes, prefix, num_tunes, max_length, top_p, temperature, seed):
+        return probs
 
-    try:
-        seed = int(seed)
-    except:
-        seed = None
+class TunesFormer(PreTrainedModel):
+    """
+    TunesFormer is a hierarchical music generation model based on bar patching. 
+    It includes a patch-level decoder and a character-level decoder.
+    It inherits PreTrainedModel from transformers.
+    """
 
-    prefix = unidecode(control_codes + prefix)
-    tokenizer = ABCTokenizer()
-    model = GPT2LMHeadModel.from_pretrained('sander-wood/tunesformer').to(device)
+    def __init__(self, encoder_config, decoder_config, share_weights=False):
+        super().__init__(encoder_config)
+        self.pad_token_id = 0
+        self.bos_token_id = 1
+        self.eos_token_id = 2
+        if share_weights:
+            max_layers = max(encoder_config.num_hidden_layers, decoder_config.num_hidden_layers)
+            max_context_size = max(encoder_config.max_length, decoder_config.max_length)
+            max_position_embeddings = max(encoder_config.max_position_embeddings, decoder_config.max_position_embeddings)
 
-    if prefix:
-        ids = tokenizer.encode(prefix)['input_ids'][:-1]
-    else:
-        ids = torch.tensor([tokenizer.bos_token_id])
+            encoder_config.num_hidden_layers = max_layers
+            encoder_config.max_length = max_context_size
+            encoder_config.max_position_embeddings = max_position_embeddings
+            decoder_config.num_hidden_layers = max_layers
+            decoder_config.max_length = max_context_size
+            decoder_config.max_position_embeddings = max_position_embeddings
 
-    random.seed(seed)
-    tunes = ""
+        self.patch_level_decoder = PatchLevelDecoder(encoder_config)
+        self.char_level_decoder = CharLevelDecoder(decoder_config)
+
+        if share_weights:
+            self.patch_level_decoder.base = self.char_level_decoder.base.transformer
+
+    def forward(self, patches: torch.Tensor, patch_sampling_batch_size: int=PATCH_SAMPLING_BATCH_SIZE):
+        """
+        The forward pass of the TunesFormer model.
+        :param patches: the patches to be both encoded and decoded
+        :return: the decoded patches
+        """
+        patches = patches.reshape(len(patches), -1, PATCH_SIZE)
+        encoded_patches = self.patch_level_decoder(patches)["last_hidden_state"]
+        
+        return self.char_level_decoder(encoded_patches.squeeze(0)[:-1, :], patches.squeeze(0)[1:, :], patch_sampling_batch_size)
+
+    def generate(self, patches: torch.Tensor,
+                 tokens: torch.Tensor,
+                 top_p: float=1,
+                 top_k: int=0,
+                 temperature: float=1,
+                 seed: int=None):
+        """
+        The generate function for generating patches based on patches.
+        :param patches: the patches to be encoded
+        :return: the generated patches
+        """
+        patches = patches.reshape(len(patches), -1, PATCH_SIZE)
+        encoded_patches = self.patch_level_decoder(patches)["last_hidden_state"]
+        if tokens==None:
+            tokens = torch.tensor([self.bos_token_id], device=self.device)
+        generated_patch = []
+        random.seed(seed)
 
-    for c_idx in range(num_tunes):
-        print("\nX:"+str(c_idx+1)+"\n", end="")
-        print(tokenizer.decode(ids[1:], skip_special_tokens=True), end="")
-        input_ids = ids.unsqueeze(0)
-        for t_idx in range(max_length):
+        while True:
             if seed!=None:
                 n_seed = random.randint(0, 1000000)
                 random.seed(n_seed)
             else:
                 n_seed = None
+            prob = self.char_level_decoder.generate(encoded_patches[0][-1], tokens).cpu().detach().numpy()
+            prob = top_p_sampling(prob, top_p=top_p, return_probs=True)
+            prob = top_k_sampling(prob, top_k=top_k, return_probs=True)
+            token = temperature_sampling(prob, temperature=temperature, seed=n_seed)
+            generated_patch.append(token)
+            if token == self.eos_token_id or len(tokens) >= PATCH_SIZE - 1:
+                break
+            else:
+                tokens = torch.cat((tokens, torch.tensor([token], device=self.device)), dim=0)
+        
+        return generated_patch, n_seed
+
+def generate_abc(prompt, num_tunes, max_patch, top_p, top_k, temperature, seed, show_control_code):
+
+    if torch.cuda.is_available():    
+        device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+
+    patchilizer = Patchilizer()
 
-            outputs = model(input_ids=input_ids.to(device))
-            probs = outputs.logits[0][-1]
-            probs = torch.nn.Softmax(dim=-1)(probs).cpu().detach().numpy()
-            sampled_id = temperature_sampling(probs=top_p_sampling(probs, 
-                                                                top_p=top_p, 
-                                                                seed=n_seed,
-                                                                return_probs=True),
-                                            seed=n_seed,
-                                            temperature=temperature)
-            input_ids = torch.cat((input_ids, torch.tensor([[sampled_id]])), 1)
-            if sampled_id!=tokenizer.eos_token_id:
-                print(tokenizer.decode([sampled_id], skip_special_tokens=True), end="")
-                continue
+    patch_config = GPT2Config(num_hidden_layers=PATCH_NUM_LAYERS, 
+                        max_length=PATCH_LENGTH, 
+                        max_position_embeddings=PATCH_LENGTH,
+                        vocab_size=1)
+    char_config = GPT2Config(num_hidden_layers=CHAR_NUM_LAYERS, 
+                        max_length=PATCH_SIZE, 
+                        max_position_embeddings=PATCH_SIZE,
+                        vocab_size=128)
+    model = TunesFormer(patch_config, char_config, share_weights=SHARE_WEIGHTS)
+
+    filename = "weights.pth"
+
+    if os.path.exists(filename):
+        print(f"Weights already exist at '{filename}'. Loading...")
+    else:
+        print(f"Downloading weights to '{filename}' from huggingface.co...")
+        try:
+            url = 'https://huggingface.co/sander-wood/tunesformer/resolve/main/weights.pth'
+            response = requests.get(url, stream=True)
+
+            total_size = int(response.headers.get('content-length', 0))
+            chunk_size = 1024
+
+            with open(filename, 'wb') as file, tqdm(
+                desc=filename,
+                total=total_size,
+                unit='B',
+                unit_scale=True,
+                unit_divisor=1024,
+            ) as bar:
+                for data in response.iter_content(chunk_size=chunk_size):
+                    size = file.write(data)
+                    bar.update(size)
+        except Exception as e:
+            print(f"Error: {e}")
+            exit()
+            
+    checkpoint = torch.load('weights.pth')
+    model.load_state_dict(checkpoint['model'])
+    model = model.to(device)
+    model.eval()
+
+    tunes = ""
+
+    print("\n"+" OUTPUT TUNES ".center(60, "#"))
+
+    start_time = time.time()
+
+    for i in range(num_tunes):
+        tune = "X:"+str(i+1) + "\n" + prompt
+        lines = re.split(r'(\n)', tune)
+        tune = ""
+        skip = False
+        for line in lines:
+            if show_control_code or line[:2] not in ["S:", "B:", "E:"]:
+                if not skip:
+                    print(line, end="")
+                    tune += line
+                skip = False
+            else:
+                skip = True
+
+        input_patches = torch.tensor([patchilizer.encode(prompt, add_special_patches=True)[:-1]], device=device)
+        if tune=="":
+            tokens = None
+        else:
+            prefix = patchilizer.decode(input_patches[0])
+            remaining_tokens = prompt[len(prefix):]
+            tokens  = torch.tensor([patchilizer.bos_token_id]+[ord(c) for c in remaining_tokens], device=device)
+        
+        while input_patches.shape[1]<max_patch:
+            predicted_patch, seed = model.generate(input_patches,
+                                                    tokens,
+                                                    top_p=top_p,
+                                                    top_k=top_k,
+                                                    temperature=temperature,
+                                                    seed=seed)
+            tokens = None
+            if predicted_patch[0]!=patchilizer.eos_token_id:
+                next_bar = patchilizer.decode([predicted_patch])
+                if show_control_code or next_bar[:2] not in ["S:", "B:", "E:"]:
+                    print(next_bar, end="")
+                    tune += next_bar
+                if next_bar=="":
+                    break
+                next_bar = remaining_tokens+next_bar
+                remaining_tokens = ""
+                predicted_patch = torch.tensor(patchilizer.bar2patch(next_bar), device=device).unsqueeze(0)
+                input_patches = torch.cat([input_patches, predicted_patch.unsqueeze(0)], dim=1)
             else:
-                tune = "X:"+str(c_idx+1)+"\n"+tokenizer.decode(input_ids.squeeze(), skip_special_tokens=True)
-                tunes += tune+"\n\n"
-                print("\n")
                 break
+
+        tunes += tune+"\n\n"
     
     return tunes
 
-input_control_codes = gr.inputs.Textbox(lines=5, label="Control Codes", default="[SECS_2][BARS_9][SIM_3][BARS_9]")
-input_prefix = gr.inputs.Textbox(lines=5, label="Prefix", default="L:1/8\nQ:1/4=114\nM:3/4\nK:D\nde | \"D\"")
+default_prompt = """S:2
+B:9
+E:4
+B:9
+L:1/8
+M:3/4
+K:D
+ de |"D" """
+
+input_prompt = gr.inputs.Textbox(lines=5, label="Prefix", default=default_prompt)
 input_num_tunes = gr.inputs.Slider(minimum=1, maximum=10, step=1, default=1, label="Number of Tunes")
-input_max_length = gr.inputs.Slider(minimum=10, maximum=1000, step=10, default=500, label="Max Length")
-input_top_p = gr.inputs.Slider(minimum=0.0, maximum=1.0, step=0.05, default=0.9, label="Top P")
-input_temperature = gr.inputs.Slider(minimum=0.0, maximum=2.0, step=0.1, default=1.0, label="Temperature")
+input_max_patch = gr.inputs.Slider(minimum=10, maximum=128, step=1, default=128, label="Max Patch")
+input_top_p = gr.inputs.Slider(minimum=0.0, maximum=1.0, step=0.05, default=0.8, label="Top P")
+input_top_k = gr.inputs.Slider(minimum=1, maximum=20, step=1, default=8, label="Top K")
+input_temperature = gr.inputs.Slider(minimum=0.0, maximum=2.0, step=0.05, default=1.2, label="Temperature")
 input_seed = gr.inputs.Textbox(lines=1, label="Seed (int)", default="None")
 output_abc = gr.outputs.Textbox(label="Generated Tunes")
 
 gr.Interface(generate_abc,
-                [input_control_codes, input_prefix, input_num_tunes, input_max_length, input_top_p, input_temperature, input_seed],
+                [input_prompt, input_num_tunes, input_max_patch, input_top_p, input_top_k, input_temperature, input_seed],
                 output_abc,
-            title="TunesFormer: Forming Tunes with Control Codes",
+            title="TunesFormer: Forming Irish Tunes with Control Codes by Bar Patching",
             description=description).launch()