Update app.py

app.py

@@ -1,6 +1,3 @@
-import json
-from pathlib import Path
-
 import gradio as gr
 import matplotlib.pyplot as plt
 import numpy as np
@@ -9,34 +6,29 @@ from datasets import load_dataset
 from sklearn.ensemble import RandomForestRegressor
 from sklearn.metrics import mean_absolute_error, r2_score
 from sklearn.model_selection import train_test_split
-
+from pathlib import Path
 
 # =========================================================
-# CONFIG
+# CONFIG & REPOSITORIES
 # =========================================================
-HF_DATASET_NAME = "QSBench/QSBench-Core-v1.0.0-demo"
-LOCAL_BENCHMARK_CSV = "noise_benchmark_results.csv"
+DATASET_MAP = {
+    "Core (Clean)": "QSBench/QSBench-Core-v1.0.0-demo",
+    "Depolarizing Noise": "QSBench/QSBench-Depolarizing-v1.0.0-demo",
+    "Amplitude Damping": "QSBench/QSBench-Amplitude-v1.0.0-demo",
+    "Transpilation (10q)": "QSBench/QSBench-Transpilation-v1.0.0-demo"
+}
 
+LOCAL_BENCHMARK_CSV = "noise_benchmark_results.csv"
 TARGET_COL = "ideal_expval_Z_global"
 
 EXCLUDE_COLS = {
-    "sample_id",
-    "sample_seed",
-    "ideal_expval_Z_global",
-    "ideal_expval_X_global",
-    "ideal_expval_Y_global",
-    "noisy_expval_Z_global",
-    "noisy_expval_X_global",
-    "noisy_expval_Y_global",
-    "error_Z_global",
-    "error_X_global",
-    "error_Y_global",
-    "sign_ideal_Z_global",
-    "sign_noisy_Z_global",
-    "sign_ideal_X_global",
-    "sign_noisy_X_global",
-    "sign_ideal_Y_global",
-    "sign_noisy_Y_global",
+    "sample_id", "sample_seed", "split",
+    "ideal_expval_Z_global", "ideal_expval_X_global", "ideal_expval_Y_global",
+    "noisy_expval_Z_global", "noisy_expval_X_global", "noisy_expval_Y_global",
+    "error_Z_global", "error_X_global", "error_Y_global",
+    "sign_ideal_Z_global", "sign_noisy_Z_global", 
+    "sign_ideal_X_global", "sign_noisy_X_global",
+    "sign_ideal_Y_global", "sign_noisy_Y_global",
 }
 
 MODEL_PARAMS = dict(
@@ -47,89 +39,50 @@ MODEL_PARAMS = dict(
     n_jobs=-1,
 )
 
+# Global cache to avoid redundant downloads
+dataset_cache = {}
 
 # =========================================================
-# DATA LOADING
+# DATA UTILS
 # =========================================================
-def load_demo_dataset() -> pd.DataFrame:
-    ds_all = load_dataset(HF_DATASET_NAME)
-    df_all = pd.DataFrame(ds_all["train"])
-    return df_all
-
-
-def split_by_split_column(df: pd.DataFrame) -> dict:
-    if "split" not in df.columns:
-        return {"all": df.reset_index(drop=True)}
-
-    splits = {}
-    for split_name in df["split"].dropna().astype(str).unique():
-        splits[split_name] = df[df["split"].astype(str) == split_name].reset_index(drop=True)
-    return splits
-
+def get_df(dataset_key):
+    if dataset_key not in dataset_cache:
+        repo_id = DATASET_MAP[dataset_key]
+        print(f"Downloading {repo_id}...")
+        ds = load_dataset(repo_id)
+        dataset_cache[dataset_key] = pd.DataFrame(ds["train"])
+    return dataset_cache[dataset_key]
 
 def get_numeric_feature_cols(df: pd.DataFrame) -> list[str]:
     numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
-    feature_cols = [c for c in numeric_cols if c not in EXCLUDE_COLS and not c.startswith("error_")]
-    return feature_cols
-
-
-def load_benchmark_results() -> pd.DataFrame:
-    path = Path(LOCAL_BENCHMARK_CSV)
-    if not path.exists():
-        return pd.DataFrame(
-            [
-                {
-                    "dataset": "noise_benchmark_results.csv not found",
-                    "split_used": "",
-                    "n_samples": 0,
-                    "r2": np.nan,
-                    "mae": np.nan,
-                    "avg_noise_prob": np.nan,
-                    "status": "missing_file",
-                }
-            ]
-        )
-
-    df = pd.read_csv(path)
-    return df
-
-
-# =========================================================
-# DATA EXPLORER TAB
-# =========================================================
-def show_data(split_name, splits_cache):
-    if not splits_cache:
-        return pd.DataFrame([{"message": "Dataset not loaded"}])
-
-    if split_name in splits_cache:
-        return splits_cache[split_name].head(10)
-
-    first_key = next(iter(splits_cache.keys()))
-    return splits_cache[first_key].head(10)
-
+    return [c for c in numeric_cols if c not in EXCLUDE_COLS and not c.startswith("error_")]
 
 # =========================================================
-# MODEL DEMO TAB
+# TAB FUNCTIONS
 # =========================================================
-def train_model_demo(df: pd.DataFrame):
-    if TARGET_COL not in df.columns:
-        return None, "Target column not found."
-
+def update_explorer(dataset_name):
+    df = get_df(dataset_name)
+    splits = df["split"].unique().tolist() if "split" in df.columns else ["all"]
+    return gr.update(choices=splits, value=splits[0]), df.head(10)
+
+def filter_explorer_by_split(dataset_name, split_name):
+    df = get_df(dataset_name)
+    if "split" in df.columns:
+        return df[df["split"] == split_name].head(10)
+    return df.head(10)
+
+def run_model_demo(dataset_name):
+    df = get_df(dataset_name)
     feature_cols = get_numeric_feature_cols(df)
-    if not feature_cols:
-        return None, "No numeric feature columns found."
-
-    work_df = df.dropna(subset=feature_cols + [TARGET_COL]).reset_index(drop=True)
-
+    
+    # Ensure target exists, fallback to noisy if clean is missing (though unlikely in your schema)
+    target = TARGET_COL if TARGET_COL in df.columns else df.filter(like="expval").columns[0]
+    
+    work_df = df.dropna(subset=feature_cols + [target]).reset_index(drop=True)
     X = work_df[feature_cols]
-    y = work_df[TARGET_COL]
-
-    if len(work_df) < 20:
-        return None, "Not enough rows for a stable demo."
+    y = work_df[target]
 
-    X_train, X_test, y_train, y_test = train_test_split(
-        X, y, test_size=0.2, random_state=42
-    )
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
 
     model = RandomForestRegressor(**MODEL_PARAMS)
     model.fit(X_train, y_train)
@@ -138,175 +91,121 @@ def train_model_demo(df: pd.DataFrame):
     r2 = r2_score(y_test, preds)
     mae = mean_absolute_error(y_test, preds)
 
-    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5))
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
 
