Added local saving to CSV and JSON (#38)

* Local save

- Added save local function to chat tab (CSV, JSON)
- Rebuild UI with new feature
- CSS edit for gr.File (perfectionism)

* Local save

* Mistake

* Update chat.py

* Local save RAG and Textcat

* Rebuild UI

* Show save_local only if save_local_dir is provided

src/synthetic_dataset_generator/app.py

@@ -12,12 +12,13 @@ css = """
 .main_ui_logged_out{opacity: 0.3; pointer-events: none}
 button[role="tab"][aria-selected="true"] { border: 0; background: var(--button-primary-background-fill); color: white; border-top-right-radius: var(--radius-md); border-top-left-radius: var(--radius-md)}
 button[role="tab"][aria-selected="true"]:hover {border-color: var(--button-primary-background-fill); background: var(var(--button-primary-background-fill-hover))}
-.tabitem { border: 0; padding-inline: 0}
+.tabitem {border: 0; padding-inline: 0}
 .gallery-item {background: var(--background-fill-secondary); text-align: left}
-.table-wrap .tbody td { vertical-align: top }
-#system_prompt_examples { color: var(--body-text-color) !important; background-color: var(--block-background-fill) !important;}
+.table-wrap .tbody td {vertical-align: top}
+#system_prompt_examples {color: var(--body-text-color) !important; background-color: var(--block-background-fill) !important;}
 .container {padding-inline: 0 !important}
-#sign_in_button { flex-grow: 0; width: auto !important; display: flex; align-items: center; justify-content: center; margin: 0 auto; }
+#sign_in_button {flex-grow: 0; width: auto !important; display: flex; align-items: center; justify-content: center; margin: 0 auto;}
+.datasets {height: 70px;}
 """
 
 image = """<br><img src="https://raw.githubusercontent.com/argilla-io/synthetic-data-generator/main/assets/logo.svg" alt="Synthetic Data Generator Logo" style="display: block; margin-left: auto; margin-right: auto; width: clamp(50%, 400px, 100%)"/>"""

src/synthetic_dataset_generator/apps/base.py

@@ -12,9 +12,13 @@ from huggingface_hub import HfApi, upload_file, repo_exists
 from unstructured.chunking.title import chunk_by_title
 from unstructured.partition.auto import partition
 
-from synthetic_dataset_generator.constants import MAX_NUM_ROWS
+from synthetic_dataset_generator.constants import MAX_NUM_ROWS, SAVE_LOCAL_DIR
 from synthetic_dataset_generator.utils import get_argilla_client
 
+if SAVE_LOCAL_DIR is not None:
+    import os
+    os.makedirs(SAVE_LOCAL_DIR, exist_ok=True)
+
 
 def validate_argilla_user_workspace_dataset(
     dataset_name: str,

src/synthetic_dataset_generator/apps/chat.py

@@ -2,6 +2,7 @@ import ast
 import json
 import random
 import uuid
+import os
 from typing import Dict, List, Union
 
 import argilla as rg
@@ -30,6 +31,7 @@ from synthetic_dataset_generator.constants import (
     MODEL,
     MODEL_COMPLETION,
     SFT_AVAILABLE,
+    SAVE_LOCAL_DIR,
 )
 from synthetic_dataset_generator.pipelines.base import get_rewritten_prompts
 from synthetic_dataset_generator.pipelines.chat import (
@@ -264,7 +266,6 @@ 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,
@@ -506,7 +507,7 @@ def push_dataset(
         num_turns=num_turns,
         num_rows=num_rows,
         temperature=temperature,
-        temperature_completion=temperature_completion
+        temperature_completion=temperature_completion,
     )
     push_dataset_to_hub(
         dataframe=dataframe,
@@ -637,6 +638,45 @@ def push_dataset(
     return ""
 
 
+def save_local(
+    repo_id: str,
+    file_paths: list[str],
+    input_type: str,
+    system_prompt: str,
+    document_column: str,
+    num_turns: int,
+    num_rows: int,
+    temperature: float,
+    repo_name: str,
+    temperature_completion: Union[float, None] = None,
+) -> pd.DataFrame:
+    if input_type == "prompt-input":
+        dataframe = _get_dataframe()
+    else:
+        dataframe, _ = load_dataset_file(
+            repo_id=repo_id,
+            file_paths=file_paths,
+            input_type=input_type,
+            num_rows=num_rows,
+        )
+    dataframe = generate_dataset(
+        input_type=input_type,
+        dataframe=dataframe,
+        system_prompt=system_prompt,
+        document_column=document_column,
+        num_turns=num_turns,
+        num_rows=num_rows,
+        temperature=temperature,
+        temperature_completion=temperature_completion
+    )
+    local_dataset = Dataset.from_pandas(dataframe)
+    output_csv = os.path.join(SAVE_LOCAL_DIR, repo_name + ".csv")
+    output_json = os.path.join(SAVE_LOCAL_DIR, repo_name + ".json")
+    local_dataset.to_csv(output_csv, index=False)
+    local_dataset.to_json(output_json, index=False)
+    return output_csv, output_json
+
+
 def show_system_prompt_visibility():
     return {system_prompt: gr.Textbox(visible=True)}
 
@@ -670,6 +710,13 @@ 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():
+    return {
+        btn_save_local: gr.Button(visible=True),
+        csv_file: gr.File(visible=True),
+        json_file: gr.File(visible=True)
+    }
 
 
 ######################
@@ -852,6 +899,11 @@ with gr.Blocks() as app:
                     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):
                     success_message = gr.Markdown(
                         visible=True,
@@ -998,6 +1050,23 @@ with gr.Blocks() as app:
         inputs=[],
         outputs=[pipeline_code_ui],
     )
+    
+    btn_save_local.click(
+        save_local,
+        inputs=[
+            search_in,
+            file_in,
+            input_type,
+            system_prompt,
+            document_column,
+            num_turns,
+            num_rows,
+            temperature,
+            repo_name,
+            temperature_completion,
+        ],
+        outputs=[csv_file, json_file]
+    )
 
     clear_dataset_btn_part.click(fn=lambda: "", inputs=[], outputs=[search_in])
     clear_file_btn_part.click(fn=lambda: None, inputs=[], outputs=[file_in])
@@ -1011,3 +1080,5 @@ with gr.Blocks() as app:
     app.load(fn=get_org_dropdown, outputs=[org_name])
     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])

src/synthetic_dataset_generator/apps/rag.py

@@ -24,7 +24,7 @@ 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
+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,
@@ -486,6 +486,49 @@ def push_dataset(
     return ""
 
 
+def save_local(
+    repo_id: str,
+    file_paths: list[str],
+    input_type: str,
+    system_prompt: str,
+    document_column: str,
+    retrieval_reranking: list[str],
+    num_rows: int,
+    temperature: float,
+    repo_name: str,
+    temperature_completion: float,
+) -> pd.DataFrame:
+    retrieval = "Retrieval" in retrieval_reranking
+    reranking = "Reranking" in retrieval_reranking
+
+    if input_type == "prompt-input":
+        dataframe = pd.DataFrame(columns=["context", "question", "response"])
+    else:
+        dataframe, _ = load_dataset_file(
+            repo_id=repo_id,
+            file_paths=file_paths,
+            input_type=input_type,
+            num_rows=num_rows,
+        )
+    dataframe = generate_dataset(
+        input_type=input_type,
+        dataframe=dataframe,
+        system_prompt=system_prompt,
+        document_column=document_column,
+        retrieval=retrieval,
+        reranking=reranking,
+        num_rows=num_rows,
+        temperature=temperature,
+        temperature_completion=temperature_completion,
+    )
+    local_dataset = Dataset.from_pandas(dataframe)
+    output_csv = os.path.join(SAVE_LOCAL_DIR, repo_name + ".csv")
+    output_json = os.path.join(SAVE_LOCAL_DIR, repo_name + ".json")
+    local_dataset.to_csv(output_csv, index=False)
+    local_dataset.to_json(output_json, index=False)
+    return output_csv, output_json
+
+
 def show_system_prompt_visibility():
     return {system_prompt: gr.Textbox(visible=True)}
 
@@ -521,6 +564,14 @@ def show_temperature_completion():
         return {temperature_completion: gr.Slider(value=0.9, visible=True)}
 
 
+def show_save_local():
+    return {
+        btn_save_local: gr.Button(visible=True),
+        csv_file: gr.File(visible=True),
+        json_file: gr.File(visible=True)
+    }
+
+
 ######################
 # Gradio UI
 ######################
@@ -674,7 +725,14 @@ 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):
                 success_message = gr.Markdown(
                     visible=True,
@@ -822,6 +880,23 @@ with gr.Blocks() as app:
         outputs=[pipeline_code_ui],
     )
 
+    btn_save_local.click(
+        save_local,
+        inputs=[
+            search_in,
+            file_in,
+            input_type,
+            system_prompt,
+            document_column,
+            retrieval_reranking,
+            num_rows,
+            temperature,
+            repo_name,
+            temperature_completion,
+        ],
+        outputs=[csv_file, json_file]
+    )
+
     clear_dataset_btn_part.click(fn=lambda: "", inputs=[], outputs=[search_in])
     clear_file_btn_part.click(fn=lambda: None, inputs=[], outputs=[file_in])
     clear_prompt_btn_part.click(fn=lambda: "", inputs=[], outputs=[dataset_description])
@@ -835,3 +910,5 @@ with gr.Blocks() as app:
     app.load(fn=get_org_dropdown, outputs=[org_name])
     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])

