Merge branch 'main' of https://github.com/argilla-io/synthetic-data-generator into feat/improve-textcat

src/synthetic_dataset_generator/apps/base.py

@@ -1,19 +1,13 @@
 import io
 import uuid
-from typing import List, Union
+from typing import Union
 
 import argilla as rg
 import gradio as gr
-import pandas as pd
-from datasets import ClassLabel, Dataset, Features, Sequence, Value
-from distilabel.distiset import Distiset
 from gradio import OAuthToken
 from huggingface_hub import HfApi, upload_file
 
-from synthetic_dataset_generator.constants import TEXTCAT_TASK
-from synthetic_dataset_generator.utils import (
-    get_argilla_client,
-)
+from synthetic_dataset_generator.utils import get_argilla_client
 
 
 def validate_argilla_user_workspace_dataset(
@@ -52,7 +46,7 @@ def push_pipeline_code_to_hub(
     oauth_token: Union[OAuthToken, None] = None,
     progress=gr.Progress(),
 ):
-    repo_id = validate_push_to_hub(org_name, repo_name)
+    repo_id: str | None = validate_push_to_hub(org_name, repo_name)
     progress(0.1, desc="Uploading pipeline code")
     with io.BytesIO(pipeline_code.encode("utf-8")) as f:
         upload_file(

src/synthetic_dataset_generator/apps/eval.py

@@ -19,6 +19,7 @@ from huggingface_hub import HfApi, repo_exists
 
 from synthetic_dataset_generator.apps.base import (
     hide_success_message,
+    push_pipeline_code_to_hub,
     show_success_message,
     validate_argilla_user_workspace_dataset,
     validate_push_to_hub,
@@ -346,7 +347,12 @@ def evaluate_sample_dataset(
 
 
 def push_dataset_to_hub(
-    dataframe: pd.DataFrame, org_name: str, repo_name: str, oauth_token, private
+    dataframe: pd.DataFrame,
+    org_name: str,
+    repo_name: str,
+    oauth_token: Union[gr.OAuthToken, None],
+    private: bool,
+    pipeline_code: str,
 ):
     repo_id = validate_push_to_hub(org_name, repo_name)
     distiset = Distiset({"default": Dataset.from_pandas(dataframe)})
@@ -357,6 +363,7 @@ def push_dataset_to_hub(
         token=oauth_token.token,
         create_pr=False,
     )
+    push_pipeline_code_to_hub(pipeline_code, org_name, repo_name, oauth_token)
 
 
 def push_dataset(
@@ -371,6 +378,7 @@ def push_dataset(
     response_instruction_response: str,
     prompt_template: str,
     structured_output: dict,
+    pipeline_code: str,
     oauth_token: Union[gr.OAuthToken, None] = None,
     progress=gr.Progress(),
 ) -> pd.DataFrame:
@@ -385,7 +393,9 @@ def push_dataset(
         structured_output=structured_output,
         num_rows=num_rows,
     )
-    push_dataset_to_hub(dataframe, org_name, repo_name, oauth_token, private)
+    push_dataset_to_hub(
+        dataframe, org_name, repo_name, oauth_token, private, pipeline_code
+    )
     try:
         progress(0.1, desc="Setting up user and workspace")
         hf_user = HfApi().whoami(token=oauth_token.token)["name"]
@@ -854,6 +864,7 @@ with gr.Blocks() as app:
             response_instruction_response,
             prompt_template,
             structured_output,
+            pipeline_code,
         ],
         outputs=[success_message],
         show_progress=True,

src/synthetic_dataset_generator/apps/sft.py

@@ -11,6 +11,7 @@ from huggingface_hub import HfApi
 
 from synthetic_dataset_generator.apps.base import (
     hide_success_message,
+    push_pipeline_code_to_hub,
     show_success_message,
     validate_argilla_user_workspace_dataset,
     validate_push_to_hub,
@@ -202,7 +203,14 @@ def generate_dataset(
     return dataframe
 
 
-def push_dataset_to_hub(dataframe, org_name, repo_name, oauth_token, private):
+def push_dataset_to_hub(
+    dataframe: pd.DataFrame,
+    org_name: str,
+    repo_name: str,
+    oauth_token: Union[gr.OAuthToken, None],
+    private: bool,
+    pipeline_code: str,
+):
     repo_id = validate_push_to_hub(org_name, repo_name)
     original_dataframe = dataframe.copy(deep=True)
     dataframe = convert_dataframe_messages(dataframe)
@@ -214,6 +222,7 @@ def push_dataset_to_hub(dataframe, org_name, repo_name, oauth_token, private):
         token=oauth_token.token,
         create_pr=False,
     )
+    push_pipeline_code_to_hub(pipeline_code, org_name, repo_name, oauth_token)
     return original_dataframe
 
 
@@ -225,6 +234,7 @@ def push_dataset(
     num_rows: int = 10,
     private: bool = False,
     temperature: float = 0.9,
+    pipeline_code: str = "",
     oauth_token: Union[gr.OAuthToken, None] = None,
     progress=gr.Progress(),
 ) -> pd.DataFrame:
@@ -234,7 +244,9 @@ def push_dataset(
         num_rows=num_rows,
         temperature=temperature,
     )
-    push_dataset_to_hub(dataframe, org_name, repo_name, oauth_token, private)
+    push_dataset_to_hub(
+        dataframe, org_name, repo_name, oauth_token, private, pipeline_code
+    )
     try:
         progress(0.1, desc="Setting up user and workspace")
         hf_user = HfApi().whoami(token=oauth_token.token)["name"]
@@ -528,6 +540,7 @@ with gr.Blocks() as app:
                 num_rows,
                 private,
                 temperature,
+                pipeline_code,
             ],
             outputs=[success_message],
             show_progress=True,

src/synthetic_dataset_generator/utils.py

@@ -39,14 +39,13 @@ def list_orgs(oauth_token: Union[OAuthToken, None] = None):
             organizations = [org for org in organizations if org != data["name"]]
             organizations = [data["name"]] + organizations
     except Exception as e:
-        data = whoami(oauth_token.token)
         warnings.warn(str(e))
         gr.Info(
             "Your user token does not have the necessary permissions to push to organizations."
             "Please check your OAuth permissions in https://huggingface.co/settings/connected-applications."
             "Update yout token permissions to include repo.write: https://huggingface.co/settings/tokens."
         )
-        return [data["name"]]
+        return []
 
     return organizations