[Ref, Add] custom css for sizing, move demo utility to its own file

app.py

@@ -1,94 +1,8 @@
-import os
 import gradio as gr
-import numpy as np
-import matplotlib.pyplot as plt
-import seaborn as sns
-from io import BytesIO
-from PIL import Image
-from datasets.exceptions import DatasetNotFoundError
 
+import src.app_util as app_util
 from src.dataloading import get_leaderboard_models_cached, get_leaderboard_datasets
-from src.similarity import load_data_and_compute_similarities
 
-# Set matplotlib backend for non-GUI environments
-plt.switch_backend('Agg')
-
-
-def create_heatmap(selected_models, selected_dataset, selected_metric):
-    if not selected_models or not selected_dataset:
-        return None
-    
-    # Sort models and get short names
-    similarities = load_data_and_compute_similarities(selected_models, selected_dataset, selected_metric)
-
-    # Check if similarity matrix contains NaN rows
-    failed_models = []
-    for i in range(len(similarities)):
-        if np.isnan(similarities[i]).all():
-            failed_models.append(selected_models[i])
-
-    if failed_models:
-        gr.Warning(f"Failed to load data for models: {', '.join(failed_models)}")
-
-    # Create figure and heatmap using seaborn
-    plt.figure(figsize=(8, 6))
-    ax = sns.heatmap(
-        similarities,
-        annot=True,
-        fmt=".2f",
-        cmap="viridis",
-        vmin=0,
-        vmax=1,
-        xticklabels=selected_models,
-        yticklabels=selected_models
-    )
-    
-    # Customize plot
-    plt.title(f"{selected_metric} for {selected_dataset}", fontsize=16)
-    plt.xlabel("Models", fontsize=14)
-    plt.ylabel("Models", fontsize=14)
-    plt.xticks(rotation=45, ha='right')
-    plt.yticks(rotation=0)
-    plt.tight_layout()
-
-    # Save to buffer
-    buf = BytesIO()
-    plt.savefig(buf, format="png", dpi=100, bbox_inches="tight")
-    plt.close()
-    
-    # Convert to PIL Image
-    buf.seek(0)
-    img = Image.open(buf).convert("RGB")
-    return img
-
-def validate_inputs(selected_models, selected_dataset):
-    if not selected_models:
-        raise gr.Error("Please select at least one model!")
-    if not selected_dataset:
-        raise gr.Error("Please select a dataset!")
-    
-
-def update_datasets_based_on_models(selected_models, current_dataset):
-    try:
-        available_datasets = get_leaderboard_datasets(selected_models) if selected_models else []
-        if current_dataset in available_datasets:
-            valid_dataset = current_dataset
-        elif "mmlu_pro" in available_datasets:
-            valid_dataset = "mmlu_pro"
-        else:
-            valid_dataset = None
-        return gr.update(
-            choices=available_datasets,
-            value=valid_dataset
-        )
-    except DatasetNotFoundError as e:
-        # Extract model name from error message
-        model_name = e.args[0].split("'")[1]
-        model_name = model_name.split("/")[-1].replace("__", "/").replace("_details", "")
-        
-        # Display a shorter warning
-        gr.Warning(f"Data for '{model_name}' is gated or unavailable.")
-        return gr.update(choices=[], value=None)
 
 links_markdown = """
 [📄 Paper](https://arxiv.org/abs/2502.04313)   |  
@@ -104,7 +18,7 @@ metric_init = "CAPA"
 
 
 # Create Gradio interface
-with gr.Blocks(title="LLM Similarity Analyzer") as demo:
+with gr.Blocks(title="LLM Similarity Analyzer", css=app_util.custom_css) as demo:
     gr.Markdown("# Model Similarity Comparison Tool")
     gr.Markdown(links_markdown)
     gr.Markdown('Demo for the recent publication ["Great Models Think Alike and this Undermines AI Oversight"](https://huggingface.co/papers/2502.04313).')
@@ -137,20 +51,20 @@ with gr.Blocks(title="LLM Similarity Analyzer") as demo:
     )
 
     model_dropdown.change(
-        fn=update_datasets_based_on_models,
+        fn=app_util.update_datasets_based_on_models,
         inputs=[model_dropdown, dataset_dropdown],
         outputs=dataset_dropdown
     )
     
     generate_btn = gr.Button("Generate Heatmap", variant="primary")
-    heatmap = gr.Image(value=create_heatmap(model_init, dataset_init, metric_init), label="Similarity Heatmap", visible=True)
+    heatmap = gr.Image(value=app_util.create_heatmap(model_init, dataset_init, metric_init), label="Similarity Heatmap", visible=True)
 
     generate_btn.click(
-        fn=validate_inputs,
+        fn=app_util.validate_inputs,
         inputs=[model_dropdown, dataset_dropdown],
         queue=False
     ).then(
-        fn=create_heatmap,
+        fn=app_util.create_heatmap,
         inputs=[model_dropdown, dataset_dropdown, metric_dropdown],
         outputs=heatmap
     )
@@ -170,7 +84,7 @@ biased towards more similar models controlling for the model's capability. (2) G
 of weak supervisors (weak-to-strong generalization) is higher when the two models are more different. (3) Concerningly, model \
 errors are getting more correlated as capabilities increase.""")
     with gr.Row():
