Update app.py

app.py

@@ -96,6 +96,7 @@ with gr.Blocks() as demo:
     prefix_box = gr.TextArea(value="Twinkle Twinkle Little Star", label="Score title / text prefix")
     with gr.Row():
         submit_btn = gr.Button("Generate")
+        continue_btn = gr.Button("Continue")
         clear_btn = gr.Button("Clear history")
     with gr.Row():
         get_audio_btn = gr.Button("Convert to audio")
@@ -119,16 +120,28 @@ with gr.Blocks() as demo:
     def user_fn(user_message, history: list):
         return "", history + [{"role": "user", "content": user_message}]
 
-    @spaces.GPU
-    def bot_fn(history: list):
-        prefix = history[-1]["content"]
+    
+    def generate_fn(history: list):
+         # continue from user input
+        prefix = history[-1]["content"]  
         # prevent the model from continuing user's score title
         if prefix != '' and '\n' not in prefix:
             # prefix is a single line --> prefix is the score title
             prefix += '\n'
             
         history.append({"role": "assistant", "content": ""})
-        history[-1]["content"] += "Generating with the given prefix...\n"
+        # history[-1]["content"] += "Generating with the given prefix...\n"
+        for history in model_fn(prefix):
+            yield history
+
+    def continue_fn(history: list):
+         # continue from the last model output
+        prefix = history[-1]["content"]
+        for history in model_fn(prefix):
+            yield history
+        
+    @spaces.GPU
+    def model_fn(prefix):
         queue = Queue(maxsize=10)
         class MyStreamer:
             def put(self, tokens):
@@ -144,7 +157,7 @@ with gr.Blocks() as demo:
                 result = pipe(
                     prefix,
                     streamer=MyStreamer(), 
-                    max_length=3000,
+                    max_new_tokens=500,
                     top_p=0.9, temperature=0.6,
                 )
             except queue.Full:
@@ -161,12 +174,15 @@ with gr.Blocks() as demo:
             yield history
 
     prefix_box.submit(user_fn, [prefix_box, chatbot_box], [prefix_box, chatbot_box], queue=False).then(
-        bot_fn, chatbot_box, chatbot_box
+        generate_fn, chatbot_box, chatbot_box
     )
     submit_event = submit_btn.click(user_fn, [prefix_box, chatbot_box], [prefix_box, chatbot_box], queue=False).then(
-        bot_fn, chatbot_box, chatbot_box
+        generate_fn, chatbot_box, chatbot_box
+    )
+    continue_event = continue_btn.click(
+        continue_fn, chatbot_box, chatbot_box
     )
-    clear_btn.click(lambda: None, inputs=[], outputs=chatbot_box, cancels=[submit_event], queue=False)
+    clear_btn.click(lambda: None, inputs=[], outputs=chatbot_box, cancels=[submit_event, continue_event], queue=False)
 
     def get_audio_fn(history):
         i = random.randint(0, 1000_000_000)