simplified

app.py

@@ -17,6 +17,7 @@ from utils.dataset import (
 )
 from utils.helpers import (
     check_token_availability,
+    format_dataset_help_text,
     format_token_status,
     get_inference_token,
     get_org_token,
@@ -39,6 +40,31 @@ from ui.tab_testing import (
 # Handlers
 # ============================================================================
 
+def prepare_save_data(test_input, current_policy, parsed, model_choice, raw_response, 
+                      reasoning, reasoning_effort, max_tokens, temperature, top_p, 
+                      system_prompt_val, response_format_val):
+    """Prepare data dict for saving to dataset."""
+    categories_and_reasoning_text = format_categories_and_reasoning(parsed)
+    policy_violation = parsed.get("label", -1)
+    
+    return {
+        "input": test_input,
+        "policy_violation": policy_violation,
+        "categories_and_reasoning": categories_and_reasoning_text,
+        "policy": current_policy,
+        "model_selection": model_choice,
+        "raw_response": raw_response,
+        "reasoning_trace": reasoning or "",
+        "reasoning_effort": reasoning_effort or "",
+        "max_tokens": int(max_tokens),
+        "temperature": float(temperature),
+        "top_p": float(top_p),
+        "system_prompt": system_prompt_val or "",
+        "response_format": response_format_val or "",
+        "timestamp": datetime.now().isoformat(),
+    }
+
+
 def handle_run_test(test_input, current_policy, model_choice, reasoning_effort, max_tokens, temperature, top_p, system_prompt_val, response_format_val, save_mode, oauth_token: gr.OAuthToken | None = None):
     """Handle test execution."""
 
@@ -81,57 +107,25 @@ def handle_run_test(test_input, current_policy, model_choice, reasoning_effort,
         org_token = get_org_token()
         if org_token:
             try:
-                categories_and_reasoning_text = format_categories_and_reasoning(parsed)
-                policy_violation = parsed.get("label", -1)
-                
-                data = {
-                    "input": test_input,
-                    "policy_violation": policy_violation,
-                    "categories_and_reasoning": categories_and_reasoning_text,
-                    "policy": current_policy,
-                    "model_selection": model_choice,
-                    "raw_response": raw_response,
-                    "reasoning_trace": reasoning or "",
-                    "reasoning_effort": reasoning_effort or "",
-                    "max_tokens": int(max_tokens),
-                    "temperature": float(temperature),
-                    "top_p": float(top_p),
-                    "system_prompt": system_prompt_val or "",
-                    "response_format": response_format_val or "",
-                    "timestamp": datetime.now().isoformat(),
-                }
-                repo_id = get_roost_dataset_repo_id()
-                save_to_dataset(repo_id, org_token, data)
+                data = prepare_save_data(
+                    test_input, current_policy, parsed, model_choice, raw_response,
+                    reasoning, reasoning_effort, max_tokens, temperature, top_p,
+                    system_prompt_val, response_format_val
+                )
+                save_to_dataset(get_roost_dataset_repo_id(), org_token, data)
             except Exception as e:
-                # Log error but don't break test execution
                 print(f"Failed to save to ROOST dataset: {e}")
     elif save_mode == "Save to Private Dataset":
         personal_token, _ = get_personal_token(oauth_token)
         if personal_token:
             try:
-                categories_and_reasoning_text = format_categories_and_reasoning(parsed)
-                policy_violation = parsed.get("label", -1)
-                
-                data = {
-                    "input": test_input,
-                    "policy_violation": policy_violation,
-                    "categories_and_reasoning": categories_and_reasoning_text,
-                    "policy": current_policy,
-                    "model_selection": model_choice,
-                    "raw_response": raw_response,
-                    "reasoning_trace": reasoning or "",
-                    "reasoning_effort": reasoning_effort or "",
-                    "max_tokens": int(max_tokens),
-                    "temperature": float(temperature),
-                    "top_p": float(top_p),
-                    "system_prompt": system_prompt_val or "",
-                    "response_format": response_format_val or "",
-                    "timestamp": datetime.now().isoformat(),
-                }
-                repo_id = get_dataset_repo_id(personal_token)
-                save_to_dataset(repo_id, personal_token, data)
+                data = prepare_save_data(
+                    test_input, current_policy, parsed, model_choice, raw_response,
+                    reasoning, reasoning_effort, max_tokens, temperature, top_p,
+                    system_prompt_val, response_format_val
+                )
+                save_to_dataset(get_dataset_repo_id(personal_token), personal_token, data)
             except Exception as e:
-                # Log error but don't break test execution
                 print(f"Failed to save to private dataset: {e}")
     
     return (
@@ -237,55 +231,19 @@ with gr.Blocks(title="Moderation Model Testing") as demo:
         outputs=model_info_display,
     )
     
-    # Token status update handler
-    def update_token_status(oauth_token: gr.OAuthToken | None = None):
-        """Update token status markdown when OAuth changes."""
-        return format_token_status(oauth_token)
-    
-    # Save mode help text update handler
-    def update_save_mode_help(oauth_token: gr.OAuthToken | None = None):
-        """Update save mode help text based on token availability."""
-        from ui.tab_testing import format_save_mode_help
-        has_personal, has_org = check_token_availability(oauth_token)
-        return format_save_mode_help(has_personal, has_org)
-    
-    # Dataset button state update handler
-    def update_dataset_button_states(oauth_token: gr.OAuthToken | None = None):
-        """Update dataset button states based on token availability."""
-        has_personal, has_org = check_token_availability(oauth_token)
-        
-        # Update help text
-        help_text = (
-            f"*Private Dataset: {'✅ Available' if has_personal else '❌ Requires personal token (OAuth login or .env)'}*\n"
-            f"*ROOST Dataset: {'✅ Available' if has_org else '⚠️ Can load if public, requires org token to save'}*"
-        )
-        
-        return (
-            gr.update(interactive=has_personal),  # refresh_private_btn
-            gr.update(interactive=True),  # refresh_roost_btn (can load if public)
-            help_text,  # dataset_help_text
-        )
-    
     # Combined handler for login button click - updates all token-dependent UI
     def handle_login_click(oauth_token: gr.OAuthToken | None = None):
         """Handle login button click and update all token-dependent UI."""
-        token_status = format_token_status(oauth_token)
-        
         from ui.tab_testing import format_save_mode_help
-        has_personal, has_org = check_token_availability(oauth_token)
-        save_help = format_save_mode_help(has_personal, has_org)
         
-        dataset_help = (
-            f"*Private Dataset: {'✅ Available' if has_personal else '❌ Requires personal token (OAuth login or .env)'}*\n"
-            f"*ROOST Dataset: {'✅ Available' if has_org else '⚠️ Can load if public, requires org token to save'}*"
-        )
+        has_personal, has_org = check_token_availability(oauth_token)
         
         return (
-            token_status,  # token_status_markdown
-            save_help,  # save_mode_help
+            format_token_status(oauth_token),  # token_status_markdown
+            format_save_mode_help(has_personal, has_org),  # save_mode_help
             gr.update(interactive=has_personal),  # refresh_private_btn
             gr.update(interactive=True),  # refresh_roost_btn
-            dataset_help,  # dataset_help_text
+            format_dataset_help_text(has_personal, has_org),  # dataset_help_text
         )
     
     login_button.click(
@@ -303,66 +261,48 @@ with gr.Blocks(title="Moderation Model Testing") as demo:
     # Dataset load handler
     def load_example_from_dataset(selected_label, cached_examples_list, dropdown_choices_list):
         """Load example from dataset and populate all fields."""
-        if (not cached_examples_list or not selected_label or 
-            not dropdown_choices_list or selected_label not in dropdown_choices_list):
-            # Return None to skip updates
-            return None, None, None, None, None, None, None, None, None, None, None, None, None, None, None
+        if not (cached_examples_list and selected_label and dropdown_choices_list and 
+                selected_label in dropdown_choices_list):
+            return [None] * 15
         
         try:
-            # Find index by matching label
             idx = dropdown_choices_list.index(selected_label)
-            if idx < 0 or idx >= len(cached_examples_list):
-                return None, None, None, None, None, None, None, None, None, None, None, None, None, None, None
+            if not (0 <= idx < len(cached_examples_list)):
+                return [None] * 15
             
             example = cached_examples_list[idx]
-            
-            # Get policy - ensure it's a string (not None)
             policy = example.get("policy", "") or ""
-            
-            # Extract saved results
             policy_violation = example.get("policy_violation", -1)
-            categories_and_reasoning = example.get("categories_and_reasoning", "")
-            raw_response = example.get("raw_response", "")
-            reasoning_trace = example.get("reasoning_trace", "")
             model_selection = example.get("model_selection", "")
             reasoning_effort_val = example.get("reasoning_effort", "")
+            reasoning_trace = example.get("reasoning_trace", "")
             
             # Format label text
-            if policy_violation == 1:
-                label_text = "## ❌ Policy Violation Detected"
-            elif policy_violation == 0:
-                label_text = "## ✅ No Policy Violation"
-            else:
-                label_text = "## ⚠️ Unable to determine label"
-            
-            # Format model info
-            model_info = format_model_info(model_selection, reasoning_effort_val)
+            emoji = "❌" if policy_violation == 1 else "✅" if policy_violation == 0 else "⚠️"
+            label_text = f"## {emoji} {'Policy Violation Detected' if policy_violation == 1 else 'No Policy Violation' if policy_violation == 0 else 'Unable to determine label'}"
             
-            # Format reasoning info
             reasoning_info_text, reasoning_info_visible = format_reasoning_info(model_selection, reasoning_trace)
-            
             reasoning_visible = bool(reasoning_trace and reasoning_trace.strip())
             
             return (
                 example.get("input", ""),
-                policy,  # current_policy_state - UI syncs automatically via change handler
+                policy,
                 example.get("model_selection", ""),
-                example.get("reasoning_effort", ""),
+                reasoning_effort_val,
                 example.get("max_tokens", 0),
                 example.get("temperature", 0.0),
                 example.get("top_p", 0.0),
                 example.get("system_prompt", ""),
                 example.get("response_format", ""),
-                # Results
-                model_info,
+                format_model_info(model_selection, reasoning_effort_val),
                 label_text,
-                categories_and_reasoning,
-                raw_response,
+                example.get("categories_and_reasoning", ""),
+                example.get("raw_response", ""),
                 gr.update(value=reasoning_info_text, visible=reasoning_info_visible),
                 gr.update(value=reasoning_trace or "", visible=reasoning_visible),
             )
         except (ValueError, IndexError):
-            return None, None, None, None, None, None, None, None, None, None, None, None, None, None, None
+            return [None] * 15
     
     example_dropdown.change(
         load_example_from_dataset,

ui/tab_dataset.py

@@ -8,7 +8,7 @@ import gradio as gr
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 from utils.dataset import get_dataset_repo_id, get_roost_dataset_repo_id, load_dataset_examples
-from utils.helpers import check_token_availability, get_org_token, get_personal_token
+from utils.helpers import check_token_availability, format_dataset_help_text, get_label_emoji, get_org_token, get_personal_token
 from utils.model_interface import extract_model_id, get_model_info
 
 
@@ -26,12 +26,8 @@ def format_preview_markdown(example: dict) -> str:
     model_name = model_info.get("name", model_id) if model_info else model_id or "Unknown"
     
     # Format label with emoji
-    if policy_violation == 1:
-        label_text = "❌ Policy Violation Detected"
-    elif policy_violation == 0:
-        label_text = "✅ No Policy Violation"
-    else:
-        label_text = "⚠️ Unable to determine label"
+    emoji = get_label_emoji(policy_violation)
+    label_text = f"{emoji} Policy Violation Detected" if policy_violation == 1 else f"{emoji} No Policy Violation" if policy_violation == 0 else f"{emoji} Unable to determine label"
     
     # Truncate policy preview
     policy_preview = policy # [:512] + "..." if len(policy) > 512 else policy
@@ -84,10 +80,7 @@ def build_dataset_tab() -> dict:
         
         # Help text explaining token requirements
         dataset_help_text = gr.Markdown(
-            value=(
-                f"*Private Dataset: {'✅ Available' if has_personal else '❌ Requires personal token (OAuth login or .env)'}*\n"
-                f"*ROOST Dataset: {'✅ Available' if has_org else '⚠️ Can load if public, requires org token to save'}*"
-            ),
+            value=format_dataset_help_text(has_personal, has_org),
             visible=True
         )
         
@@ -145,22 +138,14 @@ def build_dataset_tab() -> dict:
             
             return "*Select an example to preview*"
         
-        def refresh_private(oauth_token: gr.OAuthToken | None = None):
-            """Refresh private dataset."""
-            return refresh_dataset("private", oauth_token)
-        
-        def refresh_roost(oauth_token: gr.OAuthToken | None = None):
-            """Refresh ROOST dataset."""
-            return refresh_dataset("roost", oauth_token)
-        
         refresh_private_btn.click(
-            refresh_private,
+            lambda oauth_token=None: refresh_dataset("private", oauth_token),
             inputs=None,  # OAuth token auto-injected
             outputs=[example_dropdown, preview_markdown, cached_examples, dropdown_choices_state]
         )
         
         refresh_roost_btn.click(
-            refresh_roost,
+            lambda oauth_token=None: refresh_dataset("roost", oauth_token),
             inputs=None,  # OAuth token auto-injected
             outputs=[example_dropdown, preview_markdown, cached_examples, dropdown_choices_state]
         )

utils/constants.py

@@ -46,39 +46,6 @@ MODELS = [
     },
 ]
 
-UNUSED_MODELS = [
-    {
-        "name": "Qwen3-4B-Thinking",
-        "id": "Qwen/Qwen3-4B-Thinking-2507",
-        "is_thinking": True,
-        "supports_reasoning_level": False,
-    },
-    {
-        "name": "Qwen3-4B-Instruct",
-        "id": "Qwen/Qwen3-4B-Instruct-2507",
-        "is_thinking": False,
-        "supports_reasoning_level": False,
-    },
-    {
-        "name": "Qwen3-30B-Instruct",
-        "id": "Qwen/Qwen3-30B-A3B-Instruct-2507",
-        "is_thinking": False,
-        "supports_reasoning_level": False,
-    },
-    {
-        "name": "GPT-OSS-20B",
-        "id": "openai/gpt-oss-20b",
-        "is_thinking": True,
-        "supports_reasoning_level": True,
-    },
-    {
-        "name": "GPT-OSS-120B",
-        "id": "openai/gpt-oss-120b",
-        "is_thinking": True,
-        "supports_reasoning_level": True,
-    },
-]
-
 # Reasoning effort levels for GPT-OSS
 REASONING_EFFORTS = ["Low", "Medium", "High"]
 
@@ -105,10 +72,10 @@ RESPONSE_FORMAT = """## Response Format
 
 Given a content item, please provide a JSON-formatted response with the following fields:
 - `label`: INT - 1 for a policy violation, 0 for no violation
-- `categories`: LIST[DCT] - a list of categories from the disallowed and allowed categories that the content item falls into along with the reasoning for each category. Each dict has the following keys:
+- `categories`: LIST[DCT] - a list of categories (at least one) from the disallowed and allowed categories that the content item falls into along with the reasoning for each category. Each dict has the following keys:
   - `category`: STR - the category from the policy
   - `reasoning`: STR - the reasoning for why the content item falls into this category
-  - `policy_source`: STR - specific text from the policy that best supports the reasoning. Use [...] to concatenate multi-part citations. Make sure to quote the policy text exactly and include all relevant passages.
+  - `policy_source`: STR - specific text from the policy that best supports the reasoning. Use [...] to concatenate multi-part citations. Prioritize policy explanatory text over repeating the category name. Make sure to quote the policy text exactly and include all relevant passages.
 """
 
 # Test examples from notebook

utils/dataset.py

@@ -6,6 +6,7 @@ from datasets import Dataset, load_dataset
 from huggingface_hub import HfApi
 
 from utils.model_interface import extract_model_id, get_model_info
+from utils.helpers import get_label_emoji
 
 
 def get_username_from_token(token: str | None) -> str:
@@ -175,12 +176,7 @@ def load_dataset_examples(repo_id: str, token: str | None) -> tuple[list[dict],
         policy_violation = example.get("policy_violation", -1)
         
         # Get label emoji
-        if policy_violation == 1:
-            label_emoji = "❌"
-        elif policy_violation == 0:
-            label_emoji = "✅"
-        else:
-            label_emoji = "⚠️"
+        label_emoji = get_label_emoji(policy_violation)
         
         # Extract model name
         model_id = extract_model_id(model_selection)

utils/helpers.py

@@ -6,6 +6,25 @@ import gradio as gr
 from dotenv import load_dotenv
 
 
+def _load_token_from_env(env_var: str) -> str | None:
+    """Load token from .env file."""
+    load_dotenv()
+    return os.getenv(env_var)
+
+
+def get_label_emoji(policy_violation: int) -> str:
+    """Get emoji for policy violation label."""
+    return "❌" if policy_violation == 1 else "✅" if policy_violation == 0 else "⚠️"
+
+
+def format_dataset_help_text(has_personal: bool, has_org: bool) -> str:
+    """Format help text explaining dataset availability."""
+    return (
+        f"*Private Dataset: {'✅ Available' if has_personal else '❌ Requires personal token (OAuth login or .env)'}*\n"
+        f"*ROOST Dataset: {'✅ Available' if has_org else '⚠️ Can load if public, requires org token to save'}*"
+    )
+
+
 def get_personal_token(oauth_token: gr.OAuthToken | None) -> tuple[str | None, str]:
     """
     Get personal Hugging Face token from OAuth or .env fallback.
@@ -20,30 +39,18 @@ def get_personal_token(oauth_token: gr.OAuthToken | None) -> tuple[str | None, s
         - hf_token: Token string if available, None otherwise
         - status_message: Warning message if using local .env, empty string otherwise
     """
-    print(f"DEBUG: get_personal_token called with oauth_token type: {type(oauth_token)}")
-    
     if oauth_token is None or (isinstance(oauth_token, str) and oauth_token == "Log in to Hugging Face"):
         # Try loading from .env file
-        print("DEBUG: oauth_token is None, loading from .env")
-        load_dotenv()
-        hf_token = os.getenv("HF_TOKEN_MLSOC")
+        hf_token = _load_token_from_env("HF_TOKEN_MLSOC")
         if hf_token is None:
-            print("DEBUG: HF_TOKEN_MLSOC not found in .env")
             return None, ""
-        else:
-            print(f"DEBUG: Loaded token from .env, length: {len(hf_token)}, first 4 chars: {hf_token[:4] if len(hf_token) >= 4 else hf_token}")
-            return hf_token, "\n⚠️ Using local .env file for token (not online)"
+        return hf_token, "\n⚠️ Using local .env file for token (not online)"
     else:
         # OAuthToken object
-        print(f"DEBUG: oauth_token is OAuthToken object")
         token = oauth_token.token
-        print(f"DEBUG: Extracted token from OAuthToken, length: {len(token) if token else 0}, first 4 chars: {token[:4] if token and len(token) >= 4 else (token if token else 'None')}")
         if not token or not token.strip():
-            print("DEBUG: OAuthToken.token is empty, falling back to .env")
-            load_dotenv()
-            hf_token = os.getenv("HF_TOKEN_MLSOC")
+            hf_token = _load_token_from_env("HF_TOKEN_MLSOC")
             if hf_token:
-                print(f"DEBUG: Loaded token from .env (empty OAuth case), length: {len(hf_token)}, first 4 chars: {hf_token[:4] if len(hf_token) >= 4 else hf_token}")
                 return hf_token, "\n⚠️ Using local .env file for token (not online)"
             return None, ""
         return token, ""
@@ -64,12 +71,7 @@ def get_org_token() -> str | None:
         return org_token
     
     # Fall back to .env file
-    load_dotenv()
-    org_token = os.getenv("ROOST_TOKEN_FALLBACK")
-    if org_token:
-        return org_token
-    
-    return None
+    return _load_token_from_env("ROOST_TOKEN_FALLBACK")
 
 
 def get_inference_token(oauth_token: gr.OAuthToken | None) -> tuple[str | None, str]:
@@ -111,7 +113,7 @@ def format_token_status(oauth_token: gr.OAuthToken | None) -> str:
     has_personal, has_org = check_token_availability(oauth_token)
     
     lines = [
-        "You can log in to yout Hugging Face account to save your work in a private dataset and use the app for inference after the end of the hackathon.",
+        "You can log in to your Hugging Face account to save your work in a private dataset and use the app for inference after the end of the hackathon.",
         "### Token Status",
     ]