-    ax1.scatter(y_test, preds, alpha=0.6)
-    min_v = min(float(y_test.min()), float(np.min(preds)))
-    max_v = max(float(y_test.max()), float(np.max(preds)))
-    ax1.plot([min_v, max_v], [min_v, max_v], linestyle="--")
-    ax1.set_xlabel("True value")
-    ax1.set_ylabel("Predicted value")
-    ax1.set_title(f"Predictions vs Truth\nR² = {r2:.4f}, MAE = {mae:.4f}")
+    # Parity Plot
+    ax1.scatter(y_test, preds, alpha=0.5, color='#636EFA')
+    lims = [min(y_test.min(), preds.min()), max(y_test.max(), preds.max())]
+    ax1.plot(lims, lims, 'r--', alpha=0.75, zorder=3)
+    ax1.set_xlabel("Ground Truth")
+    ax1.set_ylabel("Predictions")
+    ax1.set_title(f"Prediction Accuracy\nR² = {r2:.4f}")
 
+    # Feature Importance
     importances = model.feature_importances_
-    top_idx = np.argsort(importances)[-10:]
-    ax2.barh(range(len(top_idx)), importances[top_idx])
-    ax2.set_yticks(range(len(top_idx)))
-    ax2.set_yticklabels([feature_cols[i] for i in top_idx])
-    ax2.set_xlabel("Feature importance")
-    ax2.set_title("Top 10 features")
-
+    indices = np.argsort(importances)[-10:]
+    ax2.barh(range(len(indices)), importances[indices], color='#EF553B')
+    ax2.set_yticks(range(len(indices)))
+    ax2.set_yticklabels([feature_cols[i] for i in indices])
+    ax2.set_title("Top 10 Structural Features")
+    
     plt.tight_layout()
