Upload inference.py

inference.py

@@ -42,15 +42,16 @@ byte_config = GPT2Config(num_hidden_layers=CHAR_NUM_LAYERS,
 
 model = NotaGenLMHeadModel(encoder_config=patch_config, decoder_config=byte_config).to(device)
 
+
 def download_model_weights():
     weights_path = "weights_notagenx_p_size_16_p_length_1024_p_layers_20_h_size_1280.pth"
     local_weights_path = os.path.join(os.getcwd(), weights_path)
-    
+
     # Check if weights already exist locally
     if os.path.exists(local_weights_path):
         logger.info(f"Model weights already exist at {local_weights_path}")
         return local_weights_path
-    
+
     logger.info("Downloading model weights from HuggingFace Hub...")
     try:
         # Download from HuggingFace
@@ -92,7 +93,7 @@ def prepare_model_for_kbit_training(model, use_gradient_checkpointing=True):
 
 model = prepare_model_for_kbit_training(
     model,
-    use_gradient_checkpointing=False  
+    use_gradient_checkpointing=False
 )
 
 print("Parameter Number: " + str(sum(p.numel() for p in model.parameters() if p.requires_grad)))
@@ -107,13 +108,12 @@ model.eval()
 
 
 def postprocess_inst_names(abc_text):
-
     with open('standard_inst_names.txt', 'r', encoding='utf-8') as f:
         standard_instruments_list = [line.strip() for line in f if line.strip()]
 
     with open('instrument_mapping.json', 'r', encoding='utf-8') as f:
         instrument_mapping = json.load(f)
-    
+
     abc_lines = abc_text.split('\n')
     abc_lines = list(filter(None, abc_lines))
     abc_lines = [line + '\n' for line in abc_lines]
@@ -123,20 +123,20 @@ def postprocess_inst_names(abc_text):
             match = re.search(r'nm="([^"]*)"', line)
             if match:
                 inst_name = match.group(1)
-                
+
                 # Check if the instrument name is already standard
                 if inst_name in standard_instruments_list:
                     continue
-                
+
                 # Find the most similar key in instrument_mapping
                 matching_key = difflib.get_close_matches(inst_name, list(instrument_mapping.keys()), n=1, cutoff=0.6)
-                
+
                 if matching_key:
                     # Replace the instrument name with the standardized version
                     replacement = instrument_mapping[matching_key[0]]
                     new_line = line.replace(f'nm="{inst_name}"', f'nm="{replacement}"')
                     abc_lines[i] = new_line
-    
+
     # Combine the lines back into a single string
     processed_abc_text = ''.join(abc_lines)
     return processed_abc_text
@@ -145,7 +145,7 @@ def postprocess_inst_names(abc_text):
 def complete_brackets(s):
     stack = []
     bracket_map = {'{': '}', '[': ']', '(': ')'}
-    
+
     # Iterate through each character, handle bracket matching
     for char in s:
         if char in bracket_map:
@@ -157,15 +157,13 @@ def complete_brackets(s):
                     if stack and stack[-1] == key:
                         stack.pop()
                     break  # Found matching right bracket, process next character
-    
+
     # Complete missing right brackets (in reverse order of remaining left brackets in stack)
     completion = ''.join(bracket_map[c] for c in reversed(stack))
     return s + completion
 
 
-
 def rest_unreduce(abc_lines):
-
     tunebody_index = None
     for i in range(len(abc_lines)):
         if abc_lines[i].startswith('%%score'):
@@ -215,7 +213,7 @@ def rest_unreduce(abc_lines):
             line_bar_dict[key] = value
 
         # calculate duration and collect barline
-        dur_dict = {}  
+        dur_dict = {}
         for symbol, bartext in line_bar_dict.items():
             right_barline = ''.join(re.split(Barline_regexPattern, bartext)[-2:])
             bartext = bartext[:-len(right_barline)]
@@ -232,7 +230,7 @@ def rest_unreduce(abc_lines):
         try:
             ref_dur = max(dur_dict, key=dur_dict.get)
         except:
-            pass    # use last ref_dur
+            pass  # use last ref_dur
 
         if i == 0:
             prefix_left_barline = line.split('[V:')[0]
@@ -256,16 +254,11 @@ def rest_unreduce(abc_lines):
     return unreduced_lines
 
 
-
-
-
-
 def inference_patch(period, composer, instrumentation):
-
-    prompt_lines=[
-    '%' + period + '\n',
-    '%' + composer + '\n',
-    '%' + instrumentation + '\n']
+    prompt_lines = [
+        '%' + period + '\n',
+        '%' + composer + '\n',
+        '%' + instrumentation + '\n']
 
     while True:
 
@@ -294,21 +287,22 @@ def inference_patch(period, composer, instrumentation):
         tunebody_flag = False
 
         with torch.inference_mode():
-            
+
             while True:
                 with torch.autocast(device_type='cuda', dtype=torch.float16):
                     predicted_patch = model.generate(input_patches.unsqueeze(0),
-                                                    top_k=TOP_K,
-                                                    top_p=TOP_P,
-                                                    temperature=TEMPERATURE)
-                if not tunebody_flag and patchilizer.decode([predicted_patch]).startswith('[r:'):  # 初次进入tunebody，必须以[r:0/开头
+                                                     top_k=TOP_K,
+                                                     top_p=TOP_P,
+                                                     temperature=TEMPERATURE)
+                if not tunebody_flag and patchilizer.decode([predicted_patch]).startswith(
+                        '[r:'):  # 初次进入tunebody，必须以[r:0/开头
                     tunebody_flag = True
                     r0_patch = torch.tensor([ord(c) for c in '[r:0/']).unsqueeze(0).to(device)
                     temp_input_patches = torch.concat([input_patches, r0_patch], axis=-1)
                     predicted_patch = model.generate(temp_input_patches.unsqueeze(0),
-                                                    top_k=TOP_K,
-                                                    top_p=TOP_P,
-                                                    temperature=TEMPERATURE)
+                                                     top_k=TOP_K,
+                                                     top_p=TOP_P,
+                                                     temperature=TEMPERATURE)
                     predicted_patch = [ord(c) for c in '[r:0/'] + predicted_patch
                 if predicted_patch[0] == patchilizer.bos_token_id and predicted_patch[1] == patchilizer.eos_token_id:
                     end_flag = True
