Upload app_onnx.py

app_onnx.py

@@ -29,6 +29,7 @@ def softmax(x, axis):
     exp_x_shifted = np.exp(x - x_max)
     return exp_x_shifted / np.sum(exp_x_shifted, axis=axis, keepdims=True)
 
+
 def sample_top_p_k(probs, p, k, generator=None):
     if generator is None:
         generator = np.random
@@ -48,9 +49,10 @@ def sample_top_p_k(probs, p, k, generator=None):
     next_token = next_token.reshape(*shape[:-1])
     return next_token
 
+
 def apply_io_binding(model: rt.InferenceSession, inputs, outputs, batch_size, past_len, cur_len):
     io_binding = model.io_binding()
-    for input_ in model.get_inputs():
+    for input_ in  model.get_inputs():
         name = input_.name
         if name.startswith("past_key_values"):
             present_name = name.replace("past_key_values", "present")
@@ -80,8 +82,7 @@ def apply_io_binding(model: rt.InferenceSession, inputs, outputs, batch_size, pa
     return io_binding
 
 def generate(model, prompt=None, batch_size=1, max_len=512, temp=1.0, top_p=0.98, top_k=20,
-             disable_patch_change=False, disable_control_change=False, disable_channels=None, 
-             repetition_penalty=1.0, generator=None):
+             disable_patch_change=False, disable_control_change=False, disable_channels=None, generator=None):
     tokenizer = model[2]
     if disable_channels is not None:
         disable_channels = [tokenizer.parameter_ids["channel"][c] for c in disable_channels]
@@ -106,7 +107,7 @@ def generate(model, prompt=None, batch_size=1, max_len=512, temp=1.0, top_p=0.98
         prompt = prompt[..., :max_token_seq]
         if prompt.shape[-1] < max_token_seq:
             prompt = np.pad(prompt, ((0, 0), (0, 0), (0, max_token_seq - prompt.shape[-1])),
-                          mode="constant", constant_values=tokenizer.pad_id)
+                            mode="constant", constant_values=tokenizer.pad_id)
         input_tensor = prompt
     cur_len = input_tensor.shape[1]
     bar = tqdm.tqdm(desc="generating", total=max_len - cur_len, disable=in_space)
@@ -161,6 +162,7 @@ def generate(model, prompt=None, batch_size=1, max_len=512, temp=1.0, top_p=0.98
                 mask = mask[:, None, :]
                 x = next_token_seq
                 if i != 0:
+                    # cached
                     if i == 1:
                         hidden = np.zeros((batch_size, 0, emb_size), dtype=np.float32)
                         model1_inputs["hidden"] = rt.OrtValue.ortvalue_from_numpy(hidden, device_type=device)
@@ -176,16 +178,6 @@ def generate(model, prompt=None, batch_size=1, max_len=512, temp=1.0, top_p=0.98
                 model[1].run_with_iobinding(io_binding)
                 io_binding.synchronize_outputs()
                 logits = model1_outputs["y"].numpy()
-
-                # Apply repetition penalty
-                if repetition_penalty != 1.0:
-                    for b in range(batch_size):
-                        if not end[b]:
-                            prev_tokens = input_tensor[b, :cur_len].tolist()
-                            used_tokens = set(prev_tokens)
-                            for token in used_tokens:
-                                logits[b, :, token] /= repetition_penalty
-
                 scores = softmax(logits / temp, -1) * mask
                 samples = sample_top_p_k(scores, top_p, top_k, generator)
                 if i == 0:
@@ -204,8 +196,8 @@ def generate(model, prompt=None, batch_size=1, max_len=512, temp=1.0, top_p=0.98
                         break
             if next_token_seq.shape[1] < max_token_seq:
                 next_token_seq = np.pad(next_token_seq,
-                                      ((0, 0), (0, max_token_seq - next_token_seq.shape[-1])),
-                                      mode="constant", constant_values=tokenizer.pad_id)
+                                        ((0, 0), (0, max_token_seq - next_token_seq.shape[-1])),
+                                        mode="constant", constant_values=tokenizer.pad_id)
             next_token_seq = next_token_seq[:, None, :]
             input_tensor = np.concatenate([input_tensor, next_token_seq], axis=1)
             past_len = cur_len
@@ -594,12 +586,10 @@ if __name__ == "__main__":
             input_temp = gr.Slider(label="temperature", minimum=0.1, maximum=1.2, step=0.01, value=1)
             input_top_p = gr.Slider(label="top p", minimum=0.1, maximum=1, step=0.01, value=0.95)
             input_top_k = gr.Slider(label="top k", minimum=1, maximum=128, step=1, value=20)
-            input_rep_penalty = gr.Slider(label="repetition penalty", minimum=1.0, maximum=2.0, 
-                                         step=0.05, value=1.0)
             input_allow_cc = gr.Checkbox(label="allow midi cc event", value=True)
             input_render_audio = gr.Checkbox(label="render audio after generation", value=True)
             example3 = gr.Examples([[1, 0.94, 128], [1, 0.98, 20], [1, 0.98, 12]],
-                                  [input_temp, input_top_p, input_top_k])
+                                   [input_temp, input_top_p, input_top_k])
         run_btn = gr.Button("generate", variant="primary")
         # stop_btn = gr.Button("stop and output")
         output_midi_seq = gr.State()
@@ -615,13 +605,13 @@ if __name__ == "__main__":
                     midi_outputs.append(output_midi)
                     audio_outputs.append(output_audio)
         run_event = run_btn.click(run, [input_model, tab_select, output_midi_seq, output_continuation_state,
-                               input_continuation_select, input_instruments, input_drum_kit, input_bpm,
-                               input_time_sig, input_key_sig, input_midi, input_midi_events,
-                               input_reduce_cc_st, input_remap_track_channel,
-                               input_add_default_instr, input_remove_empty_channels,
-                               input_seed, input_seed_rand, input_gen_events, input_temp, input_top_p,
-                               input_top_k, input_rep_penalty, input_allow_cc],
-                         [output_midi_seq, output_continuation_state, input_seed, js_msg], queue=True)
+                                        input_continuation_select, input_instruments, input_drum_kit, input_bpm,
+                                        input_time_sig, input_key_sig, input_midi, input_midi_events,
+                                        input_reduce_cc_st, input_remap_track_channel,
+                                        input_add_default_instr, input_remove_empty_channels,
+                                        input_seed, input_seed_rand, input_gen_events, input_temp, input_top_p,
+                                        input_top_k, input_allow_cc],
+                                  [output_midi_seq, output_continuation_state, input_seed, js_msg], queue=True)
         finish_run_event = run_event.then(fn=finish_run,
                                           inputs=[input_model, output_midi_seq],
                                           outputs=midi_outputs + [js_msg],