Update streamlit_app.py

streamlit_app.py

@@ -12,17 +12,179 @@ import plotly.express as px
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 import warnings
+import datasets
+import logging
 warnings.filterwarnings('ignore')
 
-# Import visualization utilities
-from visualization.utils import (
-    load_and_prepare_dataset, 
-    get_available_turn_metrics,
-    get_human_friendly_metric_name,
-    clean_metric_values,
-    PLOT_PALETTE,
-    setup_plot_style
-)
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Constants
+PLOT_PALETTE = {
+    "jailbreak": "#D000D8",  # Purple
+    "benign": "#008393",     # Cyan
+    "control": "#EF0000",    # Red
+}
+
+# Utility functions
+def load_and_prepare_dataset(dataset_config):
+    """Load the risky conversations dataset and prepare it for analysis."""
+    logger.info("Loading dataset...")
+    
+    dataset_name = dataset_config["dataset_name"]
+    logger.info(f"Loading dataset: {dataset_name}")
+    
+    # Load the dataset
+    dataset = datasets.load_dataset(dataset_name, split="train")
+    logger.info(f"Dataset loaded with {len(dataset)} conversations")
+    
+    # Convert to pandas
+    pandas_dataset = dataset.to_pandas()
+    
+    # Explode the conversation column
+    pandas_dataset_exploded = pandas_dataset.explode("conversation")
+    pandas_dataset_exploded = pandas_dataset_exploded.reset_index(drop=True)
+    
+    # Normalize conversation data
+    conversations_unfolded = pd.json_normalize(
+        pandas_dataset_exploded["conversation"],
+    )
+    conversations_unfolded = conversations_unfolded.add_prefix("turn.")
+    
+    # Ensure there's a 'conversation_metrics' column, even if empty
+    if "conversation_metrics" not in pandas_dataset_exploded.columns:
+        pandas_dataset_exploded["conversation_metrics"] = [{}] * len(
+            pandas_dataset_exploded
+        )
+    
+    # Normalize conversation metrics
+    conversations_metrics_unfolded = pd.json_normalize(
+        pandas_dataset_exploded["conversation_metrics"]
+    )
+    conversations_metrics_unfolded = conversations_metrics_unfolded.add_prefix(
+        "conversation_metrics."
+    )
+    
+    # Concatenate all dataframes
+    pandas_dataset_exploded = pd.concat(
+        [
+            pandas_dataset_exploded.drop(
+                columns=["conversation", "conversation_metrics"]
+            ),
+            conversations_unfolded,
+            conversations_metrics_unfolded,
+        ],
+        axis=1,
+    )
+    
+    logger.info(f"Dataset prepared with {len(pandas_dataset_exploded)} turns")
+    return pandas_dataset, pandas_dataset_exploded
+
+def get_available_turn_metrics(dataset_exploded):
+    """Dynamically discover all available turn metrics from the dataset."""
+    # Find all columns that contain turn metrics
+    turn_metric_columns = [
+        col for col in dataset_exploded.columns if col.startswith("turn.turn_metrics.")
+    ]
+    
+    # Extract the metric names by removing the prefix
+    metrics = [col.replace("turn.turn_metrics.", "") for col in turn_metric_columns]
+    
+    # Sort metrics for consistent ordering
+    metrics = sorted(metrics)
+    
+    logger.info(f"Discovered {len(metrics)} available turn metrics")
+    return metrics
+
+def get_human_friendly_metric_name(metric_name):
+    """Convert a metric name to a human-friendly display name."""
+    # Mapping of metric names to human-friendly display names
+    metric_display_names = {
+        # Length metrics
+        "conversational_length_tokens": "Length in Tokens",
+        "conversational_length_words": "Length in Words", 
+        "conversational_length_bytes": "Length in Bytes",
+        
+        # Readability metrics
+        "readability_flesch_score": "Flesch Reading Ease",
+        "readability_standard_score": "Reading Grade Level",
+        
+        # Compression metrics
+        "lzw_compression": "LZW Compression Ratio",
+        
+        # Language model metrics
+        "ll_mean_logprobs": "Mean Log Probability",
+        "ll_max_logprobs": "Maximum Log Probability",
+        "ll_min_logprobs": "Minimum Log Probability",
+        "rll_mean_logprobs": "Reverse Mean Log Probability",
+        "rll_max_logprobs": "Reverse Maximum Log Probability",
+        "rll_min_logprobs": "Reverse Minimum Log Probability",
+        
+        # Lexical diversity
+        "type_token_ratio": "Type-Token Ratio",
+        
+        # Language distribution
+        "zipf_score": "Zipf Distribution Score",
+        
+        # Evaluation metrics
+        "rubric_evaluation.Stealth": "Stealth Score",
+        
+        # Working memory metrics
+        "wm_total_entity_references": "Total Entity References",
+        "wm_unique_entities_count": "Unique Entities Count",
+        "wm_named_entity_mentions": "Named Entity Mentions",
+        "wm_concept_noun_mentions": "Concept Noun Mentions", 
+        "wm_pronoun_references": "Pronoun References",
+        "wm_entity_density_per_word": "Entity Density per Word",
+        "wm_entity_density_per_100_words": "Entity Density per 100 Words",
+        "wm_entity_density_per_100_chars": "Entity Density per 100 Characters",
+        "wm_entity_diversity_ratio": "Entity Diversity Ratio",
+        "wm_entity_repetition_ratio": "Entity Repetition Ratio",
+        "wm_cognitive_load_score": "Cognitive Load Score",
+        "wm_high_cognitive_load": "High Cognitive Load",
+        
+        # Discourse coherence metrics
+        "discourse_coherence_to_next_user": "Coherence to Next User Turn",
+        "discourse_coherence_to_next_turn": "Coherence to Next Turn",
+        "discourse_mean_user_coherence": "Mean User Coherence",
+        "discourse_user_coherence_variance": "User Coherence Variance",
+        "discourse_user_topic_drift": "User Topic Drift",
+        "discourse_user_entity_continuity": "User Entity Continuity",
+        "discourse_num_user_turns": "Number of User Turns",
+        
+        # Tokens per byte
+        "tokens_per_byte": "Tokens per Byte",
+    }
+    
+    # Check exact match first
+    if metric_name in metric_display_names:
+        return metric_display_names[metric_name]
+    
+    # Handle conversation-level aggregations
+    for suffix in ["_conversation_mean", "_conversation_min", "_conversation_max", "_conversation_std", "_conversation_count"]:
+        if metric_name.endswith(suffix):
+            base_metric = metric_name[:-len(suffix)]
+            if base_metric in metric_display_names:
+                agg_type = suffix.split("_")[-1].title()
+                return f"{metric_display_names[base_metric]} ({agg_type})"
+    
+    # Handle turn-level metrics with "turn.turn_metrics." prefix
+    if metric_name.startswith("turn.turn_metrics."):
+        base_metric = metric_name[len("turn.turn_metrics."):]
+        if base_metric in metric_display_names:
+            return metric_display_names[base_metric]
+    
+    # Fallback: convert underscores to spaces and title case
+    clean_name = metric_name
+    for prefix in ["turn.turn_metrics.", "conversation_metrics.", "turn_metrics."]:
+        if clean_name.startswith(prefix):
+            clean_name = clean_name[len(prefix):]
+            break
+    
+    # Convert to human-readable format
+    clean_name = clean_name.replace("_", " ").title()
+    return clean_name
 
 # Setup page config
 st.set_page_config(
@@ -113,6 +275,11 @@ def main():
     if not data_loaded:
         st.stop()
     
+    # Check if we have data after filtering
+    if len(filtered_df_exploded) == 0:
+        st.error("No data available with current filters. Please adjust your selection.")
+        st.stop()
+    
     # Sidebar controls
     st.sidebar.header("🎛️ Controls")
     
@@ -127,13 +294,32 @@ def main():
     
     # Role filter
     if 'turn.role' in df_exploded.columns:
-        roles = df_exploded['turn.role'].unique()
-        selected_roles = st.sidebar.multiselect(
-            "Select Roles",
-            options=roles,
-            default=roles,
-            help="Filter by turn role"
-        )
+        roles = df_exploded['turn.role'].dropna().unique()
+        # Assert only user and assistant roles exist
+        expected_roles = {'user', 'assistant'}
+        actual_roles = set(roles)
+        assert actual_roles.issubset(expected_roles), f"Unexpected roles found: {actual_roles - expected_roles}. Expected only 'user' and 'assistant'"
+        
+        st.sidebar.subheader("👥 Role Filter")
+        col1, col2 = st.sidebar.columns(2)
+        
+        with col1:
+            include_user = st.checkbox("User", value=True, help="Include user turns")
+        with col2:
+            include_assistant = st.checkbox("Assistant", value=True, help="Include assistant turns")
+        
+        # Build selected roles list
+        selected_roles = []
+        if include_user and 'user' in roles:
+            selected_roles.append('user')
+        if include_assistant and 'assistant' in roles:
+            selected_roles.append('assistant')
+            
+        # Show selection info
+        if selected_roles:
+            st.sidebar.success(f"Including: {', '.join(selected_roles)}")
+        else:
+            st.sidebar.warning("No roles selected")
     else:
         selected_roles = None
     
@@ -303,6 +489,9 @@ def main():
     
     if selected_roles and 'turn.role' in filtered_df_exploded.columns:
         filtered_df_exploded = filtered_df_exploded[filtered_df_exploded['turn.role'].isin(selected_roles)]
+    elif selected_roles is not None and len(selected_roles) == 0:
+        # If roles exist but none are selected, show empty dataset
+        filtered_df_exploded = filtered_df_exploded.iloc[0:0]  # Empty dataframe with same structure
     
     # Main content tabs
     tab1, tab2, tab3, tab4 = st.tabs(["📊 Distributions", "🔗 Correlations", "📈 Comparisons", "🎯 Details"])