add cancel timeout buttons and update ui

llama_lora/ui/finetune_ui.py

@@ -393,6 +393,8 @@ def do_abort_training():
 
 
 def finetune_ui():
+    things_that_might_timeout = []
+
     with gr.Blocks() as finetune_ui_blocks:
         with gr.Column(elem_id="finetune_ui_content"):
             with gr.Tab("Prepare"):
@@ -466,10 +468,13 @@ def finetune_ui():
                                 elem_id="dataset_plain_text_data_separator",
                                 placeholder=default_dataset_plain_text_data_separator,
                                 value=default_dataset_plain_text_data_separator)
-                        dataset_text_format.change(fn=handle_switch_dataset_text_format, inputs=[
-                                                   dataset_text_format], outputs=[dataset_plain_text_separators_group])
-                    dataset_text_load_sample_button.click(fn=load_sample_dataset_to_text_input, inputs=[
-                                                          dataset_text_format], outputs=[dataset_text])
+                        things_that_might_timeout.append(
+                            dataset_text_format.change(fn=handle_switch_dataset_text_format, inputs=[
+                                dataset_text_format], outputs=[dataset_plain_text_separators_group]))
+
+                    things_that_might_timeout.append(
+                        dataset_text_load_sample_button.click(fn=load_sample_dataset_to_text_input, inputs=[
+                            dataset_text_format], outputs=[dataset_text]))
                 gr.Markdown(
                     "💡 Switch to the \"Preview\" tab to verify that your inputs are correct.")
             with gr.Tab("Preview"):
@@ -484,14 +489,15 @@ def finetune_ui():
                     )
                 finetune_dataset_preview = gr.Dataframe(
                     wrap=True, elem_id="finetune_dataset_preview")
