add auth

app.py

@@ -7,6 +7,8 @@ from supabase.client import ClientOptions
 from enum import Enum
 from datasets import get_dataset_infos
 from transformers import AutoConfig
+from huggingface_hub import whoami
+from typing import Optional, List, Tuple, Union
 
 """
  Still TODO:
@@ -16,6 +18,30 @@ from transformers import AutoConfig
 """
 
 
+def verify_pro_status(token: Optional[Union[gr.OAuthToken, str]]) -> bool:
+    """Verifies if the user is a Hugging Face PRO user or part of an enterprise org."""
+    if not token:
+        return False
+    
+    if isinstance(token, gr.OAuthToken):
+        token_str = token.token
+    elif isinstance(token, str):
+        token_str = token
+    else:
+        return False
+    
+    try:
+        user_info = whoami(token=token_str)
+        return (
+            user_info.get("isPro", False) or
+            any(org.get("isEnterprise", False) for org in user_info.get("orgs", []))
+        )
+    except Exception as e:
+        print(f"Could not verify user's PRO/Enterprise status: {e}")
+        return False
+
+
+
 class GenerationStatus(Enum):
     PENDING = "PENDING"
     RUNNING = "RUNNING"
@@ -64,17 +90,23 @@ def validate_request(request: GenerationRequest) -> GenerationRequest:
     # check that the input dataset split exists
     if request.input_dataset_split not in input_dataset_info.splits:
         raise Exception(f"Dataset split {request.input_dataset_split} does not exist in dataset {request.input_dataset_name}. Available splits: {list(input_dataset_info.splits.keys())}")
-    
-    # check that the number of samples is less than MAX_SAMPLES
-    if input_dataset_info.splits[request.input_dataset_split].num_examples > MAX_SAMPLES:
+
+    # if num_output_examples is 0, set it to the number of examples in the input dataset split
+    if request.num_output_examples == 0:
         request.num_output_examples = input_dataset_info.splits[request.input_dataset_split].num_examples
-        raise Exception(f"Dataset split {request.input_dataset_split} in dataset {request.input_dataset_name} exceeds max sample limit of {MAX_SAMPLES}.")
+    else: 
+        if request.num_output_examples > input_dataset_info.splits[request.input_dataset_split].num_examples:
+            raise Exception(f"Requested number of output examples {request.num_output_examples} exceeds the number of examples in the input dataset split {input_dataset_info.splits[request.input_dataset_split].num_examples}.")
+        request.input_dataset_split = f"{request.input_dataset_split}[:{request.num_output_examples}]"
+
+
+    if request.num_output_examples > MAX_SAMPLES:
+        raise Exception(f"Requested number of output examples {request.num_output_examples} exceeds the max limit of {MAX_SAMPLES}.")
 
     # check the prompt column exists in the dataset
     if request.prompt_column not in input_dataset_info.features:
         raise Exception(f"Prompt column {request.prompt_column} does not exist in dataset {request.input_dataset_name}. Available columns: {list(input_dataset_info.features.keys())}")
 
-
     # check the models exists
     try:
         model_config = AutoConfig.from_pretrained(request.model_name_or_path, revision=request.model_revision, token=request.model_token)
@@ -146,127 +178,151 @@ def add_request_to_db(request: GenerationRequest):
         raise Exception("Failed to add request to database")
 
 
-def create_gradio_interface():
-    with gr.Blocks(title="Synthetic Data Generation") as interface:
-        with gr.Group():
-            with gr.Row():
-                gr.Markdown("# Synthetic Data Generation Request")    
-            with gr.Row():
-                gr.Markdown("""
-                Welcome to the Synthetic Data Generation service! This tool allows you to generate synthetic data using large language models. Generation is FREE for Hugging Face PRO users and uses idle GPUs on the HF science cluster.\n
-                """)
-        with gr.Group():
-            with gr.Row():
-                gr.Markdown("""
-                **How it works:**
-                1. Provide an input dataset with prompts
-                2. Select a public language model for generation
-                3. Configure generation parameters
-                4. Submit your request and receive generated data
-                """)
-                gr.Markdown("""              
-                
-                **Requirements:**
-                - Input dataset must be publicly accessible
-                - Output dataset repository must exist and you must have write access
-                - Model must be accessible (public or with valid token)
-                - Maximum 10,000 samples per dataset
-                - Maximum of 8192 generation tokens
-                """)
+
+
+
+
+def main():
+    with gr.Blocks(title="Synthetic Data Generation") as demo:
+        gr.HTML("<h3 style='text-align:center'>Hugging Face PRO users can use the Synthetic generation service. <a href='http://huggingface.co/subscribe/pro?source=synthetic-data-universe' target='_blank'>Subscribe to PRO</a></h3>", elem_id="sub_title")
         