-
-    explanation = f"""
-**R² score:** {r2:.4f}  
-**MAE:** {mae:.4f}
-
-This is a lightweight baseline on the demo dataset. The point is not to get a perfect score, but to show that the dataset contains real structure and can support quantum ML experiments.
-"""
-
-    return fig, explanation
-
-
-# =========================================================
-# BENCHMARK TAB
-# =========================================================
-def make_bar_plot(df: pd.DataFrame, value_col: str, title: str, ylabel: str):
-    fig, ax = plt.subplots(figsize=(9, 4.8))
-    if df.empty or value_col not in df.columns or "dataset" not in df.columns:
-        ax.text(0.5, 0.5, "No benchmark data available", ha="center", va="center")
-        ax.axis("off")
-        return fig
-
-    plot_df = df.copy()
-    plot_df = plot_df.dropna(subset=[value_col])
-
-    ax.bar(plot_df["dataset"].astype(str), plot_df[value_col].astype(float))
-    ax.set_title(title)
-    ax.set_xlabel("Dataset")
-    ax.set_ylabel(ylabel)
-    ax.tick_params(axis="x", rotation=20)
-    ax.axhline(0, linewidth=1)
+    
+    summary = f"""
+    ### Model Performance: {dataset_name}
+    - **R² Score:** {r2:.4f}
+    - **Mean Absolute Error (MAE):** {mae:.4f}
+    
+    *This baseline demonstrates that structural circuit metrics (entropy, gate counts, etc.) hold predictive power for quantum expectation values.*
+    """
+    return fig, summary
+
+def load_benchmark():
+    path = Path(LOCAL_BENCHMARK_CSV)
+    if not path.exists():
+        return pd.DataFrame([{"info": "Benchmark file not found"}]), None, None
+    
+    df = pd.read_csv(path)
+    
+    # R2 Plot
+    fig_r2, ax = plt.subplots(figsize=(8, 4))
+    ax.bar(df["dataset"], df["r2"], color='skyblue')
+    ax.set_title("Cross-Dataset Robustness (R² Score)")
+    ax.set_ylabel("R²")
+    plt.xticks(rotation=15)
     plt.tight_layout()
-    return fig
-
-
-def build_benchmark_dashboard():
-    df = load_benchmark_results()
-
-    explanation = """
-### Noise robustness benchmark
-
-This dashboard shows how a model trained on clean circuits behaves on:
-- **core_clean**
-- **depolarizing**
-- **amplitude_damping**
-- **transpilation**
-
-A sharp drop in R² indicates strong distribution shift. That is exactly the value of the larger QSBench packs.
-"""
-
-    r2_fig = make_bar_plot(df, "r2", "Noise Robustness Benchmark — R²", "R²")
-    mae_fig = make_bar_plot(df, "mae", "Noise Robustness Benchmark — MAE", "MAE")
 
-    return df, r2_fig, mae_fig, explanation
+    # MAE Plot
+    fig_mae, ax = plt.subplots(figsize=(8, 4))
+    ax.bar(df["dataset"], df["mae"], color='salmon')
+    ax.set_title("Cross-Dataset Error (MAE)")
+    ax.set_ylabel("MAE")
+    plt.xticks(rotation=15)
+    plt.tight_layout()
 
+    return df, fig_r2, fig_mae
 
 # =========================================================
-# APP
+# INTERFACE
 # =========================================================
