Adding stop button

app.py

@@ -22,6 +22,13 @@ def validate_inputs(*args):
             return False
     return True
 
+stop_event = asyncio.Event()
+
+def stop_generation():
+    """
+    Stop the dataset generation process.
+    """
+    stop_event.set()
 
 def generate_synthetic_dataset(
     llm_model,
@@ -114,10 +121,18 @@ def generate_synthetic_dataset(
     else:
         return "Invalid dataset type"
 
-    dataset = asyncio.run(generator.agenerate_dataset())
-    dataset.save_dataset(hf_repo_name=hf_repo_name)
-    return "Dataset generated and saved successfully."
+    stop_event.clear()
+
+    async def generate():
+        if stop_event.is_set():
+            return "Dataset generation stopped."
+        dataset = await generator.agenerate_dataset()
+        if stop_event.is_set():
+            return "Dataset generation stopped."
+        dataset.save_dataset(hf_repo_name=hf_repo_name)
+        return "Dataset generated and saved successfully."
 
+    return asyncio.run(generate())
 
 def ui_main():
     """
@@ -206,6 +221,7 @@ def ui_main():
             )
 
         generate_button = gr.Button("Generate Dataset")
+        stop_button = gr.Button("Stop Generation")
         output = gr.Textbox(label="Operation Result", value="")
 
         generate_button.click(
@@ -227,6 +243,11 @@ def ui_main():
             outputs=output,
         )
 
+        stop_button.click(
+            stop_generation,
+            outputs=output,
+        )
+
     demo.launch(inbrowser=True, favicon_path=None)