-            load_dataset_from.change(
-                fn=handle_switch_dataset_source,
-                inputs=[load_dataset_from],
-                outputs=[
-                    dataset_text_input_group,
-                    dataset_from_data_dir_group
-                ]
-            )
+            things_that_might_timeout.append(
+                load_dataset_from.change(
+                    fn=handle_switch_dataset_source,
+                    inputs=[load_dataset_from],
+                    outputs=[
+                        dataset_text_input_group,
+                        dataset_from_data_dir_group
+                    ]
+                ))
 
             dataset_inputs = [
                 template,
@@ -506,21 +512,23 @@ def finetune_ui():
             dataset_preview_inputs = dataset_inputs + \
                 [finetune_dataset_preview_show_actual_prompt]
             for i in dataset_preview_inputs:
-                i.change(
-                    fn=refresh_preview,
-                    inputs=dataset_preview_inputs,
-                    outputs=[finetune_dataset_preview,
-                             finetune_dataset_preview_info_message,
-                             dataset_from_text_message,
-                             dataset_from_data_dir_message
-                             ]
-                )
-
-            reload_selections_button.click(
+                things_that_might_timeout.append(
+                    i.change(
+                        fn=refresh_preview,
+                        inputs=dataset_preview_inputs,
+                        outputs=[finetune_dataset_preview,
+                                 finetune_dataset_preview_info_message,
+                                 dataset_from_text_message,
+                                 dataset_from_data_dir_message
+                                 ]
+                    ))
+
+            things_that_might_timeout.append(reload_selections_button.click(
                 reload_selections,
                 inputs=[template, dataset_from_data_dir],
                 outputs=[template, dataset_from_data_dir],
             )
+            )
 
             max_seq_length = gr.Slider(
                 minimum=1, maximum=4096, value=512,
@@ -621,6 +629,13 @@ def finetune_ui():
             inputs=None, outputs=None,
             cancels=[train_progress])
 
+        stop_timeoutable_btn = gr.Button(
+            "stop not-responding elements",
+            elem_id="inference_stop_timeoutable_btn",
+            elem_classes="foot_stop_timeoutable_btn")
+        stop_timeoutable_btn.click(
+            fn=None, inputs=None, outputs=None, cancels=things_that_might_timeout)
+
     finetune_ui_blocks.load(_js="""
     function finetune_ui_blocks_js() {
       // Auto load options

llama_lora/ui/inference_ui.py

@@ -175,6 +175,8 @@ def update_prompt_preview(prompt_template,
 
 
 def inference_ui():
+    things_that_might_timeout = []
+
     with gr.Blocks() as inference_ui_blocks:
         with gr.Row():
             lora_model = gr.Dropdown(
@@ -203,13 +205,20 @@ def inference_ui():
                         placeholder="Tell me about alpecas and llamas.",
                         elem_id="inference_variable_0"
                     )
-                    variable_1 = gr.Textbox(lines=2, label="", visible=False, elem_id="inference_variable_1")
-                    variable_2 = gr.Textbox(lines=2, label="", visible=False, elem_id="inference_variable_2")
-                    variable_3 = gr.Textbox(lines=2, label="", visible=False, elem_id="inference_variable_3")
-                    variable_4 = gr.Textbox(lines=2, label="", visible=False, elem_id="inference_variable_4")
-                    variable_5 = gr.Textbox(lines=2, label="", visible=False, elem_id="inference_variable_5")
-                    variable_6 = gr.Textbox(lines=2, label="", visible=False, elem_id="inference_variable_6")
-                    variable_7 = gr.Textbox(lines=2, label="", visible=False, elem_id="inference_variable_7")
+                    variable_1 = gr.Textbox(
+                        lines=2, label="", visible=False, elem_id="inference_variable_1")
+                    variable_2 = gr.Textbox(
+                        lines=2, label="", visible=False, elem_id="inference_variable_2")
+                    variable_3 = gr.Textbox(
+                        lines=2, label="", visible=False, elem_id="inference_variable_3")
+                    variable_4 = gr.Textbox(
+                        lines=2, label="", visible=False, elem_id="inference_variable_4")
+                    variable_5 = gr.Textbox(
+                        lines=2, label="", visible=False, elem_id="inference_variable_5")
+                    variable_6 = gr.Textbox(
+                        lines=2, label="", visible=False, elem_id="inference_variable_6")
+                    variable_7 = gr.Textbox(
+                        lines=2, label="", visible=False, elem_id="inference_variable_7")
 
                     with gr.Accordion("Preview", open=False, elem_id="inference_preview_prompt_container"):
                         preview_prompt = gr.Textbox(
@@ -303,19 +312,22 @@ def inference_ui():
                             interactive=False,
                             elem_id="inference_raw_output")
 
-        show_raw.change(
+        show_raw_change_event = show_raw.change(
             fn=lambda show_raw: gr.Accordion.update(visible=show_raw),
             inputs=[show_raw],
             outputs=[raw_output_group])
+        things_that_might_timeout.append(show_raw_change_event)
 
-        reload_selections_button.click(
+        reload_selections_event = reload_selections_button.click(
             reload_selections,
             inputs=[lora_model, prompt_template],
             outputs=[lora_model, prompt_template],
         )
+        things_that_might_timeout.append(reload_selections_event)
 
-        prompt_template.change(fn=handle_prompt_template_change, inputs=[prompt_template], outputs=[
-                               variable_0, variable_1, variable_2, variable_3, variable_4, variable_5, variable_6, variable_7])
+        prompt_template_change_event = prompt_template.change(fn=handle_prompt_template_change, inputs=[prompt_template], outputs=[
+            variable_0, variable_1, variable_2, variable_3, variable_4, variable_5, variable_6, variable_7])
+        things_that_might_timeout.append(prompt_template_change_event)
 
         generate_event = generate_btn.click(
             fn=do_inference,
@@ -339,9 +351,17 @@ def inference_ui():
         stop_btn.click(fn=None, inputs=None, outputs=None,
                        cancels=[generate_event])
 
-        update_prompt_preview_btn.click(fn=update_prompt_preview, inputs=[prompt_template,
-                                                                          variable_0, variable_1, variable_2, variable_3,
-                                                                          variable_4, variable_5, variable_6, variable_7,], outputs=preview_prompt)
+        update_prompt_preview_event = update_prompt_preview_btn.click(fn=update_prompt_preview, inputs=[prompt_template,
+                                                                                                        variable_0, variable_1, variable_2, variable_3,
+                                                                                                        variable_4, variable_5, variable_6, variable_7,], outputs=preview_prompt)
+        things_that_might_timeout.append(update_prompt_preview_event)
+
+        stop_timeoutable_btn = gr.Button(
+            "stop not-responding elements",
+            elem_id="inference_stop_timeoutable_btn",
+            elem_classes="foot_stop_timeoutable_btn")
+        stop_timeoutable_btn.click(
+            fn=None, inputs=None, outputs=None, cancels=things_that_might_timeout)
 
     inference_ui_blocks.load(_js="""
     function inference_ui_blocks_js() {

llama_lora/ui/main_page.py

@@ -395,6 +395,30 @@ def main_page_custom_css():
         }
     }
 
+    #tokenizer_encoded_tokens_input_textbox .codemirror-wrapper,
+    #tokenizer_decoded_text_input_textbox .codemirror-wrapper {
+        margin-bottom: -20px;
+    }
+
+    .foot_stop_timeoutable_btn {
+        align-self: flex-end;
+        border: 0 !important;
+        width: auto !important;
+        background: transparent !important;
+        box-shadow: none !important;
+        padding: 0 !important;
+        font-weight: 100 !important;
+        font-size: 80% !important;
+        text-decoration: underline;
+        opacity: 0.3;
+    }
+    .foot_stop_timeoutable_btn:hover {
+        opacity: 0.8;
+    }
+    .foot_stop_timeoutable_btn:active {
+        opacity: 1;
+    }
+
     .tippy-box[data-animation=scale-subtle][data-placement^=top]{transform-origin:bottom}.tippy-box[data-animation=scale-subtle][data-placement^=bottom]{transform-origin:top}.tippy-box[data-animation=scale-subtle][data-placement^=left]{transform-origin:right}.tippy-box[data-animation=scale-subtle][data-placement^=right]{transform-origin:left}.tippy-box[data-animation=scale-subtle][data-state=hidden]{transform:scale(.8);opacity:0}
     """
     return css

llama_lora/ui/tokenizer_ui.py

@@ -31,6 +31,8 @@ def handle_encode(decoded_tokens):
 
 
 def tokenizer_ui():
+    things_that_might_timeout = []
+
     with gr.Blocks() as tokenizer_ui_blocks:
         with gr.Row():
             with gr.Column():
@@ -50,7 +52,6 @@ def tokenizer_ui():
                 encode_btn = gr.Button("⬅️ Encode")
                 decoded_tokens_error_message = gr.Markdown(
                     "", visible=False, elem_classes="error-message")
-            stop_btn = gr.Button("Stop")
 
             decoding = decode_btn.click(
                 fn=handle_decode,
@@ -62,7 +63,15 @@ def tokenizer_ui():
                 inputs=[decoded_tokens],
                 outputs=[encoded_tokens, decoded_tokens_error_message],
             )
-            stop_btn.click(fn=None, inputs=None, outputs=None, cancels=[decoding, encoding])
+            things_that_might_timeout.append(decoding)
+            things_that_might_timeout.append(encoding)
+
+            stop_timeoutable_btn = gr.Button(
+                "stop not-responding elements",
+                elem_id="inference_stop_timeoutable_btn",
+                elem_classes="foot_stop_timeoutable_btn")
+            stop_timeoutable_btn.click(
+                fn=None, inputs=None, outputs=None, cancels=things_that_might_timeout)
 
     tokenizer_ui_blocks.load(_js="""
     function tokenizer_ui_blocks_js() {