Merge branch 'main' of https://huggingface.co/spaces/HuggingFaceH4/syngen

app.py

@@ -8,7 +8,7 @@ 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
+from typing import Optional, Union
 
 """
  Still TODO:
@@ -17,6 +17,16 @@ from typing import Optional, List, Tuple, Union
  - validate max model params
 """
 
+SUPPORTED_MODELS = [
+    "Qwen/Qwen3-4B-Instruct-2507",
+    "Qwen/Qwen3-30B-A3B-Instruct-2507",
+    "meta-llama/Llama-3.2-1B-Instruct",
+    "meta-llama/Llama-3.2-3B-Instruct",
+    "baidu/ERNIE-4.5-21B-A3B-Thinking",
+    "LLM360/K2-Think",
+    "openai/gpt-oss-20b",
+]
+
 
 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."""
@@ -49,7 +59,8 @@ class GenerationStatus(Enum):
     FAILED = "FAILED"
 
 
-MAX_SAMPLES = 10000  # max number of samples in the input dataset
+MAX_SAMPLES_PRO = 10000  # max number of samples for PRO/Enterprise users
+MAX_SAMPLES_FREE = 100   # max number of samples for free users
 MAX_TOKENS = 8192
 MAX_MODEL_PARAMS = 20_000_000_000  # 20 billion parameters (for now)
 
@@ -79,7 +90,7 @@ class GenerationRequest:
     private: bool = False
     num_retries: int = 0
 
-def validate_request(request: GenerationRequest) -> GenerationRequest:
+def validate_request(request: GenerationRequest, oauth_token: Optional[Union[gr.OAuthToken, str]] = None) -> GenerationRequest:
     # checks that the request is valid
     # - input dataset exists and can be accessed with the provided token
     try:
@@ -101,8 +112,17 @@ def validate_request(request: GenerationRequest) -> GenerationRequest:
 
 
 
-    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 user tier and apply appropriate limits
+    # Anonymous users (oauth_token is None) are treated as free tier
+    is_pro = verify_pro_status(oauth_token) if oauth_token else False
+    max_samples = MAX_SAMPLES_PRO if is_pro else MAX_SAMPLES_FREE
+    
+    if request.num_output_examples > max_samples:
+        if oauth_token is None:
+            user_tier = "non-signed-in"
+        else:
+            user_tier = "PRO/Enterprise" if is_pro else "free"
+        raise Exception(f"Requested number of output examples {request.num_output_examples} exceeds the max limit of {max_samples} for {user_tier} users.")
 
     # check the prompt column exists in the dataset
     if request.prompt_column not in input_dataset_info.features:
@@ -117,8 +137,34 @@ def validate_request(request: GenerationRequest) -> GenerationRequest:
     try:
         output_dataset_info = get_dataset_infos(request.output_dataset_name, token=request.output_dataset_token)
         raise Exception(f"Output dataset {request.output_dataset_name} already exists. Please choose a different name.")
-    except Exception as e:
+    except Exception:
         pass  # dataset does not exist, which is expected
+    
+    # check the output dataset name doesn't already exist in the database
+    try:
+        url = os.getenv("SUPABASE_URL")
+        key = os.getenv("SUPABASE_KEY")
+        
+        if url and key:
+            supabase = create_client(
+                url,
+                key,
+                options=ClientOptions(
+                    postgrest_client_timeout=10,
+                    storage_client_timeout=10,
+                    schema="public",
+                )
+            )
+            
+            existing_request = supabase.table("gen-requests").select("id").eq("output_dataset_name", request.output_dataset_name).execute()
+            if existing_request.data:
+                raise Exception(f"Output dataset {request.output_dataset_name} is already being generated or has been requested. Please choose a different name.")
+    except Exception as e:
+        # If it's our custom exception about dataset already existing, re-raise it
+        if "already being generated" in str(e):
+            raise e
+        # Otherwise, ignore database connection errors and continue
+        pass
 
     # check the models exists
     try:
@@ -145,7 +191,7 @@ def validate_request(request: GenerationRequest) -> GenerationRequest:
     if request.top_p < 0.0 or request.top_p > 1.0:
         raise Exception("Top P must be between 0.0 and 1.0")   
     
-    # check valid email address TODO: use py3-validate-email https://stackoverflow.com/questions/8022530/how-to-check-for-valid-email-address
+    # check valid email address TODO: could use py3-validate-email https://stackoverflow.com/questions/8022530/how-to-check-for-valid-email-address
     if "@" not in request.email or "." not in request.email.split("@")[-1]:
         raise Exception("Invalid email address")
     
@@ -197,14 +243,58 @@ def add_request_to_db(request: GenerationRequest):
 
 
 
+def get_generation_stats_safe():
+    """Safely fetch generation request statistics with proper error handling"""
+    try:
+        url = os.getenv("SUPABASE_URL")
+        key = os.getenv("SUPABASE_KEY")
+        
+        if not url or not key:
+            raise Exception("Missing SUPABASE_URL or SUPABASE_KEY environment variables")
+        
+        supabase = create_client(
+            url,
+            key,
+            options=ClientOptions(
+                postgrest_client_timeout=10,
+                storage_client_timeout=10,
+                schema="public",
+            )
+        )
+        
+        # Fetch data excluding sensitive token fields
+        response = supabase.table("gen-requests").select(
+            "id, created_at, status, input_dataset_name, input_dataset_config, "
+            "input_dataset_split, output_dataset_name, prompt_column, "
+            "model_name_or_path, model_revision, max_tokens, temperature, "
+            "top_k, top_p, username, num_output_examples, private"
+        ).order("created_at", desc=True).limit(50).execute()
+        
+        return {"status": "success", "data": response.data}
+        
+    except Exception as e:
+        return {"status": "error", "message": str(e), "data": []}
+
+
+# Old commented code removed - replaced with DatabaseManager and get_generation_stats_safe()
 
 
 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")
+        gr.HTML("<h3 style='text-align:center'>Generate synthetic data with AI models. Free to use! Sign up for PRO benefits (10k samples vs 100). <a href='http://huggingface.co/subscribe/pro?source=synthetic-data-universe' target='_blank'>Upgrade to PRO</a></h3>", elem_id="sub_title")
+        
+        # Add sign-in button at the top
+        with gr.Row():
+            gr.Markdown("")  # Empty space for alignment
+            login_button = gr.LoginButton(value="🔑 Sign in", size="sm")
+            gr.Markdown("")  # Empty space for alignment
         
-        pro_message = gr.Markdown(visible=False)
+        signin_message = gr.Markdown("## 🔑 Sign In Required\n\nPlease sign in with your Hugging Face account to access the synthetic data generation service. Click the **Sign in** button above to continue.", visible=True)
         main_interface = gr.Column(visible=False)
+        
+        # Store the current oauth token for use in submit_request
+        current_oauth_token = gr.State(None)
+        
         with main_interface:
             with gr.Group():
                 with gr.Row():
@@ -214,7 +304,7 @@ def main():
                     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.Accordion("More Information", open=False):
                 with gr.Row():
                     gr.Markdown("""
                     **How it works:**
