Display loading animation at runtime

app.css

@@ -104,11 +104,43 @@
 }
 
 .solution-content .solution-content {
-  overflow-y: auto;
+  overflow-y: scroll;
   padding: 12px;
+  scroll-snap-align: end;
 }
 
 .run-btn {
   background: linear-gradient(to right, #ce7e53, #bb470b);
   color: white;
 }
+
+.running-btn {
+  background: linear-gradient(to right, #ce7e53, #bb470b);
+  color: white;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  position: relative;
+}
+
+.running-btn::before {
+  content: "";
+  position: absolute;
+  width: 22px;
+  height: 22px;
+  border-radius: 50%;
+  border: 3px solid white;
+  border-top-color: transparent;
+  /* border-right-color: transparent;
+  border-bottom-color: transparent; */
+  animation: spin 1s linear infinite;
+}
+
+@keyframes spin {
+  0% {
+    transform: rotate(0deg);
+  }
+  100% {
+    transform: rotate(360deg);
+  }
+}

app.py

@@ -661,6 +661,9 @@ def process_code(
 
 
 def solve_problem(problem, temperature, progress=gr.Progress()):
+    """
+    yield token: string, stop: bool
+    """
     problem = apply_template({"prompt": problem}, prompt=config.system_prompt)
     print(f"Problem: {problem}")
 
@@ -690,9 +693,10 @@ def solve_problem(problem, temperature, progress=gr.Progress()):
         for reponse_message, error in generate(messages, temperature):
             if reponse_message is not None:
                 step_reponse += reponse_message
-                yield step_reponse
+                yield step_reponse, False
 
                 if error:
+                    yield step_reponse, True
                     return
 
         sample["gen_texts"] = step_reponse
@@ -711,12 +715,12 @@ def solve_problem(problem, temperature, progress=gr.Progress()):
         for output_mseeage in run_code_reponse:
             if output_mseeage is not None:
                 step_reponse += output_mseeage
-                yield step_reponse
+                yield step_reponse, False
 
         if sample["should_prune"]:
             break
 
-    yield sample["gen_texts"]
+    yield sample["gen_texts"], True
 
 
 example_data = datasets.load_dataset(
@@ -754,7 +758,11 @@ def clear():
     return "", 0.1, "", problem_example_text_display, problem_example_text
 
 
+running_success = False
+
 with gr.Blocks(css=css, title="Math Olympiad Solver") as demo:
+    btn_list = []
+
     problem_example_text, problem_example_text_display = update_example_problem()
 
     with gr.Row(elem_classes="title"):
@@ -789,16 +797,51 @@ with gr.Blocks(css=css, title="Math Olympiad Solver") as demo:
             with gr.Accordion("Advanced Options", open=False):
                 temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.1, step=0.1, label="Temperature")
 
-            with gr.Row():
-                btn_clear = gr.Button("Clear")
+            with gr.Row() as btn_area:
+                btn_clear = gr.Button("Clear", elem_classes="clear-btn")
                 btn_run = gr.Button("Run", elem_classes="run-btn")
+                btn_list.append(btn_clear)
+                btn_list.append(btn_run)
 
         with gr.Column(scale=1, elem_classes="right"):
             gr.HTML("Solution", elem_classes="solution-title-content")
             out = gr.Markdown(elem_classes="solution-content", latex_delimiters=latex_delimiters)
 
         problem_example_text_hidden = gr.Markdown(value=problem_example_text, visible=False)
-        btn_run.click(fn=solve_problem, inputs=[inp, temperature], outputs=out)
+
+        def solve_problem_wrapper(inp_text, temperature):
+            global running_success
+            for token, stop in solve_problem(inp_text, temperature):
+                yield token
+
+                if stop:
+                    running_success = True
+
+        def mount_run_btn(btn):
+            btn.click(fn=solve_problem_wrapper, inputs=[inp, temperature], outputs=out)
+            btn.click(get_running_btns, None, outputs=btn_list)
+
+        def get_running_btns():
+            global running_success
+
+            btn_clear = gr.Button("Clear")
+            btn_run = gr.Button("", elem_classes="run-btn running-btn")
+            yield [btn_clear, btn_run]
+
+            time.sleep(3)
+
+            btn_clear = gr.Button("Clear")
+            btn_run = gr.Button("Run", elem_classes="run-btn")
+
+            while True:
+                if running_success:
+                    running_success = False
+                    yield [btn_clear, btn_run]
+                    time.sleep(1)
+                    mount_run_btn(btn_run)
+                    break
+
+                time.sleep(1)
 
         copy_btn.click(fn=lambda example: example, inputs=[problem_example_text_hidden], outputs=[inp])
 
@@ -814,6 +857,8 @@ with gr.Blocks(css=css, title="Math Olympiad Solver") as demo:
             ],
         )
 
+        mount_run_btn(btn_run)
+
         demo.load(
             update_example_problem,
             inputs=None,