Update app.py

app.py

@@ -1,11 +1,12 @@
-#=========================================================================
-# https://huggingface.co/spaces/asigalov61/Parsons-Code-Melody-Transformer
-#=========================================================================
+#==================================================================
+# https://huggingface.co/spaces/asigalov61/Popular-Hook-Transformer
+#==================================================================
 
 import time as reqtime
 import datetime
 from pytz import timezone
 
+import statistics
 import re
 import tqdm
 
@@ -21,29 +22,113 @@ import matplotlib.pyplot as plt
     
 #=====================================================================================
 
-def parsons_code_to_tokens(parsons_code_str):
+print('=' * 70)
+print('Popular Hook Transformer')
+print('=' * 70)
 
-  tokens = [388]
+print('Loading Popular Hook Transformer training data...')
 
-  for chr in parsons_code_str[1:]:
+#====================================================================================
+
+SEQ_LEN = 512
+PAD_IDX = 918
+DEVICE = 'cpu'
+
+#====================================================================================
+
+def str_strip(string):
+    return re.sub(r'[^A-Za-z-]+', '', string).rstrip('-')
+
+def mode_time(seq):
+    return statistics.mode([t for t in seq if 0 < t < 128])
+
+def mode_dur(seq):
+    return statistics.mode([t-128 for t in seq if 128 < t < 256])
+
+def mode_pitch(seq):
+    return statistics.mode([t % 128 for t in seq if 256 < t < 512])
+
+parts_dict = sorted(set([str_strip(s[2]).rstrip('-') for s in melody_chords_f]))
+
+train_data = []
+
+for m in tqdm.tqdm(melody_chords_f):
+    
+    if 64 < len(m[5]) < 506:
+        
+        for tv in range(-3, 3):
+        
+            part = str_strip(m[2])
+            part_tok = parts_dict.index(part)
+            
+            score = [t+tv if 256 < t < 512 else t for t in m[5]]
+
+            seq = [916] + [part_tok+512, mode_time(score)+532, mode_dur(score)+660, mode_pitch(score)+tv+788]
+
+            seq += score
+
+            seq += [917]
+
+            seq = seq + [PAD_IDX] * (SEQ_LEN - len(seq))
+
+            train_data.append(seq)
+
+#====================================================================================
+
+print('Done!')
+print('=' * 70)
+print('All data is good:', len(max(train_data, key=len)) == len(min(train_data, key=len)))
+print('=' * 70)
+print('Randomizing training data...')
+random.shuffle(train_data)
+print('Done!')
+print('=' * 70)
+print('Total length of training data:', len(train_data))
+print('=' * 70)
+
+#====================================================================================
+
+print('Loading Popular Hook Transformer pre-trained model...')
+print('=' * 70)
+
+print('Instantiating model...')
+
+model = TransformerWrapper(
+    num_tokens = PAD_IDX+1,
+    max_seq_len = SEQ_LEN,
+    attn_layers = Decoder(dim = 1024, 
+                          depth = 4, 
+                          heads = 32,
+                          rotary_pos_emb = True,
+                          attn_flash = True
+                         )
+    )
+
+model = AutoregressiveWrapper(model, ignore_index = PAD_IDX, pad_value=PAD_IDX)
+
+print('=' * 70)
+print('Loading model checkpoint...')
+
+model_path = 'Popular_Hook_Transformer_Small_Trained_Model_10869_steps_0.2308_loss_0.9252_acc.pth'
+
+model.load_state_dict(torch.load(model_path))
 
-    if chr == 'D':
-      tokens.extend([385])
+print('=' * 70)
 
-    elif chr == 'R':
-      tokens.extend([386])
+model.to(DEVICE)
+model.eval()
 
-    elif chr == 'U':
-      tokens.extend([387])
+ctx = torch.amp.autocast(device_type=DEVICE, dtype=torch.bfloat16)
 
-  return tokens
+print('Done!')
+print('=' * 70)
 
 #====================================================================================
 
-def Generate_Melody(input_parsons_code, 
-                   input_first_note_duration, 
-                   iinput_first_note_MIDI_pitch
-                 ):
+def Generate_POP_Section(input_parsons_code, 
+                           input_first_note_duration, 
+                           iinput_first_note_MIDI_pitch
+                         ):
     
     print('=' * 70)
     print('Req start time: {:%Y-%m-%d %H:%M:%S}'.format(datetime.datetime.now(PDT)))