-def main():
-    print("Loading demo dataset...")
-    df_all = load_demo_dataset()
-    splits_cache = split_by_split_column(df_all)
-    split_choices = list(splits_cache.keys())
-
-    default_split = split_choices[0] if split_choices else None
-
-    with gr.Blocks(title="QSBench Demo Explorer") as demo:
-        gr.Markdown(
-            """
-# QSBench Demo Explorer
-
-Interactive demo for the QSBench Core demo dataset and precomputed noise robustness benchmark.
-"""
-        )
-
-        with gr.Tabs():
-            with gr.TabItem("Data Explorer"):
-                gr.Markdown("Inspect the demo dataset split by split.")
-                split_selector = gr.Dropdown(
-                    choices=split_choices,
-                    value=default_split,
-                    label="Choose a split",
-                )
-                data_table = gr.Dataframe(label="First 10 rows", interactive=False)
-
-                split_selector.change(
-                    fn=lambda s: show_data(s, splits_cache),
-                    inputs=split_selector,
-                    outputs=data_table,
-                )
-
-                demo.load(
-                    fn=lambda: show_data(default_split, splits_cache),
-                    inputs=[],
-                    outputs=data_table,
-                )
-
-            with gr.TabItem("Model Demo"):
-                gr.Markdown(
-                    """
-Train a lightweight Random Forest baseline on the demo data and inspect predictions.
-"""
-                )
-                train_button = gr.Button("Train model")
-                plot_output = gr.Plot()
-                text_output = gr.Markdown()
-
-                train_button.click(
-                    fn=lambda: train_model_demo(df_all),
-                    inputs=[],
-                    outputs=[plot_output, text_output],
-                )
-
-            with gr.TabItem("Noise Robustness Benchmark"):
-                gr.Markdown(
-                    """
-This tab loads the precomputed local benchmark results from `noise_benchmark_results.csv`.
-"""
-                )
-                refresh_button = gr.Button("Load benchmark results")
-                benchmark_table = gr.Dataframe(label="Benchmark results", interactive=False)
-                r2_plot = gr.Plot(label="R² plot")
-                mae_plot = gr.Plot(label="MAE plot")
-                benchmark_text = gr.Markdown()
-
-                refresh_button.click(
-                    fn=build_benchmark_dashboard,
-                    inputs=[],
-                    outputs=[benchmark_table, r2_plot, mae_plot, benchmark_text],
-                )
-
-                demo.load(
-                    fn=build_benchmark_dashboard,
-                    inputs=[],
-                    outputs=[benchmark_table, r2_plot, mae_plot, benchmark_text],
-                )
-
-        gr.Markdown("---")
-        gr.Markdown(
-            """
-### What this demo shows
-
-- Data Explorer: inspect the dataset splits
-- Model Demo: quick baseline on the demo data
-- Noise Robustness Benchmark: precomputed results that show how performance changes across clean, noisy, and transpiled datasets
-"""
-        )
+with gr.Blocks(title="QSBench Unified Explorer", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 🌌 QSBench: Quantum Synthetic Benchmark Explorer
+        **Unified interface for Core, Noise-Affected, and Hardware-Transpiled Quantum Datasets.**
+        
+        Browse the demo datasets from the QSBench family, run baseline ML models, and analyze noise robustness across different distributions.
+        """
+    )
 
-    demo.launch()
+    with gr.Tabs():
+        # TAB 1: DATA EXPLORER
+        with gr.TabItem("🔎 Dataset Explorer"):
+            with gr.Row():
+                ds_selector = gr.Dropdown(choices=list(DATASET_MAP.keys()), value="Core (Clean)", label="Select Dataset Pack")
+                split_selector = gr.Dropdown(choices=["train", "test", "validation"], value="train", label="Split")
+            
+            data_table = gr.Dataframe(label="Sample Data (First 10 rows)", interactive=False)
+            
+            ds_selector.change(update_explorer, inputs=[ds_selector], outputs=[split_selector, data_table])
+            split_selector.change(filter_explorer_by_split, inputs=[ds_selector, split_selector], outputs=[data_table])
+
+        # TAB 2: ML BASELINE
+        with gr.TabItem("🤖 ML Baseline Demo"):
+            gr.Markdown("Select a dataset and train a Random Forest regressor to predict expectation values from circuit metadata.")
+            model_ds_selector = gr.Dropdown(choices=list(DATASET_MAP.keys()), value="Core (Clean)", label="Target Dataset")
+            train_btn = gr.Button("Train Baseline Model", variant="primary")
+            
+            with gr.Row():
+                plot_output = gr.Plot(label="Model Metrics")
+                text_output = gr.Markdown(label="Stats")
+            
+            train_btn.click(run_model_demo, inputs=[model_ds_selector], outputs=[plot_output, text_output])
+
+        # TAB 3: BENCHMARKING
+        with gr.TabItem("📊 Noise Robustness Benchmark"):
+            gr.Markdown("Analysis of model performance degradation under distribution shifts (Clean → Noisy → Hardware).")
+            bench_btn = gr.Button("Load Precomputed Benchmark Results")
+            bench_table = gr.Dataframe(interactive=False)
+            with gr.Row():
+                r2_plot = gr.Plot()
+                mae_plot = gr.Plot()
+            
+            bench_btn.click(load_benchmark, outputs=[bench_table, r2_plot, mae_plot])
+
+    gr.Markdown(
+        """
+        ---
+        ### About QSBench
+        QSBench is a collection of high-quality synthetic datasets designed for **Quantum Machine Learning** research. 
+        It provides paired ideal/noisy data, structural circuit metrics, and transpilation metadata.
+        
+        🔗 [Website](https://qsbench.github.io) | 🤗 [Hugging Face](https://huggingface.co/QSBench) | 🛠️ [GitHub](https://github.com/QSBench)
+        """
+    )
 
+    # Initial load
+    demo.load(update_explorer, inputs=[ds_selector], outputs=[split_selector, data_table])
 
 if __name__ == "__main__":
-    main()
+    demo.launch()