-        with gr.Group():
-            gr.Markdown("##  Dataset information")
-            with gr.Column():
+        pro_message = gr.Markdown(visible=False)
+        main_interface = gr.Column(visible=False)
+        with main_interface:
+            with gr.Group():
                 with gr.Row():
-                    input_dataset_name = gr.Textbox(label="Input Dataset Name", placeholder="e.g., simplescaling/s1K-1.1")
-                    input_dataset_split = gr.Textbox(label="Input Dataset Split", value="train", placeholder="e.g., train, test, validation")
-                    input_dataset_config = gr.Textbox(label="Input Dataset Config", value="default", placeholder="e.g., default, custom")
-                    prompt_column = gr.Textbox(label="Prompt Column", placeholder="e.g., text, prompt, question")
-                with gr.Column():
-                    output_dataset_name = gr.Textbox(label="Output Dataset Name", placeholder="e.g., my-generated-dataset, must be unique. Will be created under the org 'synthetic-data-universe'")
-        with gr.Group():
-            gr.Markdown("##  Model information")
-            with gr.Column():
+                    gr.Markdown("# Synthetic Data Generation Request")    
+                with gr.Row():
+                    gr.Markdown("""
+                    Welcome to the Synthetic Data Generation service! This tool allows you to generate synthetic data using large language models. Generation is FREE for Hugging Face PRO users and uses idle GPUs on the HF science cluster.\n
+                    Outputs from this service will be PUBLIC and available on the Hugging Face Hub under the organization [synthetic-data-universe](https://huggingface.co/synthetic-data-universe).\n
+                    """)
+            with gr.Group():
                 with gr.Row():
-                    model_name_or_path = gr.Textbox(label="Model Name or Path", placeholder="e.g., Qwen/Qwen3-4B-Instruct-2507")
-                    model_revision = gr.Textbox(label="Model Revision", value="main", placeholder="e.g., main, v1.0")
-                # model_token = gr.Textbox(label="Model Token (Optional)", type="password", placeholder="Your HF token with read/write access to the model...")
-        with gr.Group():
-            gr.Markdown("##  Generation Parameters")
-            with gr.Row():
+                    gr.Markdown("""
+                    **How it works:**
+                    1. Provide an input dataset with prompts
+                    2. Select a public language model for generation
+                    3. Configure generation parameters
+                    4. Submit your request.
+                    """)
+                    gr.Markdown("""              
+                    
+                    **Requirements:**
+                    - Input dataset must be publicly accessible (for now)
+                    - Model must be accessible (public and note gated, for now)
+                    - Maximum 10,000 samples per dataset (for now)
+                    - Maximum of 8192 generation tokens (for now)
+                    """)
+            
+            with gr.Group():
+                gr.Markdown("##  Dataset information")
                 with gr.Column():
                     with gr.Row():
-                        max_tokens = gr.Slider(label="Max Tokens", value=512, minimum=256, maximum=MAX_TOKENS, step=256)
-                        temperature = gr.Slider(label="Temperature", minimum=0.0, maximum=2.0, value=0.7, step=0.1)
-                    with gr.Row():
-                        top_k = gr.Slider(label="Top K", value=50, minimum=5, maximum=100, step=5)
-                        top_p = gr.Slider(label="Top P", minimum=0.0, maximum=1.0, value=0.95, step=0.05)
-                    with gr.Row():
-                        system_prompt = gr.Textbox(label="System Prompt (Optional)", lines=3, placeholder="Optional system prompt... e.g., You are a helpful assistant.")
-
-        with gr.Group():
-            gr.Markdown("##  User Information, for notification when your job is completed")
-            with gr.Row():
+                        input_dataset_name = gr.Textbox(label="Input Dataset Name", placeholder="e.g., simplescaling/s1K-1.1")
+                        input_dataset_split = gr.Textbox(label="Input Dataset Split", value="train", placeholder="e.g., train, test, validation")
+                        input_dataset_config = gr.Textbox(label="Input Dataset Config", value="default", placeholder="e.g., default, custom")
+                        prompt_column = gr.Textbox(label="Prompt Column", placeholder="e.g., text, prompt, question")
+                        
+                    with gr.Column():
+                        with gr.Row():
+                            output_dataset_name = gr.Textbox(label="Output Dataset Name", placeholder="e.g., my-generated-dataset, must be unique. Will be created under the org 'synthetic-data-universe'")
+                            num_output_samples = gr.Slider(label="Number of samples, leave as '0' for all", value=0, minimum=0, maximum=MAX_SAMPLES, step=1)
+            with gr.Group():
+                gr.Markdown("##  Model information")
                 with gr.Column():
                     with gr.Row():
