Update src/distilabel_dataset_generator/apps/sft.py

src/distilabel_dataset_generator/apps/sft.py

@@ -18,7 +18,6 @@ from src.distilabel_dataset_generator.pipelines.sft import (
 from src.distilabel_dataset_generator.utils import (
     get_login_button,
     get_org_dropdown,
-    get_token,
 )
 
 
@@ -66,8 +65,8 @@ def generate_dataset(
     private=True,
     org_name=None,
     repo_name=None,
-    token=None,
     progress=gr.Progress(),
+    oauth_token: Union[gr.OAuthToken, None]
 ):
     repo_id = (
         f"{org_name}/{repo_name}"
@@ -79,12 +78,6 @@ def generate_dataset(
             raise gr.Error(
                 "Please provide a repo_name and org_name to push the dataset to."
             )
-        try:
-            whoami(token=token)
-        except Exception:
-            raise gr.Error(
-                "Provide a Hugging Face token with write access to the organization you want to push the dataset to."
-            )
 
     if num_turns > 4:
         num_turns = 4
@@ -133,7 +126,7 @@ def generate_dataset(
             repo_id=repo_id,
             private=private,
             include_script=False,
-            token=token,
+            token=oauth_token.token,
         )
 
     # If not pushing to hub generate the dataset directly
@@ -215,7 +208,7 @@ with gr.Blocks(
         # Add a header for the full dataset generation section
         gr.Markdown("## Generate full dataset")
         gr.Markdown(
-            "Once you're satisfied with the sample, generate a larger dataset and push it to the hub. Get <a href='https://huggingface.co/settings/tokens' target='_blank'>a Hugging Face token</a> with write access to the organization you want to push the dataset to. A OAuth login resets the session state, so watch out not to lose your generated system prompt!"
+            "Once you're satisfied with the sample, generate a larger dataset and push it to the Hub."
         )
 
         with gr.Column() as push_to_hub_ui:
@@ -281,7 +274,6 @@ with gr.Blocks(
             private,
             org_name,
             repo_name,
-            hf_token,
         ],
         outputs=[table],
         show_progress=True,
@@ -300,5 +292,4 @@ with gr.Blocks(
             label="Distilabel Pipeline Code",
         )
 
-    app.load(get_token, outputs=[hf_token])
     app.load(get_org_dropdown, outputs=[org_name])