Update app.py

app.py

@@ -29,8 +29,8 @@ def classify_GPU(input_data):
 
     print('Loading model...')
 
-    SEQ_LEN = 1024
-    PAD_IDX = 14627
+    SEQ_LEN = 1026
+    PAD_IDX = 940
     DEVICE = 'cuda' # 'cuda'
 
     # instantiate the model
@@ -38,18 +38,21 @@ def classify_GPU(input_data):
     model = TransformerWrapper(
         num_tokens = PAD_IDX+1,
         max_seq_len = SEQ_LEN,
-        attn_layers = Decoder(dim = 1024, depth = 12, heads = 16, attn_flash = True)
-        )
+        attn_layers = Decoder(dim = 1024, depth = 24, heads = 32, attn_flash = True)
+    )
+    
+    model = AutoregressiveWrapper(model, ignore_index=PAD_IDX, pad_value=PAD_IDX)
     
-    model = AutoregressiveWrapper(model, ignore_index = PAD_IDX)
+    model = torch.nn.DataParallel(model)
 
     model.to(DEVICE)
+    
     print('=' * 70)
 
     print('Loading model checkpoint...')
 
     model.load_state_dict(
-        torch.load('Annotated_MIDI_Dataset_Classifier_Trained_Model_21269_steps_0.4335_loss_0.8716_acc.pth',
+        torch.load('Ultimate_MIDI_Classifier_Trained_Model_29886_steps_0.556_loss_0.8339_acc.pth',
                    map_location=DEVICE))
     print('=' * 70)
 
@@ -67,12 +70,8 @@ def classify_GPU(input_data):
 
     #==================================================================
 
-    number_of_batches = 1 # @param {type:"slider", min:1, max:100, step:1}
-    
-    # @markdown NOTE: You can increase the number of batches on high-ram GPUs for better classification
-    
     print('=' * 70)
-    print('Annotated MIDI Dataset Classifier')
+    print('Ultimate MIDI Classifier')
     print('=' * 70)
     print('Classifying...')
     
@@ -80,29 +79,20 @@ def classify_GPU(input_data):
     
     model.eval()
     
-    results = []
-    
-    for input in input_data:
-    
-        x = torch.tensor([input[:1022]] * number_of_batches, dtype=torch.long, device='cuda')
-    
-        with ctx:
-          out = model.generate(x,
-                              1,
-                              temperature=0.3,
-                              filter_logits_fn=top_k,
-                              filter_kwargs={'k': 1},
-                              return_prime=False,
-                              verbose=False)
-    
-        y = out.tolist()
-    
-        output = [l[0] for l in y]
-        result = mode(output)
-    
-        results.append(result)
+    x = torch.tensor(input_data[:1022], dtype=torch.long, device=DEVICE)
+
+    with ctx:
+      out = model.module.generate(x,
+                                  2,
+                                  filter_logits_fn=top_k,
+                                  filter_kwargs={'k': 1},
+                                  temperature=0.9,
+                                  return_prime=False,
+                                  verbose=False)
 
-    return output, results
+    result = tuple(out[0].tolist())
+    
+    return result
                        
 # =================================================================================================