"""Main Gradio app for moderation model testing.""" import os import sys import gradio as gr sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from datetime import datetime from utils.dataset import ( format_categories_and_reasoning, get_dataset_repo_id, get_roost_dataset_repo_id, save_to_dataset, ) from utils.helpers import ( check_token_availability, format_dataset_help_text, format_token_status, get_inference_token, get_org_token, get_personal_token, ) from utils.model_interface import extract_model_id, run_test from ui.sidebar import build_sidebar from ui.tab_config import build_config_tab from ui.tab_dataset import build_dataset_tab from ui.tab_policy import build_policy_tab from ui.tab_testing import ( build_testing_tab, format_model_info, format_reasoning_info, format_test_result, ) # ============================================================================ # 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.""" if not test_input or not test_input.strip(): raise gr.Error("Please enter test content before running a test.") if not current_policy or current_policy == "*No policy loaded*": raise gr.Error("Please load a policy first. Go to the Policy Definition tab to upload or select a policy.") # OAuth token is automatically injected by Gradio - we don't pass login_button as input # Use inference token (org preferred, falls back to personal) hf_token, _ = get_inference_token(oauth_token) if hf_token is None: raise gr.Error("Please log in or set tokens to use Inference Providers. Check the sidebar for authentication options.") model_id = extract_model_id(model_choice) try: result = run_test( model_id=model_id, test_input=test_input, policy=current_policy, hf_token=hf_token, reasoning_effort=reasoning_effort, max_tokens=int(max_tokens), temperature=float(temperature), top_p=float(top_p), system_prompt=system_prompt_val, response_format=response_format_val, ) except gr.Error: raise # Re-raise Gradio errors except Exception as e: raise gr.Error(f"Unexpected error during model inference: {str(e)}. Please try again.") label_text, parsed, cat_text, reasoning, raw_response = format_test_result(result) reasoning_visible = bool(reasoning and reasoning.strip()) model_info = format_model_info(model_choice, reasoning_effort) reasoning_info_text, reasoning_info_visible = format_reasoning_info(model_choice, reasoning) # Save to dataset if enabled if save_mode == "Save to ROOST Dataset": org_token = get_org_token() if org_token: try: 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 ) success, message = save_to_dataset(get_roost_dataset_repo_id(), org_token, data) if not success: raise gr.Error(f"Failed to save to ROOST dataset: {message}. Please check your token permissions.") except gr.Error: raise # Re-raise Gradio errors except Exception as e: raise gr.Error(f"Failed to save to ROOST dataset: {str(e)}. Please check your token permissions and try again.") elif save_mode == "Save to Private Dataset": personal_token, _ = get_personal_token(oauth_token) if personal_token: try: 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 ) success, message = save_to_dataset(get_dataset_repo_id(personal_token), personal_token, data) if not success: raise gr.Error(f"Failed to save to private dataset: {message}. Please check your token permissions.") except gr.Error: raise # Re-raise Gradio errors except Exception as e: raise gr.Error(f"Failed to save to private dataset: {str(e)}. Please check your token permissions and try again.") return ( model_info, label_text, cat_text, raw_response, gr.update(value=reasoning_info_text, visible=reasoning_info_visible), gr.update(value=reasoning or "", visible=reasoning_visible), ) # ============================================================================ # UI Components # ============================================================================ with gr.Blocks(title="Moderation Model Testing") as demo: gr.Markdown("# Moderation Model Testing Interface") gr.Markdown( "Test moderation models with custom content policies. Define your policy, select a model, " "and evaluate how different models classify content according to your rules. " "Supports reasoning models that provide detailed explanations for their decisions." ) # Sidebar (collapsible) sidebar_components = build_sidebar() login_button = sidebar_components["login_button"] token_status_markdown = sidebar_components["token_status"] # Main content area with tabs with gr.Tabs(): # Build tabs testing_components = build_testing_tab() test_input = testing_components["test_input"] run_test_btn = testing_components["run_test_btn"] save_mode = testing_components["save_mode"] save_mode_help = testing_components["save_mode_help"] model_info_display = testing_components["model_info_display"] label_display = testing_components["label_display"] categories_display = testing_components["categories_display"] model_response_display = testing_components["model_response_display"] reasoning_info = testing_components["reasoning_info"] reasoning_display = testing_components["reasoning_display"] policy_components = build_policy_tab(os.path.dirname(__file__)) current_policy_state = policy_components["current_policy_state"] config_components = build_config_tab() model_dropdown = config_components["model_dropdown"] reasoning_effort = config_components["reasoning_effort"] max_tokens = config_components["max_tokens"] temperature = config_components["temperature"] top_p = config_components["top_p"] system_prompt_textbox = config_components["system_prompt_textbox"] response_format_textbox = config_components["response_format_textbox"] dataset_components = build_dataset_tab() example_dropdown = dataset_components["example_dropdown"] cached_examples = dataset_components["cached_examples"] dropdown_choices_state = dataset_components["dropdown_choices_state"] refresh_private_btn = dataset_components["refresh_private_btn"] refresh_roost_btn = dataset_components["refresh_roost_btn"] dataset_help_text = dataset_components["dataset_help_text"] # ============================================================================ # Event Handlers # ============================================================================ # Cross-tab handler: Run test (needs components from all tabs) run_test_btn.click( handle_run_test, inputs=[ test_input, current_policy_state, model_dropdown, reasoning_effort, max_tokens, temperature, top_p, system_prompt_textbox, response_format_textbox, save_mode, ], outputs=[ model_info_display, label_display, categories_display, model_response_display, reasoning_info, reasoning_display, ], ) model_dropdown.change( format_model_info, inputs=[model_dropdown, reasoning_effort], outputs=model_info_display, ) reasoning_effort.change( format_model_info, inputs=[model_dropdown, reasoning_effort], outputs=model_info_display, ) # 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.""" from ui.tab_testing import format_save_mode_help has_personal, has_org = check_token_availability(oauth_token) return ( 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 format_dataset_help_text(has_personal, has_org), # dataset_help_text ) login_button.click( handle_login_click, inputs=None, # OAuth token auto-injected outputs=[ token_status_markdown, save_mode_help, refresh_private_btn, refresh_roost_btn, dataset_help_text, ] ) # Update token status on app load (when OAuth token is already available) demo.load( handle_login_click, inputs=None, # OAuth token auto-injected outputs=[ token_status_markdown, save_mode_help, refresh_private_btn, refresh_roost_btn, dataset_help_text, ] ) # 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 and selected_label and dropdown_choices_list and selected_label in dropdown_choices_list): return [None] * 15 try: idx = dropdown_choices_list.index(selected_label) if not (0 <= idx < len(cached_examples_list)): raise gr.Warning("Selected example index is out of range. Please refresh the dataset.") example = cached_examples_list[idx] policy = example.get("policy", "") or "" policy_violation = example.get("policy_violation", -1) model_selection = example.get("model_selection", "") reasoning_effort_val = example.get("reasoning_effort", "") reasoning_trace = example.get("reasoning_trace", "") # Format label text 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'}" 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, example.get("model_selection", ""), 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", ""), format_model_info(model_selection, reasoning_effort_val), label_text, 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 gr.Warning: raise # Re-raise Gradio warnings except (ValueError, IndexError) as e: raise gr.Warning(f"Failed to load example: {str(e)}. Please try selecting a different example or refresh the dataset.") example_dropdown.change( load_example_from_dataset, inputs=[example_dropdown, cached_examples, dropdown_choices_state], outputs=[ test_input, current_policy_state, # UI components sync automatically via change handler model_dropdown, reasoning_effort, max_tokens, temperature, top_p, system_prompt_textbox, response_format_textbox, # Results model_info_display, label_display, categories_display, model_response_display, reasoning_info, reasoning_display, ], ) # Update token status on app load (when OAuth token is already available) demo.load( handle_login_click, inputs=None, # OAuth token auto-injected outputs=[ token_status_markdown, save_mode_help, refresh_private_btn, refresh_roost_btn, dataset_help_text, ] ) if __name__ == "__main__": demo.launch(ssr_mode=False)