feat: Add support for textcat

src/distilabel_dataset_generator/apps/textcat.py

@@ -1,4 +1,205 @@
-from src.distilabel_dataset_generator.utils import get_base_app
+from typing import List
 
-with get_base_app() as app:
-    pass
+import gradio as gr
+import pandas as pd
+
+from src.distilabel_dataset_generator.apps.base import (
+    get_main_ui,
+    get_pipeline_code_ui,
+    hide_success_message,
+    push_dataset_to_hub,
+    push_pipeline_code_to_hub,
+    show_success_message_argilla,
+    show_success_message_hub,
+    validate_argilla_user_workspace_dataset,
+)
+from src.distilabel_dataset_generator.pipelines.textcat import (
+    DEFAULT_DATASET_DESCRIPTIONS,
+    DEFAULT_DATASETS,
+    DEFAULT_SYSTEM_PROMPTS,
+    generate_pipeline_code,
+)
+
+
+def push_dataset_to_argilla(dataset: pd.DataFrame, dataset_name: str) -> pd.DataFrame:
+    return dataset
+
+
+def generate_system_prompt(dataset_description: str) -> str:
+    return dataset_description
+
+
+def generate_dataset(
+    system_prompt: str, labels: List[str], multi_label: bool
+) -> pd.DataFrame:
+    return pd.DataFrame({"prompt": [system_prompt], "completion": [system_prompt]})
+
+
+(
+    app,
+    main_ui,
+    custom_input_ui,
+    dataset_description,
+    examples,
+    btn_generate_system_prompt,
+    system_prompt,
+    sample_dataset,
+    btn_generate_sample_dataset,
+    dataset_name,
+    add_to_existing_dataset,
+    btn_generate_full_dataset_copy,
+    btn_generate_and_push_to_argilla,
+    btn_push_to_argilla,
+    org_name,
+    repo_name,
+    private,
+    btn_generate_full_dataset,
+    btn_generate_and_push_to_hub,
+    btn_push_to_hub,
+    final_dataset,
+    success_message,
+) = get_main_ui(
+    default_dataset_descriptions=DEFAULT_DATASET_DESCRIPTIONS,
+    default_system_prompts=DEFAULT_SYSTEM_PROMPTS,
+    default_datasets=DEFAULT_DATASETS,
+    fn_generate_system_prompt=generate_system_prompt,
+    fn_generate_dataset=generate_dataset,
+)
+
+with app:
+    with main_ui:
+        with custom_input_ui:
+            labels = gr.Dropdown(
+                choices=[],
+                allow_custom_value=True,
+                interactive=True,
+                label="Labels",
+                multiselect=True,
+            )
+            num_labels = gr.Number(
+                label="Number of labels", value=2, 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)
+        )
+
+    # define app triggers
+    gr.on(
+        triggers=[
+            btn_generate_full_dataset.click,
+            btn_generate_full_dataset_copy.click,
+        ],
+        fn=hide_success_message,
+        outputs=[success_message],
+    ).then(
+        fn=generate_dataset,
+        inputs=[system_prompt, labels, multi_label],
+        outputs=[final_dataset],
+        show_progress=True,
+    )
+
+    btn_generate_and_push_to_argilla.click(
+        fn=validate_argilla_user_workspace_dataset,
+        inputs=[dataset_name, final_dataset, add_to_existing_dataset],
+        outputs=[final_dataset],
+        show_progress=True,
+    ).success(
+        fn=hide_success_message,
+        outputs=[success_message],
+    ).success(
+        fn=generate_dataset,
+        inputs=[system_prompt, labels, multi_label],
+        outputs=[final_dataset],
+        show_progress=True,
+    ).success(
+        fn=push_dataset_to_argilla,
+        inputs=[final_dataset, dataset_name],
+        outputs=[final_dataset],
+        show_progress=True,
+    ).success(
+        fn=show_success_message_argilla,
+        inputs=[],
+        outputs=[success_message],
+    )
+
+    btn_generate_and_push_to_hub.click(
+        fn=hide_success_message,
+        outputs=[success_message],
+    ).then(
+        fn=generate_dataset,
+        inputs=[system_prompt, labels, multi_label],
+        outputs=[final_dataset],
+        show_progress=True,
+    ).then(
+        fn=push_dataset_to_hub,
+        inputs=[final_dataset, private, org_name, repo_name],
+        outputs=[final_dataset],
+        show_progress=True,
+    ).then(
+        fn=push_pipeline_code_to_hub,
+        inputs=[pipeline_code, org_name, repo_name],
+        outputs=[],
+        show_progress=True,
+    ).success(
+        fn=show_success_message_hub,
+        inputs=[org_name, repo_name],
+        outputs=[success_message],
+    )
+
+    btn_push_to_hub.click(
+        fn=hide_success_message,
+        outputs=[success_message],
+    ).then(
+        fn=push_dataset_to_hub,
+        inputs=[final_dataset, private, org_name, repo_name],
+        outputs=[final_dataset],
+        show_progress=True,
+    ).then(
+        fn=push_pipeline_code_to_hub,
+        inputs=[pipeline_code, org_name, repo_name],
+        outputs=[],
+        show_progress=True,
+    ).success(
+        fn=show_success_message_hub,
+        inputs=[org_name, repo_name],
+        outputs=[success_message],
+    )
+
+    btn_push_to_argilla.click(
+        fn=hide_success_message,
+        outputs=[success_message],
+    ).success(
+        fn=validate_argilla_user_workspace_dataset,
+        inputs=[dataset_name, final_dataset, add_to_existing_dataset],
+        outputs=[final_dataset],
+        show_progress=True,
+    ).success(
+        fn=push_dataset_to_argilla,
+        inputs=[final_dataset, dataset_name],
+        outputs=[final_dataset],
+        show_progress=True,
+    ).success(
+        fn=show_success_message_argilla,
+        inputs=[],
+        outputs=[success_message],
+    )
+
+    system_prompt.change(
+        fn=generate_pipeline_code,
+        inputs=[system_prompt, labels, multi_label],
+        outputs=[pipeline_code],
+    )
+    labels.change(
+        fn=generate_pipeline_code,
+        inputs=[system_prompt, labels, multi_label],
+        outputs=[pipeline_code],
+    )
+    multi_label.change(
+        fn=generate_pipeline_code,
+        inputs=[system_prompt, labels, multi_label],
+        outputs=[pipeline_code],
+    )

