pitch_index argument to select pitch with nth greatest likelihood

Author: victor-shepardson

examples/notepredictor/autopitch.scd

@@ -92,7 +92,7 @@ MIDIdef.noteOn(\input, {
     // fixing dt and vel to performed values so just pitch is predicted
     b.sendMsg("/predictor/predict",
         \pitch, ~last_pitch, \time, ~last_dt, \vel, ~last_vel,
-        \fix_time, dt, \fix_vel, vel);
+        \index_pitch, pitch, \fix_time, dt, \fix_vel, vel);
 
     ~last_dt = dt;
     ~last_vel = vel;
@@ -122,12 +122,9 @@ OSCdef(\return, {
 
 ~reset.()
 // send a note manually if you don't have a midi controller
-MIDIdef.all[\input].func.(64)
+MIDIdef.all[\input].func.(64, 16) //velocity, "pitch"
 
-(
-MIDIdef.all[\input].func.(64);
-MIDIdef.all[\input].func.(64);
-)
+"abc"+0
 
 // load another model
 // b.sendMsg("/predictor/load", "/path/to/checkpoint");

notepredictor/notepredictor/model.py

@@ -193,13 +193,16 @@ def cell_state_names(self):
     def cell_state(self):
         return tuple(getattr(self, n) for n in self.cell_state_names())
 
-    def get_samplers(self, allow_start=False, allow_end=False):
+    def get_samplers(self, index_pitch=None, allow_start=False, allow_end=False):
         def sample_pitch(x):
             if not allow_start:
                 x[...,self.start_token] = -np.inf
             if not allow_end:
                 x[...,self.end_token] = -np.inf
-            return D.Categorical(logits=x).sample()
+            if index_pitch is not None:
+                return x.argsort(-1, True)[...,index_pitch]
+            else:
+                return D.Categorical(logits=x).sample()
 
         return (
             sample_pitch, 
@@ -294,7 +297,7 @@ def forward(self, pitches, times, velocities, validation=False):
     def predict(self, 
             pitch, time, vel, 
             fix_pitch=None, fix_time=None, fix_vel=None, 
-            allow_end=False, allow_start=False):
+            index_pitch=None, allow_start=False, allow_end=False):
         """
         supply the most recent note and return a prediction for the next note.
 
@@ -331,7 +334,7 @@ def predict(self,
 
             modalities = list(zip(
                 self.projections,
-                self.get_samplers(allow_start=allow_start, allow_end=allow_end),
+                self.get_samplers(index_pitch, allow_start, allow_end),
                 self.embeddings,
                 ))