-                        email = gr.Textbox(label="Email", placeholder="your.email@example.com")
-                    # with gr.Row():
-                        # input_dataset_token = gr.Textbox(label="Input dataset token", type="password", placeholder="Your HF token with read access to the input dataset, leave blank if public dataset")
-                        # output_dataset_token = gr.Textbox(label="Output dataset token", type="password", placeholder="Your HF token with write access to the output dataset")
+                        model_name_or_path = gr.Textbox(label="Model Name or Path", placeholder="e.g., Qwen/Qwen3-4B-Instruct-2507")
+                        model_revision = gr.Textbox(label="Model Revision", value="main", placeholder="e.g., main, v1.0")
+                    # model_token = gr.Textbox(label="Model Token (Optional)", type="password", placeholder="Your HF token with read/write access to the model...")
+            with gr.Group():
+                gr.Markdown("##  Generation Parameters")
+                with gr.Row():
+                    with gr.Column():
+                        with gr.Row():
+                            max_tokens = gr.Slider(label="Max Tokens", value=512, minimum=256, maximum=MAX_TOKENS, step=256)
+                            temperature = gr.Slider(label="Temperature", minimum=0.0, maximum=2.0, value=0.7, step=0.1)
+                        with gr.Row():
+                            top_k = gr.Slider(label="Top K", value=50, minimum=5, maximum=100, step=5)
+                            top_p = gr.Slider(label="Top P", minimum=0.0, maximum=1.0, value=0.95, step=0.05)
+                        with gr.Row():
+                            system_prompt = gr.Textbox(label="System Prompt (Optional)", lines=3, placeholder="Optional system prompt... e.g., You are a helpful assistant.")
 
-        submit_btn = gr.Button("Submit Generation Request", variant="primary")
-        output_status = gr.Textbox(label="Status", interactive=False)
+            with gr.Group():
+                gr.Markdown("##  User Information, for notification when your job is completed (still TODO)")
+                with gr.Row():
+                    with gr.Column():
+                        with gr.Row():
+                            email = gr.Textbox(label="Email", placeholder="your.email@example.com")
+                        # with gr.Row():
+                            # input_dataset_token = gr.Textbox(label="Input dataset token", type="password", placeholder="Your HF token with read access to the input dataset, leave blank if public dataset")
+                            # output_dataset_token = gr.Textbox(label="Output dataset token", type="password", placeholder="Your HF token with write access to the output dataset")
 
-        def submit_request(input_dataset_name, input_split, input_dataset_config, output_dataset_name, prompt_col, model_name, model_rev, sys_prompt,
-                            max_tok, temp, top_k_val, top_p_val, email_addr):
+            submit_btn = gr.Button("Submit Generation Request", variant="primary")
+            output_status = gr.Textbox(label="Status", interactive=False)
 