src/distilabel_dataset_generator/pipelines/textcat.py

@@ -0,0 +1,44 @@
+from typing import List
+
+import pandas as pd
+
+DEFAULT_DATASET_DESCRIPTIONS = [
+    "A dataset covering customer reviews for an e-commerce website.",
+    "A dataset covering news articles about various topics.",
+]
+
+DEFAULT_DATASETS = [
+    pd.DataFrame.from_dict(
+        {
+            "text": [
+                "I love the product! It's amazing and I'll buy it again.",
+                "The product was okay, but I wouldn't buy it again.",
+            ],
+            "label": ["positive", "negative"],
+        }
+    ),
+    pd.DataFrame.from_dict(
+        {
+            "text": [
+                "Yesterday, the US stock market had a significant increase.",
+                "New research suggests that the Earth is not a perfect sphere.",
+            ],
+            "label": [["economy", "politics"], ["science", "environment"]],
+        }
+    ),
+]
+
+DEFAULT_SYSTEM_PROMPTS = [
+    "Classify the following customer review as positive or negative.",
+    "Classify the following news article into one or more categories.",
+]
+
+
+def generate_pipeline_code(
+    system_prompt: str, labels: List[str], multi_label: bool
+) -> str:
+    return """
+    from distilabel import Distilabel
+
+    #### PIPELINE CODE HERE
+    """