Update app.py

app.py

@@ -10,17 +10,6 @@ from samplings import top_p_sampling, top_k_sampling, temperature_sampling
 from transformers import GPT2Config, 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>
@@ -61,288 +50,72 @@ A general recommendation is to adjust the desired musical structure and form thr
 Please make sure to operate according to the provided formats and guidelines to fully leverage the capabilities of TunesFormer and achieve a satisfying music generation experience.
 """
 
-class Patchilizer:
-    """
-    A class for converting music bars to patches and vice versa. 
-    """
+class ABCTokenizer():
     def __init__(self):
-        self.delimiters = ["|:", "::", ":|", "[|", "||", "|]", "|"]
-        self.regexPattern = '(' + '|'.join(map(re.escape, self.delimiters)) + ')'
         self.pad_token_id = 0
-        self.bos_token_id = 1
-        self.eos_token_id = 2
+        self.bos_token_id = 2
+        self.eos_token_id = 3
+        self.merged_tokens = []
 
-    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 __len__(self):
+        return 128+len(self.merged_tokens)
     
-    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 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)
-        
-        # 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)
-
-        return probs
-
-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.
-    """
-
-    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)
-
-            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
-
-        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)
-
-        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)
+    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!=2 and i!=3:
+                txt += chr(i)
+        return txt
+
+    def txt2ids(self, text, merged_tokens):
+        text = unidecode(text)
+        ids = [str(ord(c)) for c in text]
+        if torch.max(torch.tensor([ord(c) for c in text]))>=128:
+            return [128+len(self.merged_tokens)]
+        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))
         
-        return generated_patch, n_seed
+        txt_ids = txt_ids.split(' ')
+        txt_ids = [int(i) for i in txt_ids]
+        return [self.bos_token_id]+txt_ids+[self.eos_token_id]
 
 def generate_abc(prompt,
                  num_tunes,
-                 max_patch,
+                 max_length,
                  top_p,
                  top_k,
                  temperature,
-                 seed,
-                 show_control_code=True):
+                 seed):
 
     if torch.cuda.is_available():    
         device = torch.device("cuda")
     else:
         device = torch.device("cpu")
 
-    patchilizer = Patchilizer()
+    tokenizer = ABCTokenizer()
+    config = GPT2Config(vocab_size=len(tokenizer))
+    model = GPT2LMHeadModel(config).to(device)
 
-    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"
+    filename = "pytorch_model.bin"
 
     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'
+            url = 'https://huggingface.co/sander-wood/tunesformer/resolve/main/pytorch_model.bin'
             response = requests.get(url, stream=True)
 
             total_size = int(response.headers.get('content-length', 0))
@@ -362,59 +135,43 @@ def generate_abc(prompt,
             print(f"Error: {e}")
             exit()
             
-    checkpoint = torch.load('weights.pth')
-    model.load_state_dict(checkpoint['model'])
-    model = model.to(device)
+    model.load_state_dict(torch.load('pytorch_model.bin'))
     model.eval()
 
     tunes = ""
+        
+    if prompt:
+        ids = tokenizer.encode(prompt)['input_ids'][:-1]
+    else:
+        ids = torch.tensor([tokenizer.bos_token_id])
 
-    print("\n"+" OUTPUT TUNES ".center(60, "#"))
+    random.seed(seed)
 
+    print("\n"+" OUTPUT TUNES ".center(60, "#"))
+    
     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
+        print(tune, end="")
+        input_ids = ids.unsqueeze(0)
+        for t_idx in range(max_length):
+            if seed!=None:
+                n_seed = random.randint(0, 1000000)
+                random.seed(n_seed)
             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)
+                n_seed = None
+            outputs = model(input_ids=input_ids.to(device))
+            probs = outputs.logits[0][-1]
+            probs = torch.nn.Softmax(dim=-1)(probs).cpu().detach().numpy()
+            probs = top_p_sampling(probs, top_p=top_p, return_probs=True)
+            probs = top_k_sampling(probs, top_k=top_k, return_probs=True)
+            sampled_id = temperature_sampling(probs, temperature=temperature, seed=n_seed)
+            input_ids = torch.cat((input_ids, torch.tensor([[sampled_id]])), 1)
+            if sampled_id!=3:
+                tune += tokenizer.decode([sampled_id], skip_special_tokens=True)
+                print(tune[-1], end="")
+                continue
             else:
                 break
-
         tunes += tune+"\n\n"
     
     return tunes
@@ -430,7 +187,7 @@ K:D
 
 input_prompt = gr.inputs.Textbox(lines=5, label="Prompt", default=default_prompt)
 input_num_tunes = gr.inputs.Slider(minimum=1, maximum=10, step=1, default=1, label="Number of Tunes")
-input_max_patch = gr.inputs.Slider(minimum=10, maximum=128, step=1, default=128, label="Max Patch")
+input_max_length = gr.inputs.Slider(minimum=64, maximum=1024, step=1, default=1024, label="Max Length")
 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")
@@ -438,7 +195,7 @@ 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_prompt, input_num_tunes, input_max_patch, input_top_p, input_top_k, input_temperature, input_seed],
+                [input_prompt, input_num_tunes, input_max_length, input_top_p, input_top_k, input_temperature, input_seed],
                 output_abc,
             title="TunesFormer: Forming Irish Tunes with Control Codes by Bar Patching",
             description=description).launch()