-        gr.Image(value="data/table_capa.png", label="Comparison of different similarity metrics for multiple-choice questions", interactive=False, scale=1)
+        gr.Image(value="data/table_capa.png", label="Comparison of different similarity metrics for multiple-choice questions", elem_classes="image_container", interactive=False)
     gr.Markdown("""
 - **Datasets**: [Open LLM Leaderboard v2](https://huggingface.co/spaces/open-llm-leaderboard/open_llm_leaderboard#/) benchmark datasets \n
     - Some datasets are not multiple-choice - for these, the metrics are not applicable. \n

src/app_util.py

@@ -0,0 +1,99 @@
+import gradio as gr
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+from io import BytesIO
+from PIL import Image
+from datasets.exceptions import DatasetNotFoundError
+
+from src.dataloading import get_leaderboard_datasets
+from src.similarity import load_data_and_compute_similarities
+
+# Set matplotlib backend for non-GUI environments
+plt.switch_backend('Agg')
+
+def create_heatmap(selected_models, selected_dataset, selected_metric):
+    if not selected_models or not selected_dataset:
+        return None
+    
+    # Sort models and get short names
+    similarities = load_data_and_compute_similarities(selected_models, selected_dataset, selected_metric)
+
+    # Check if similarity matrix contains NaN rows
+    failed_models = []
+    for i in range(len(similarities)):
+        if np.isnan(similarities[i]).all():
+            failed_models.append(selected_models[i])
+
+    if failed_models:
+        gr.Warning(f"Failed to load data for models: {', '.join(failed_models)}")
+
+    # Create figure and heatmap using seaborn
+    plt.figure(figsize=(8, 6))
+    ax = sns.heatmap(
+        similarities,
+        annot=True,
+        fmt=".2f",
+        cmap="viridis",
+        vmin=0,
+        vmax=1,
+        xticklabels=selected_models,
+        yticklabels=selected_models
+    )
+    
+    # Customize plot
+    plt.title(f"{selected_metric} for {selected_dataset}", fontsize=16)
+    plt.xlabel("Models", fontsize=14)
+    plt.ylabel("Models", fontsize=14)
+    plt.xticks(rotation=45, ha='right')
+    plt.yticks(rotation=0)
+    plt.tight_layout()
+
+    # Save to buffer
+    buf = BytesIO()
+    plt.savefig(buf, format="png", dpi=100, bbox_inches="tight")
+    plt.close()
+    
+    # Convert to PIL Image
+    buf.seek(0)
+    img = Image.open(buf).convert("RGB")
+    return img
+
+def validate_inputs(selected_models, selected_dataset):
+    if not selected_models:
+        raise gr.Error("Please select at least one model!")
+    if not selected_dataset:
+        raise gr.Error("Please select a dataset!")
+    
+
+def update_datasets_based_on_models(selected_models, current_dataset):
+    try:
+        available_datasets = get_leaderboard_datasets(selected_models) if selected_models else []
+        if current_dataset in available_datasets:
+            valid_dataset = current_dataset
+        elif "mmlu_pro" in available_datasets:
+            valid_dataset = "mmlu_pro"
+        else:
+            valid_dataset = None
+        return gr.update(
+            choices=available_datasets,
+            value=valid_dataset
+        )
+    except DatasetNotFoundError as e:
+        # Extract model name from error message
+        model_name = e.args[0].split("'")[1]
+        model_name = model_name.split("/")[-1].replace("__", "/").replace("_details", "")
+        
+        # Display a shorter warning
+        gr.Warning(f"Data for '{model_name}' is gated or unavailable.")
+        return gr.update(choices=[], value=None)
+    
+custom_css = """
+.image-container img {
+    width: 80% !important;  /* Make it 80% of the parent container */
+    height: auto !important; /* Maintain aspect ratio */
+    max-width: 800px; /* Optional: Set a max limit */
+    display: block;
+    margin: auto; /* Center the image */
+}
+"""