Update app.py

app.py

@@ -22,6 +22,33 @@ model.load_state_dict(torch.load('another_epoch_1.75total.ckpt', map_location=de
 stoi = pickle.load(open('stoi.pkl', 'rb'))
 itos = pickle.load(open('itos.pkl', 'rb'))
 
+# Post-process generation
+# Completion
+def completion_to_song(c):
+  lines = c.split('\n')
+  kept_lines = []
+  notes = False
+  for l in lines:
+
+    # Record if we've hit music
+    if '|' in l:
+      notes = True
+
+    # Stop if we then go back to the start of another song
+    if 'T' in l and notes:
+      break
+    if 'T' in l and notes:
+      break
+
+    # Stop on an empty line
+    if len(l.strip()) < 2 and notes:
+      break
+    
+    # Otherwise keep the line
+    kept_lines.append(l)
+
+  return '\n'.join(kept_lines)
+
 # Generate function
 def generate_song(randomize, title, nu, ks, key):
   # Start sequence
@@ -56,7 +83,7 @@ demo = gr.Blocks()
 with demo:
     gr.Markdown("Quick demo for [WhistleGen v2](https://wandb.ai/johnowhitaker/whistlegen_v2/reports/WhistleGen-v2--VmlldzoyMTAwNjAz) which lets you generate folk music using a transformer model. I can't get the javascript needed for rendering and playback working with gradio, so this shows the raw ABC notation from the model and a link to view it properly in an external editor.")
     with gr.Row():
-        title = gr.Text(label='Title', value='The Song of AI')
+        title = gr.Text(label='Title', value='The March of AI')
         with gr.Column():
           nu = gr.Text(label='Note unit', value='1/8')
     with gr.Row():
@@ -66,7 +93,7 @@ with demo:
     with gr.Row():
       randomize = gr.Checkbox(label='Randomize (ignores settings above)', value=True)
     with gr.Row():
-        out = gr.HTML(label="Output", value='Output...')
+        out = gr.HTML(label="Output", value='Output should appear here (takes ~30s)')
     btn = gr.Button("Run")
     btn.click(fn=generate_song, inputs=[randomize, title, nu, key_signature, key], outputs=out)