-            MASTER_ORG = "synthetic-data-universe/"
-            model_token = None # This is currently not supported
-            input_dataset_token = None # This is currently not supported
-            output_dataset_token = os.getenv("OUTPUT_DATASET_TOKEN")
-            try:
-                request = GenerationRequest(
-                    id="",  # Will be generated when adding to the database
-                    created_at="",  # Will be set when adding to the database
-                    status=GenerationStatus.PENDING,
-                    input_dataset_name=input_dataset_name,
-                    input_dataset_split=input_split,
-                    input_dataset_config=input_dataset_config,
-                    output_dataset_name=MASTER_ORG + output_dataset_name,
-                    prompt_column=prompt_col,
-                    model_name_or_path=model_name,
-                    model_revision=model_rev,
-                    model_token=model_token if model_token else None,
-                    system_prompt=sys_prompt if sys_prompt else None,
-                    max_tokens=int(max_tok),
-                    temperature=temp,
-                    top_k=int(top_k_val),
-                    top_p=top_p_val,
-                    input_dataset_token=input_dataset_token if input_dataset_token else None,
-                    output_dataset_token=output_dataset_token,
-                    num_output_examples=0,  # will be set after validating the input dataset
-                    username="user",
-                    email=email_addr
-                )
+            def submit_request(input_dataset_name, input_split, input_dataset_config, output_dataset_name, prompt_col, model_name, model_rev, sys_prompt,
+                                max_tok, temp, top_k_val, top_p_val, email_addr, num_output_samples):
+
+                MASTER_ORG = "synthetic-data-universe/"
+                model_token = None # This is currently not supported
+                input_dataset_token = None # This is currently not supported
+                output_dataset_token = os.getenv("OUTPUT_DATASET_TOKEN")
                 
-                # check the input dataset exists and can be accessed with the provided token
-                request = validate_request(request)
-                add_request_to_db(request)
+                try:
+                    request = GenerationRequest(
+                        id="",  # Will be generated when adding to the database
+                        created_at="",  # Will be set when adding to the database
+                        status=GenerationStatus.PENDING,
+                        input_dataset_name=input_dataset_name,
+                        input_dataset_split=input_split,
+                        input_dataset_config=input_dataset_config,
+                        output_dataset_name=MASTER_ORG + output_dataset_name,
+                        prompt_column=prompt_col,
+                        model_name_or_path=model_name,
+                        model_revision=model_rev,
+                        model_token=model_token if model_token else None,
+                        system_prompt=sys_prompt if sys_prompt else None,
+                        max_tokens=int(max_tok),
+                        temperature=temp,
+                        top_k=int(top_k_val),
+                        top_p=top_p_val,
+                        input_dataset_token=input_dataset_token if input_dataset_token else None,
+                        output_dataset_token=output_dataset_token,
+                        num_output_examples=num_output_samples,  # will be set after validating the input dataset
+                        username="user",
+                        email=email_addr
+                    )
+                    
+                    # check the input dataset exists and can be accessed with the provided token
+                    request = validate_request(request)
+                    add_request_to_db(request)
 
-                return "Request submitted successfully!"
-            except Exception as e:
-                return f"Error: {str(e)}"
+                    return "Request submitted successfully!"
+                except Exception as e:
+                    return f"Error: {str(e)}"
+            
+            submit_btn.click(
+                submit_request,
+                inputs=[input_dataset_name, input_dataset_split, input_dataset_config, output_dataset_name, prompt_column, model_name_or_path,
+                        model_revision, system_prompt, max_tokens, temperature, top_k, top_p, email, num_output_samples],
+                outputs=output_status
+            )
         
-        submit_btn.click(
-            submit_request,
-            inputs=[input_dataset_name, input_dataset_split, input_dataset_config, output_dataset_name, prompt_column, model_name_or_path,
-                    model_revision, system_prompt, max_tokens, temperature, top_k, top_p,
-                     email],
-            outputs=output_status
-        )
-    
-    return interface
+        def control_access(profile: Optional[gr.OAuthProfile] = None, oauth_token: Optional[gr.OAuthToken] = None):
+            if not profile: return gr.update(visible=False), gr.update(visible=False)
+            if verify_pro_status(oauth_token): return gr.update(visible=True), gr.update(visible=False)
+            else:
+                message = (
+                    "## ✨ Exclusive Access for PRO Users\n\n"
+                    "Thank you for your interest! This app is available exclusively for our Hugging Face **PRO** members.\n\n"
+                    "To unlock this and many other cool stuff, please consider upgrading your account.\n\n"
+                    "### [**Become a PRO Today!**](http://huggingface.co/subscribe/pro?source=synthetic-data-universe)"
+                )
+                return gr.update(visible=False), gr.update(visible=True, value=message)    
+        
+        demo.load(control_access, inputs=None, outputs=[main_interface, pro_message]) 
+        demo.queue(max_size=None, default_concurrency_limit=None).launch(show_error=True)  
 
 if __name__ == "__main__":
-    app = create_gradio_interface()
-    app.launch()
+    main()