@@ -336,7 +330,7 @@ def inference_patch(period, composer, instrumentation):
                 if len(byte_list) > 102400:
                     failure_flag = True
                     break
-                if time.time() - start_time > 10 * 60: 
+                if time.time() - start_time > 10 * 60:
                     failure_flag = True
                     break
 
@@ -347,16 +341,19 @@ def inference_patch(period, composer, instrumentation):
                     context_tunebody = ''.join(context_tunebody_byte_list)
 
                     if '\n' not in context_tunebody:
-                        break   # Generated content is all metadata, abandon
+                        break  # Generated content is all metadata, abandon
+
+                    context_tunebody_lines = context_tunebody.strip().split('\n')
 
-                    context_tunebody_liness = context_tunebody.split('\n')
                     if not context_tunebody.endswith('\n'):
-                        context_tunebody_liness = [context_tunebody_liness[i] + '\n' for i in range(len(context_tunebody_liness) - 1)] + [context_tunebody_liness[-1]]
+                        context_tunebody_lines = [context_tunebody_lines[i] + '\n' for i in
+                                                  range(len(context_tunebody_lines) - 1)] + [context_tunebody_lines[-1]]
                     else:
-                        context_tunebody_liness = [context_tunebody_liness[i] + '\n' for i in range(len(context_tunebody_liness))]
+                        context_tunebody_lines = [context_tunebody_lines[i] + '\n' for i in
+                                                  range(len(context_tunebody_lines))]
 
-                    cut_index = len(context_tunebody_liness) // 2
-                    abc_code_slice = metadata + ''.join(context_tunebody_liness[-cut_index:])
+                    cut_index = len(context_tunebody_lines) // 2
+                    abc_code_slice = metadata + ''.join(context_tunebody_lines[-cut_index:])
 
                     input_patches = patchilizer.encode_generate(abc_code_slice)
 
@@ -379,15 +376,13 @@ def inference_patch(period, composer, instrumentation):
                     failure_flag = True
                     pass
                 else:
-                    unreduced_abc_lines = [line for line in unreduced_abc_lines if not(line.startswith('%') and not line.startswith('%%'))]
+                    unreduced_abc_lines = [line for line in unreduced_abc_lines if
+                                           not (line.startswith('%') and not line.startswith('%%'))]
                     unreduced_abc_lines = ['X:1\n'] + unreduced_abc_lines
                     unreduced_abc_text = ''.join(unreduced_abc_lines)
                     return unreduced_abc_text
 
 
-
-
 if __name__ == '__main__':
-
     inference_patch('Classical', 'Beethoven, Ludwig van', 'Orchestral')