@@ -227,59 +317,169 @@ def main():
                     
                     **Requirements:**
                     - Input dataset must be publicly accessible
-                    - Model must be accessible (public and not gated)
+                    - Model must be publicly accessible (and not gated)
                     - Maximum 10,000 samples per dataset
-                    - Maximum of 8192 generation tokens
+                    - Maximum of 8192 generated tokens
                     """)
             
-            with gr.Group():
-                gr.Markdown("##  Dataset information")
-                with gr.Column():
-                    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.Tabs():
+                with gr.TabItem("Generate Synthetic Data"):
+                    with gr.Group():
+                        gr.Markdown("##  Model information")
+                        with gr.Column():
+                            with gr.Row():
+                                model_name_or_path = gr.Dropdown(
+                                    choices=SUPPORTED_MODELS,
+                                    label="Select Model",
+                                    value="Qwen/Qwen3-4B-Instruct-2507",
+                                    info="Choose from popular instruction-tuned models under 40B parameters"
+                                )
+                            # 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("##  Dataset information")
+                        # Dynamic user limit info - default to anonymous user
+                        user_limit_info = gr.Markdown(value="👤 **Anonymous User**: You can generate up to 100 samples per request. Use the sign-in button above for PRO benefits (10,000 samples).", visible=True)
+                        with gr.Row():
+                            with gr.Column():
+                                input_dataset_name = gr.Textbox(label="Input Dataset Name", placeholder="e.g., simplescaling/s1K-1.1")
+                                prompt_column = gr.Textbox(label="Prompt Column", placeholder="e.g., text, prompt, question", value=None, interactive=False, info="Click Load Info to populate")
+                                
+                            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'", value=None, interactive=False, info="Click Load Info to populate")
                         
-                    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.Column():
+                                input_dataset_config = gr.Dropdown(label="Dataset Config", choices=[], value=None, interactive=False, info="Click Load Info to populate")
+                                prompt_column = gr.Dropdown(label="Prompt Column", choices=[], value=None, interactive=False, info="Click Load Info to populate")
+                                
+                            with gr.Column():
+                                input_dataset_split = gr.Dropdown(label="Dataset Split", choices=[], value=None, interactive=False, info="Click Load Info to populate")
+                                num_output_samples = gr.Slider(label="Number of samples, leave as '0' for all", value=0, minimum=0, maximum=MAX_SAMPLES_FREE, step=1, interactive=False, info="Click Load Info to populate")
+                    
+                    gr.Markdown("### Generation Parameters")
                     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():
-                    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.Column():
+                            with gr.Row():
+                                max_tokens = gr.Slider(label="Max Tokens", value=1024, 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 (still TODO)")
-                with gr.Row():
-                    with gr.Column():
+                    with gr.Group():
+                        gr.Markdown("##  User Information, for notification when your job is completed (still TODO)")
                         with gr.Row():
-                            email = gr.Textbox(label="Email", placeholder="[email protected]")
-                        # 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")
-
-            submit_btn = gr.Button("Submit Generation Request", variant="primary")
-            output_status = gr.Textbox(label="Status", interactive=False)
+                            with gr.Column():
+                                with gr.Row():
+                                    email = gr.Textbox(label="Email", placeholder="[email protected]")
+                                # 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")
+
+                    submit_btn = gr.Button("Submit Generation Request", variant="primary")
+                    output_status = gr.Textbox(label="Status", interactive=False)
+                
+                with gr.TabItem("Statistics Dashboard"):
+                    gr.Markdown("## Generation Requests Statistics")
+                    gr.Markdown("📊 View recent synthetic data generation requests and their status.")
+                    
+                    with gr.Row():
+                        refresh_stats_btn = gr.Button("🔄 Refresh Statistics", size="sm", variant="secondary")
+                        clear_stats_btn = gr.Button("🗑️ Clear Display", size="sm")
+                    
+                    stats_status = gr.Markdown("Click 'Refresh Statistics' to load recent generation requests.", visible=True)
+                    
+                    stats_dataframe = gr.Dataframe(
+                        headers=["ID", "Created", "Status", "Input Dataset", "Output Dataset", "Model", "Samples", "User"],
+                        datatype=["str", "str", "str", "str", "str", "str", "number", "str"],
+                        interactive=False,
+                        wrap=True,
+                        value=[],
+                        label="Recent Generation Requests (Last 50)",
+                        visible=False
+                    )
+                    
+                    def load_statistics():
+                        """Load and format statistics data"""
+                        try:
+                            # Use the new safe database function
+                            result = get_generation_stats_safe()
+                            
+                            if result["status"] == "error":
+                                return (
+                                    f"❌ **Error loading statistics**: {result['message']}",
+                                    gr.update(visible=False),
+                                    gr.update(visible=True)
+                                )
+                            
+                            data = result["data"]
+                            if not data:
+                                return (
+                                    "📝 **No data found**: The database appears to be empty or the table doesn't exist yet.",
+                                    gr.update(visible=False),
+                                    gr.update(visible=True)
+                                )
+                            
+                            # Format data for display
+                            formatted_data = []
+                            for item in data:
+                                # Format timestamp
+                                created_at = item.get('created_at', 'Unknown')
+                                if created_at and created_at != 'Unknown':
+                                    try:
+                                        from datetime import datetime
+                                        dt = datetime.fromisoformat(created_at.replace('Z', '+00:00'))
+                                        created_at = dt.strftime('%Y-%m-%d %H:%M')
+                                    except:
+                                        pass
+                                
+                                formatted_data.append([
+                                    str(item.get('id', ''))[:8] + "..." if len(str(item.get('id', ''))) > 8 else str(item.get('id', '')),
+                                    created_at,
+                                    item.get('status', 'Unknown'),
+                                    (item.get('input_dataset_name', '')[:30] + "...") if len(item.get('input_dataset_name', '')) > 30 else item.get('input_dataset_name', ''),
+                                    (item.get('output_dataset_name', '')[:30] + "...") if len(item.get('output_dataset_name', '')) > 30 else item.get('output_dataset_name', ''),
+                                    (item.get('model_name_or_path', '')[:25] + "...") if len(item.get('model_name_or_path', '')) > 25 else item.get('model_name_or_path', ''),
+                                    item.get('num_output_examples', 0),
+                                    item.get('username', 'Anonymous')
+                                ])
+                            
+                            return (
+                                f"✅ **Statistics loaded successfully**: Found {len(formatted_data)} recent requests.",
+                                gr.update(value=formatted_data, visible=True),
+                                gr.update(visible=True)
+                            )
+                            
+                        except Exception as e:
+                            return (
+                                f"❌ **Unexpected error**: {str(e)}",
+                                gr.update(visible=False),
+                                gr.update(visible=True)
+                            )
+                    
+                    def clear_statistics():
+                        """Clear the statistics display"""
+                        return (
+                            "Click 'Refresh Statistics' to load recent generation requests.",
+                            gr.update(value=[], visible=False),
+                            gr.update(visible=True)
+                        )
+                    
+                    # Connect buttons to functions
+                    refresh_stats_btn.click(
+                        load_statistics,
+                        outputs=[stats_status, stats_dataframe, stats_status]
+                    )
+                    
+                    clear_stats_btn.click(
+                        clear_statistics,
+                        outputs=[stats_status, stats_dataframe, stats_status]
+                    )
 
-            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):
+            def submit_request(input_dataset_name, input_split, input_dataset_config, output_dataset_name, prompt_col, model_name, sys_prompt,
+                                max_tok, temp, top_k_val, top_p_val, email_addr, num_output_samples, oauth_token=None):
 
                 MASTER_ORG = "synthetic-data-universe/"
                 model_token = False # This is currently not supported
@@ -297,7 +497,7 @@ def main():
                         output_dataset_name=MASTER_ORG + output_dataset_name,
                         prompt_column=prompt_col,
                         model_name_or_path=model_name,
-                        model_revision=model_rev,
+                        model_revision="main",
                         model_token=model_token,
                         system_prompt=sys_prompt if sys_prompt else None,
                         max_tokens=int(max_tok),
@@ -312,7 +512,7 @@ def main():
                     )
                     
                     # check the input dataset exists and can be accessed with the provided token
-                    request = validate_request(request)
+                    request = validate_request(request, oauth_token)
                     add_request_to_db(request)
 
                     return "Request submitted successfully!"
@@ -322,25 +522,57 @@ def main():
             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],
+                         system_prompt, max_tokens, temperature, top_k, top_p, email, num_output_samples, current_oauth_token],
                 outputs=output_status
             )
         
+        def update_user_limits(oauth_token):
+            if oauth_token is None:
+                return "👤 **Anonymous User**: You can generate up to 100 samples per request. Use the sign-in button above for PRO benefits (10,000 samples)."
+            
+            is_pro = verify_pro_status(oauth_token)
+            if is_pro:
+                return "✨ **PRO User**: You can generate up to 10,000 samples per request."
+            else:
+                return "👤 **Free User**: You can generate up to 100 samples per request. [Upgrade to PRO](http://huggingface.co/subscribe/pro?source=synthetic-data-universe) for 10,000 samples."
+
         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)
+            # Require users to be signed in
+            if oauth_token is None:
+                # User is not signed in - show sign-in prompt, hide main interface
+                return (
+                    gr.update(visible=False),  # main_interface
+                    gr.update(visible=True),   # signin_message
+                    oauth_token,               # current_oauth_token
+                    "",                        # user_limit_info (empty when not signed in)
+                    gr.update(),               # num_output_samples (no change)
+                    gr.update(value="🔑 Sign in")  # login_button
+                )
             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)"
+                # User is signed in - show main interface, hide sign-in prompt
+                limit_msg = update_user_limits(oauth_token)
+                is_pro = verify_pro_status(oauth_token)
+                max_samples = MAX_SAMPLES_PRO if is_pro else MAX_SAMPLES_FREE
+                
+                if is_pro:
+                    button_text = f"✨ Signed in as PRO ({profile.name if profile else 'User'})"
+                else:
+                    button_text = f"👤 Signed in as {profile.name if profile else 'User'}"
+                
+                return (
+                    gr.update(visible=True),   # main_interface
+                    gr.update(visible=False),  # signin_message
+                    oauth_token,               # current_oauth_token
+                    limit_msg,                 # user_limit_info
+                    gr.update(maximum=max_samples),  # num_output_samples
+                    gr.update(value=button_text)     # login_button
                 )
-                return gr.update(visible=False), gr.update(visible=True, value=message)
 
-        login_button = gr.LoginButton() # this is required or AUTH will not work
 
-        demo.load(control_access, inputs=None, outputs=[main_interface, pro_message])
+        # Handle login state changes - LoginButton automatically handles auth state changes
+        # The demo.load will handle both initial load and auth changes
+        
+        demo.load(control_access, inputs=None, outputs=[main_interface, signin_message, current_oauth_token, user_limit_info, num_output_samples, login_button])
         demo.queue(max_size=None, default_concurrency_limit=None).launch(show_error=True)
 
 if __name__ == "__main__":