src/synthetic_dataset_generator/apps/textcat.py

@@ -1,3 +1,4 @@
+import os
 import json
 import random
 import uuid
@@ -19,7 +20,7 @@ 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
+from synthetic_dataset_generator.constants import DEFAULT_BATCH_SIZE, SAVE_LOCAL_DIR
 from synthetic_dataset_generator.pipelines.base import get_rewritten_prompts
 from synthetic_dataset_generator.pipelines.embeddings import (
     get_embeddings,
@@ -406,6 +407,33 @@ def push_dataset(
     return ""
 
 
+def save_local(
+    system_prompt: str,
+    difficulty: str,
+    clarity: str,
+    labels: List[str],
+    multi_label: bool,
+    num_rows: int,
+    temperature: float,
+    repo_name: str,
+) -> pd.DataFrame:
+    dataframe = generate_dataset(
+        system_prompt=system_prompt,
+        difficulty=difficulty,
+        clarity=clarity,
+        multi_label=multi_label,
+        labels=labels,
+        num_rows=num_rows,
+        temperature=temperature,
+    )
+    local_dataset = Dataset.from_pandas(dataframe)
+    output_csv = os.path.join(SAVE_LOCAL_DIR, repo_name + ".csv")
+    output_json = os.path.join(SAVE_LOCAL_DIR, repo_name + ".json")
+    local_dataset.to_csv(output_csv, index=False)
+    local_dataset.to_json(output_json, index=False)
+    return output_csv, output_json
+
+
 def validate_input_labels(labels: List[str]) -> List[str]:
     if (
         not labels
@@ -425,6 +453,14 @@ def hide_pipeline_code_visibility():
     return {pipeline_code_ui: gr.Accordion(visible=False)}
 
 
+def show_save_local():
+    return {
+        btn_save_local: gr.Button(visible=True),
+        csv_file: gr.File(visible=True),
+        json_file: gr.File(visible=True)
+    }
+
+
 ######################
 # Gradio UI
 ######################
@@ -543,7 +579,14 @@ 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):
                 success_message = gr.Markdown(
                     visible=True,
@@ -643,6 +686,21 @@ with gr.Blocks() as app:
         inputs=[],
         outputs=[pipeline_code_ui],
     )
+    
+    btn_save_local.click(
+        save_local,
+        inputs=[
+            system_prompt,
+            difficulty,
+            clarity,
+            labels,
+            multi_label,
+            num_rows,
+            temperature,
+            repo_name,
+        ],
+        outputs=[csv_file, json_file]
+    )
 
     gr.on(
         triggers=[clear_btn_part.click, clear_btn_full.click],
@@ -660,3 +718,5 @@ with gr.Blocks() as app:
     app.load(fn=swap_visibility, outputs=main_ui)
     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])

src/synthetic_dataset_generator/constants.py

@@ -8,6 +8,9 @@ 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
+SAVE_LOCAL_DIR = os.getenv(key="SAVE_LOCAL_DIR", default=None)
+
 # Models
 MODEL = os.getenv("MODEL", "meta-llama/Meta-Llama-3.1-8B-Instruct")
 TOKENIZER_ID = os.getenv(key="TOKENIZER_ID", default=None)