feat: Add buttons to align with textcat and textcatgenerator arguments

src/distilabel_dataset_generator/apps/textcat.py

@@ -30,7 +30,12 @@ def generate_system_prompt(dataset_description: str) -> str:
 
 
 def generate_dataset(
-    system_prompt: str, labels: List[str], multi_label: bool
+    system_prompt: str,
+    difficulty: str,
+    clarity: str,
+    labels: List[str],
+    num_labels: int,
+    num_rows: int,
 ) -> pd.DataFrame:
     return pd.DataFrame({"prompt": [system_prompt], "completion": [system_prompt]})
 
@@ -69,6 +74,29 @@ def generate_dataset(
 with app:
     with main_ui:
         with custom_input_ui:
+            difficulty = gr.Dropdown(
+                choices=[
+                    ("High School", "high school"),
+                    ("College", "college"),
+                    ("PhD", "phd"),
+                    ("Mixed", "mixed"),
+                ],
+                value="mixed",
+                label="Difficulty",
+            )
+            clarity = gr.Dropdown(
+                choices=[
+                    ("Clear", "Clear"),
+                    (
+                        "Understandable",
+                        "understandable with some effort",
+                    ),
+                    ("Ambiguous", "Ambiguous"),
+                    ("Mixed", "mixed"),
+                ],
+                value="mixed",
+                label="Clarity",
+            )
             labels = gr.Dropdown(
                 choices=[],
                 allow_custom_value=True,
@@ -77,14 +105,21 @@ with app:
                 multiselect=True,
             )
             num_labels = gr.Number(
-                label="Number of labels", value=2, minimum=1, maximum=10
+                label="Number of labels", value=1, minimum=1, maximum=10
             )
             num_rows = gr.Number(
                 label="Number of rows", value=10, minimum=1, maximum=500
             )
 
         pipeline_code = get_pipeline_code_ui(
-            generate_pipeline_code(system_prompt.value, labels.value, multi_label.value)
+            generate_pipeline_code(
+                system_prompt.value,
+                difficulty=difficulty.value,
+                clarity=clarity.value,
+                labels=labels.value,
+                num_labels=num_labels.value,
+                num_rows=num_rows.value,
+            )
         )
 
     # define app triggers
@@ -97,7 +132,7 @@ with app:
         outputs=[success_message],
     ).then(
         fn=generate_dataset,
-        inputs=[system_prompt, labels, multi_label],
+        inputs=[system_prompt, difficulty, clarity, labels, num_labels, num_rows],
         outputs=[final_dataset],
         show_progress=True,
     )
@@ -112,7 +147,7 @@ with app:
         outputs=[success_message],
     ).success(
         fn=generate_dataset,
-        inputs=[system_prompt, labels, multi_label],
+        inputs=[system_prompt, difficulty, clarity, labels, num_labels, num_rows],
         outputs=[final_dataset],
         show_progress=True,
     ).success(
@@ -131,7 +166,7 @@ with app:
         outputs=[success_message],
     ).then(
         fn=generate_dataset,
-        inputs=[system_prompt, labels, multi_label],
+        inputs=[system_prompt, difficulty, clarity, labels, num_labels, num_rows],
         outputs=[final_dataset],
         show_progress=True,
     ).then(
@@ -190,16 +225,26 @@ with app:
 
     system_prompt.change(
         fn=generate_pipeline_code,
-        inputs=[system_prompt, labels, multi_label],
+        inputs=[system_prompt, difficulty, clarity, labels, num_labels, num_rows],
+        outputs=[pipeline_code],
+    )
+    difficulty.change(
+        fn=generate_pipeline_code,
+        inputs=[system_prompt, difficulty, clarity, labels, num_labels, num_rows],
+        outputs=[pipeline_code],
+    )
+    clarity.change(
+        fn=generate_pipeline_code,
+        inputs=[system_prompt, difficulty, clarity, labels, num_labels, num_rows],
         outputs=[pipeline_code],
     )
     labels.change(
         fn=generate_pipeline_code,
-        inputs=[system_prompt, labels, multi_label],
+        inputs=[system_prompt, difficulty, clarity, labels, num_labels, num_rows],
         outputs=[pipeline_code],
     )
-    multi_label.change(
+    num_labels.change(
         fn=generate_pipeline_code,
-        inputs=[system_prompt, labels, multi_label],
+        inputs=[system_prompt, difficulty, clarity, labels, num_labels, num_rows],
         outputs=[pipeline_code],
     )

src/distilabel_dataset_generator/pipelines/textcat.py

@@ -35,7 +35,12 @@ DEFAULT_SYSTEM_PROMPTS = [
 
 
 def generate_pipeline_code(
-    system_prompt: str, labels: List[str], multi_label: bool
+    system_prompt: str,
+    difficulty: str,
+    clarity: str,
+    labels: List[str],
+    num_labels: int,
+    num_rows: int,
 ) -> str:
     return """
     from distilabel import Distilabel