refactor: add local save to README and improve layout

README.md

@@ -104,6 +104,10 @@ Optionally, you can also push your datasets to Argilla for further curation by s
 - `ARGILLA_API_KEY`: Your Argilla API key to push your datasets to Argilla.
 - `ARGILLA_API_URL`: Your Argilla API URL to push your datasets to Argilla.
 
+To save the generated datasets to a local directory instead of pushing them to the Hugging Face Hub, set the following environment variable:
+
+- `SAVE_LOCAL_DIR`: The local directory to save the generated datasets to.
+
 ### Argilla integration
 
 Argilla is an open source tool for data curation. It allows you to annotate and review datasets, and push curated datasets to the Hugging Face Hub. You can easily get started with Argilla by following the [quickstart guide](https://docs.argilla.io/latest/getting_started/quickstart/).

src/synthetic_dataset_generator/apps/chat.py

@@ -1,8 +1,8 @@
 import ast
 import json
+import os
 import random
 import uuid
-import os
 from typing import Dict, List, Union
 
 import argilla as rg
@@ -30,8 +30,8 @@ from synthetic_dataset_generator.constants import (
     DEFAULT_BATCH_SIZE,
     MODEL,
     MODEL_COMPLETION,
-    SFT_AVAILABLE,
     SAVE_LOCAL_DIR,
+    SFT_AVAILABLE,
 )
 from synthetic_dataset_generator.pipelines.base import get_rewritten_prompts
 from synthetic_dataset_generator.pipelines.chat import (
@@ -266,6 +266,7 @@ def generate_dataset_from_prompt(
     progress(1.0, desc="Dataset generation completed")
     return dataframe
 
+
 def generate_dataset_from_seed(
     dataframe: pd.DataFrame,
     document_column: str,
@@ -369,7 +370,9 @@ def generate_dataset_from_seed(
                 follow_up_instructions = list(
                     follow_up_generator_instruction.process(inputs=conversations_batch)
                 )
-                for conv, follow_up in zip(conversations_batch, follow_up_instructions[0]):
+                for conv, follow_up in zip(
+                    conversations_batch, follow_up_instructions[0]
+                ):
                     conv["messages"].append(
                         {"role": "user", "content": follow_up["generation"]}
                     )
@@ -667,7 +670,7 @@ def save_local(
         num_turns=num_turns,
         num_rows=num_rows,
         temperature=temperature,
-        temperature_completion=temperature_completion
+        temperature_completion=temperature_completion,
     )
     local_dataset = Dataset.from_pandas(dataframe)
     output_csv = os.path.join(SAVE_LOCAL_DIR, repo_name + ".csv")
@@ -710,12 +713,30 @@ def hide_pipeline_code_visibility():
 def show_temperature_completion():
     if MODEL != MODEL_COMPLETION:
         return {temperature_completion: gr.Slider(value=0.9, visible=True)}
-        
+
+
+def show_save_local_button():
+    return {btn_save_local: gr.Button(visible=True)}
+
+
+def hide_save_local_button():
+    return {btn_save_local: gr.Button(visible=False)}
+
+
 def show_save_local():
+    gr.update(success_message, min_height=0)
     return {
-        btn_save_local: gr.Button(visible=True),
         csv_file: gr.File(visible=True),
-        json_file: gr.File(visible=True)
+        json_file: gr.File(visible=True),
+        success_message: success_message
+    }
+
+def hide_save_local():
+    gr.update(success_message, min_height=100)
+    return {
+        csv_file: gr.File(visible=False),
+        json_file: gr.File(visible=False),
+        success_message: success_message,
     }
 
 
@@ -902,12 +923,20 @@ with gr.Blocks() as app:
                     btn_save_local = gr.Button(
                         "Save locally", variant="primary", scale=2, visible=False
                     )
-                    csv_file = gr.File(label="CSV", elem_classes="datasets", visible=False)
-                    json_file = gr.File(label="JSON", elem_classes="datasets", visible=False)
                 with gr.Column(scale=3):
+                    csv_file = gr.File(
+                        label="CSV",
+                        elem_classes="datasets",
+                        visible=False,
+                    )
+                    json_file = gr.File(
+                        label="JSON",
+                        elem_classes="datasets",
+                        visible=False,
+                    )
                     success_message = gr.Markdown(
-                        visible=True,
-                        min_height=100,  # don't remove this otherwise progress is not visible
+                        visible=False,
+                        min_height=0 # don't remove this otherwise progress is not visible
                     )
                     with gr.Accordion(
                         "Customize your pipeline with distilabel",
@@ -1005,6 +1034,9 @@ with gr.Blocks() as app:
         fn=validate_push_to_hub,
         inputs=[org_name, repo_name],
         outputs=[success_message],
+    ).success(
+        fn=hide_save_local,
+        outputs=[csv_file, json_file, success_message],
     ).success(
         fn=hide_success_message,
         outputs=[success_message],
@@ -1050,8 +1082,19 @@ with gr.Blocks() as app:
         inputs=[],
         outputs=[pipeline_code_ui],
     )
-    
+
     btn_save_local.click(
+        fn=hide_success_message,
+        outputs=[success_message],
+    ).success(
+        fn=hide_pipeline_code_visibility,
+        inputs=[],
+        outputs=[pipeline_code_ui],
+    ).success(
+        fn=show_save_local,
+        inputs=[],
+        outputs=[csv_file, json_file, success_message],
+    ).success(
         save_local,
         inputs=[
             search_in,
@@ -1065,7 +1108,22 @@ with gr.Blocks() as app:
             repo_name,
             temperature_completion,
         ],
-        outputs=[csv_file, json_file]
+        outputs=[csv_file, json_file],
+    ).success(
+        fn=generate_pipeline_code,
+        inputs=[
+            search_in,
+            input_type,
+            system_prompt,
+            document_column,
+            num_turns,
+            num_rows,
+        ],
+        outputs=[pipeline_code],
+    ).success(
+        fn=show_pipeline_code_visibility,
+        inputs=[],
+        outputs=[pipeline_code_ui],
     )
 
     clear_dataset_btn_part.click(fn=lambda: "", inputs=[], outputs=[search_in])
@@ -1081,4 +1139,4 @@ with gr.Blocks() as app:
     app.load(fn=get_random_repo_name, outputs=[repo_name])
     app.load(fn=show_temperature_completion, outputs=[temperature_completion])
     if SAVE_LOCAL_DIR is not None:
-        app.load(fn=show_save_local, outputs=[btn_save_local, csv_file, json_file])
+        app.load(fn=show_save_local_button, outputs=btn_save_local)

src/synthetic_dataset_generator/apps/rag.py

@@ -24,7 +24,12 @@ from synthetic_dataset_generator.apps.base import (
     validate_argilla_user_workspace_dataset,
     validate_push_to_hub,
 )
-from synthetic_dataset_generator.constants import DEFAULT_BATCH_SIZE, MODEL, MODEL_COMPLETION, SAVE_LOCAL_DIR
+from synthetic_dataset_generator.constants import (
+    DEFAULT_BATCH_SIZE,
+    MODEL,
+    MODEL_COMPLETION,
+    SAVE_LOCAL_DIR,
+)
 from synthetic_dataset_generator.pipelines.base import get_rewritten_prompts
 from synthetic_dataset_generator.pipelines.embeddings import (
     get_embeddings,
@@ -156,7 +161,7 @@ def generate_dataset(
         is_sample=is_sample,
     )
     response_generator = get_response_generator(
-        temperature = temperature_completion or temperature , is_sample=is_sample
+        temperature=temperature_completion or temperature, is_sample=is_sample
     )
     if reranking:
         reranking_generator = get_sentence_pair_generator(
@@ -564,11 +569,29 @@ def show_temperature_completion():
         return {temperature_completion: gr.Slider(value=0.9, visible=True)}
 
 
+def show_save_local_button():
+    return {btn_save_local: gr.Button(visible=True)}
+
+
+def hide_save_local_button():
+    return {btn_save_local: gr.Button(visible=False)}
+
+
 def show_save_local():
+    gr.update(success_message, min_height=0)
     return {
-        btn_save_local: gr.Button(visible=True),
         csv_file: gr.File(visible=True),
-        json_file: gr.File(visible=True)
+        json_file: gr.File(visible=True),
+        success_message: success_message,
+    }
+
+
+def hide_save_local():
+    gr.update(success_message, min_height=100)
+    return {
+        csv_file: gr.File(visible=False),
+        json_file: gr.File(visible=False),
+        success_message: success_message,
     }
 
 
@@ -725,18 +748,24 @@ with gr.Blocks() as app:
                     interactive=True,
                     scale=1,
                 )
-                btn_push_to_hub = gr.Button(
-                    "Push to Hub", variant="primary", scale=2
-                )
+                btn_push_to_hub = gr.Button("Push to Hub", variant="primary", scale=2)
                 btn_save_local = gr.Button(
                     "Save locally", variant="primary", scale=2, visible=False
                 )
-                csv_file = gr.File(label="CSV", elem_classes="datasets", visible=False)
-                json_file = gr.File(label="JSON", elem_classes="datasets", visible=False)
             with gr.Column(scale=3):
+                csv_file = gr.File(
+                    label="CSV",
+                    elem_classes="datasets",
+                    visible=False,
+                )
+                json_file = gr.File(
+                    label="JSON",
+                    elem_classes="datasets",
+                    visible=False,
+                )
                 success_message = gr.Markdown(
-                    visible=True,
-                    min_height=100,  # don't remove this otherwise progress is not visible
+                    visible=False,
+                    min_height=0,  # don't remove this otherwise progress is not visible
                 )
                 with gr.Accordion(
                     "Customize your pipeline with distilabel",
@@ -834,6 +863,9 @@ with gr.Blocks() as app:
         fn=validate_push_to_hub,
         inputs=[org_name, repo_name],
         outputs=[success_message],
+    ).success(
+        fn=hide_save_local,
+        outputs=[csv_file, json_file, success_message],
     ).success(
         fn=hide_success_message,
         outputs=[success_message],
@@ -881,6 +913,17 @@ with gr.Blocks() as app:
     )
 
     btn_save_local.click(
+        fn=hide_success_message,
+        outputs=[success_message],
+    ).success(
+        fn=hide_pipeline_code_visibility,
+        inputs=[],
+        outputs=[pipeline_code_ui],
+    ).success(
+        fn=show_save_local,
+        inputs=[],
+        outputs=[csv_file, json_file, success_message],
+    ).success(
         save_local,
         inputs=[
             search_in,
@@ -894,7 +937,22 @@ with gr.Blocks() as app:
             repo_name,
             temperature_completion,
         ],
-        outputs=[csv_file, json_file]
+        outputs=[csv_file, json_file],
+    ).success(
+        fn=generate_pipeline_code,
+        inputs=[
+            search_in,
+            input_type,
+            system_prompt,
+            document_column,
+            retrieval_reranking,
+            num_rows,
+        ],
+        outputs=[pipeline_code],
+    ).success(
+        fn=show_pipeline_code_visibility,
+        inputs=[],
+        outputs=[pipeline_code_ui],
     )
 
     clear_dataset_btn_part.click(fn=lambda: "", inputs=[], outputs=[search_in])
@@ -911,4 +969,4 @@ with gr.Blocks() as app:
     app.load(fn=get_random_repo_name, outputs=[repo_name])
     app.load(fn=show_temperature_completion, outputs=[temperature_completion])
     if SAVE_LOCAL_DIR is not None:
-        app.load(fn=show_save_local, outputs=[btn_save_local, csv_file, json_file])
+        app.load(fn=show_save_local_button, outputs=btn_save_local)

src/synthetic_dataset_generator/apps/textcat.py

@@ -1,5 +1,5 @@
-import os
 import json
+import os
 import random
 import uuid
 from typing import List, Union
@@ -453,11 +453,29 @@ def hide_pipeline_code_visibility():
     return {pipeline_code_ui: gr.Accordion(visible=False)}
 
 
+def show_save_local_button():
+    return {btn_save_local: gr.Button(visible=True)}
+
+
+def hide_save_local_button():
+    return {btn_save_local: gr.Button(visible=False)}
+
+
 def show_save_local():
+    gr.update(success_message, min_height=0)
     return {
-        btn_save_local: gr.Button(visible=True),
         csv_file: gr.File(visible=True),
-        json_file: gr.File(visible=True)
+        json_file: gr.File(visible=True),
+        success_message: success_message,
+    }
+
+
+def hide_save_local():
+    gr.update(success_message, min_height=100)
+    return {
+        csv_file: gr.File(visible=False),
+        json_file: gr.File(visible=False),
+        success_message: success_message,
     }
 
 
@@ -579,18 +597,24 @@ with gr.Blocks() as app:
                     interactive=True,
                     scale=1,
                 )
-                btn_push_to_hub = gr.Button(
-                    "Push to Hub", variant="primary", scale=2
-                )
+                btn_push_to_hub = gr.Button("Push to Hub", variant="primary", scale=2)
                 btn_save_local = gr.Button(
                     "Save locally", variant="primary", scale=2, visible=False
                 )
-                csv_file = gr.File(label="CSV", elem_classes="datasets", visible=False)
-                json_file = gr.File(label="JSON", elem_classes="datasets", visible=False)
             with gr.Column(scale=3):
+                csv_file = gr.File(
+                    label="CSV",
+                    elem_classes="datasets",
+                    visible=False,
+                )
+                json_file = gr.File(
+                    label="JSON",
+                    elem_classes="datasets",
+                    visible=False,
+                )
                 success_message = gr.Markdown(
-                    visible=True,
-                    min_height=100,  # don't remove this otherwise progress is not visible
+                    visible=False,
+                    min_height=0,  # don't remove this otherwise progress is not visible
                 )
                 with gr.Accordion(
                     "Customize your pipeline with distilabel",
@@ -643,6 +667,9 @@ with gr.Blocks() as app:
         fn=validate_input_labels,
         inputs=[labels],
         outputs=[labels],
+    ).success(
+        fn=hide_save_local,
+        outputs=[csv_file, json_file, success_message],
     ).success(
         fn=hide_success_message,
         outputs=[success_message],
@@ -686,8 +713,19 @@ with gr.Blocks() as app:
         inputs=[],
         outputs=[pipeline_code_ui],
     )
-    
+
     btn_save_local.click(
+        fn=hide_success_message,
+        outputs=[success_message],
+    ).success(
+        fn=hide_pipeline_code_visibility,
+        inputs=[],
+        outputs=[pipeline_code_ui],
+    ).success(
+        fn=show_save_local,
+        inputs=[],
+        outputs=[csv_file, json_file, success_message],
+    ).success(
         save_local,
         inputs=[
             system_prompt,
@@ -699,7 +737,22 @@ with gr.Blocks() as app:
             temperature,
             repo_name,
         ],
-        outputs=[csv_file, json_file]
+        outputs=[csv_file, json_file],
+    ).success(
+        fn=generate_pipeline_code,
+        inputs=[
+            system_prompt,
+            difficulty,
+            clarity,
+            labels,
+            multi_label,
+            num_rows,
+        ],
+        outputs=[pipeline_code],
+    ).success(
+        fn=show_pipeline_code_visibility,
+        inputs=[],
+        outputs=[pipeline_code_ui],
     )
 
     gr.on(
@@ -719,4 +772,4 @@ with gr.Blocks() as app:
     app.load(fn=get_org_dropdown, outputs=[org_name])
     app.load(fn=get_random_repo_name, outputs=[repo_name])
     if SAVE_LOCAL_DIR is not None:
-        app.load(fn=show_save_local, outputs=[btn_save_local, csv_file, json_file])
+        app.load(fn=show_save_local_button, outputs=btn_save_local)

src/synthetic_dataset_generator/constants.py

@@ -8,7 +8,7 @@ MAX_NUM_TOKENS = int(os.getenv("MAX_NUM_TOKENS", 2048))
 MAX_NUM_ROWS = int(os.getenv("MAX_NUM_ROWS", 1000))
 DEFAULT_BATCH_SIZE = int(os.getenv("DEFAULT_BATCH_SIZE", 5))
 
-# Directory for outputs
+# Directory to locally save the generated data
 SAVE_LOCAL_DIR = os.getenv(key="SAVE_LOCAL_DIR", default=None)
 
 # Models