Update app.py

app.py

@@ -1,5 +1,7 @@
 import ast
+import glob
 import logging
+import os
 import re
 from typing import Dict, List, Optional, Tuple
 
@@ -7,43 +9,43 @@ import gradio as gr
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
-from datasets import load_dataset
 from sklearn.ensemble import HistGradientBoostingClassifier
 from sklearn.impute import SimpleImputer
 from sklearn.inspection import permutation_importance
 from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, f1_score
 from sklearn.model_selection import train_test_split
 from sklearn.pipeline import Pipeline
-from sklearn.preprocessing import StandardScaler
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 APP_TITLE = "Noise Detection"
 APP_SUBTITLE = (
-    "Classify quantum circuits into clean, depolarizing, amplitude_damping, or hardware-aware noise conditions."
+    "Detect hardware-aware transpilation artifacts versus all other circuit conditions using structural circuit features."
 )
 
+DATA_DIR = os.getenv("QS_DATA_DIR", "data")
+
 REPO_CONFIG = {
     "clean": {
         "label": "clean",
-        "repo": "QSBench/QSBench-Core-v1.0.0-demo",
+        "path": os.getenv("QS_CLEAN_PATH", os.path.join(DATA_DIR, "core")),
     },
     "depolarizing": {
         "label": "depolarizing",
-        "repo": "QSBench/QSBench-Depolarizing-Demo-v1.0.0",
+        "path": os.getenv("QS_DEPOLARIZING_PATH", os.path.join(DATA_DIR, "depolarizing")),
     },
     "amplitude_damping": {
         "label": "amplitude_damping",
-        "repo": "QSBench/QSBench-Amplitude-v1.0.0-demo",
+        "path": os.getenv("QS_AMPLITUDE_PATH", os.path.join(DATA_DIR, "amplitude")),
     },
     "hardware_aware": {
         "label": "hardware_aware",
-        "repo": "QSBench/QSBench-Transpilation-v1.0.0-demo",
+        "path": os.getenv("QS_HARDWARE_AWARE_PATH", os.path.join(DATA_DIR, "transpilation")),
     },
 }
 
-CLASS_ORDER = ["clean", "depolarizing", "amplitude_damping", "hardware_aware"]
+CLASS_ORDER = ["other", "hardware_aware"]
 
 NON_FEATURE_COLS = {
     "sample_id",
@@ -66,12 +68,14 @@ NON_FEATURE_COLS = {
     "meyer_wallach",
     "cx_count",
     "noise_label",
+    "source_dataset",
+    "target_label",
 }
 
 SOFT_EXCLUDE_PATTERNS = ["ideal_", "noisy_", "error_", "sign_ideal_", "sign_noisy_"]
 
 _ASSET_CACHE: Dict[str, pd.DataFrame] = {}
-_COMBINED_CACHE: Optional[pd.DataFrame] = None
+_COMBINED_CACHE: Dict[Tuple[str, ...], pd.DataFrame] = {}
 
 
 def safe_parse(value):
@@ -94,6 +98,7 @@ def adjacency_features(adj_value) -> Dict[str, float]:
             "adj_degree_mean": np.nan,
             "adj_degree_std": np.nan,
         }
+
     try:
         arr = np.array(parsed, dtype=float)
         n = arr.shape[0]
@@ -126,6 +131,7 @@ def qasm_features(qasm_value) -> Dict[str, float]:
             "qasm_measure_count": np.nan,
             "qasm_comment_count": np.nan,
         }
+
     text = qasm_value
     lines = [line for line in text.splitlines() if line.strip()]
     gate_keywords = re.findall(
@@ -135,6 +141,7 @@ def qasm_features(qasm_value) -> Dict[str, float]:
     )
     measure_count = len(re.findall(r"\bmeasure\b", text, flags=re.IGNORECASE))
     comment_count = sum(1 for line in lines if line.strip().startswith("//"))
+
     return {
         "qasm_length": float(len(text)),
         "qasm_line_count": float(len(lines)),
@@ -147,37 +154,63 @@ def qasm_features(qasm_value) -> Dict[str, float]:
 def enrich_dataframe(df: pd.DataFrame) -> pd.DataFrame:
     """Add derived numeric features for classification."""
     df = df.copy()
+
     if "adjacency" in df.columns:
         adj_df = df["adjacency"].apply(adjacency_features).apply(pd.Series)
         df = pd.concat([df, adj_df], axis=1)
+
     qasm_source = "qasm_transpiled" if "qasm_transpiled" in df.columns else "qasm_raw"
     if qasm_source in df.columns:
         qasm_df = df[qasm_source].apply(qasm_features).apply(pd.Series)
         df = pd.concat([df, qasm_df], axis=1)
+
     return df
 
 
+def _resolve_path(dataset_key: str) -> str:
+    path = REPO_CONFIG[dataset_key]["path"]
+    if not os.path.exists(path):
+        raise FileNotFoundError(
+            f"Local dataset path not found for '{dataset_key}': {path}. "
+            "Set the matching environment variable or place the parquet directory at this path."
+        )
+    return path
+
+
+def _read_parquet_source(path: str) -> pd.DataFrame:
+    """Read a parquet file or a directory of parquet shards."""
+    if os.path.isdir(path):
+        files = sorted(glob.glob(os.path.join(path, "*.parquet")))
+        if not files:
+            raise FileNotFoundError(f"No parquet files found in directory: {path}")
+        frames = [pd.read_parquet(file_path) for file_path in files]
+        return pd.concat(frames, ignore_index=True)
+
+    return pd.read_parquet(path)
+
+
 def load_single_dataset(dataset_key: str) -> pd.DataFrame:
-    """Load a dataset shard from Hugging Face and cache it in memory."""
+    """Load a local parquet dataset and cache it in memory."""
     if dataset_key not in _ASSET_CACHE:
-        logger.info("Loading dataset: %s", dataset_key)
-        ds = load_dataset(REPO_CONFIG[dataset_key]["repo"])
-        df = pd.DataFrame(ds["train"])
+        path = _resolve_path(dataset_key)
+        logger.info("Loading local dataset: %s -> %s", dataset_key, path)
+        df = _read_parquet_source(path)
         df = enrich_dataframe(df)
         df["noise_label"] = REPO_CONFIG[dataset_key]["label"]
+        df["source_dataset"] = dataset_key
         _ASSET_CACHE[dataset_key] = df
     return _ASSET_CACHE[dataset_key]
 
 
-def load_combined_dataset() -> pd.DataFrame:
-    """Load and merge all four noise-condition datasets."""
-    global _COMBINED_CACHE
-    if _COMBINED_CACHE is None:
-        frames = [load_single_dataset(key) for key in REPO_CONFIG.keys()]
+def load_combined_dataset(dataset_keys: List[str]) -> pd.DataFrame:
+    """Load and merge selected local datasets."""
+    cache_key = tuple(sorted(dataset_keys))
+    if cache_key not in _COMBINED_CACHE:
+        frames = [load_single_dataset(key) for key in dataset_keys]
         combined = pd.concat(frames, ignore_index=True)
-        combined = combined[combined["noise_label"].isin(CLASS_ORDER)].copy()
-        _COMBINED_CACHE = combined
-    return _COMBINED_CACHE
+        combined = combined.copy()
+        _COMBINED_CACHE[cache_key] = combined
+    return _COMBINED_CACHE[cache_key]
 
 
 def load_guide_content() -> str:
@@ -232,6 +265,7 @@ def make_classification_figure(
     """Create a compact classification summary figure."""
     fig = plt.figure(figsize=(20, 6))
     gs = fig.add_gridspec(1, 3)
+
     ax1 = fig.add_subplot(gs[0, 0])
     ax2 = fig.add_subplot(gs[0, 1])
     ax3 = fig.add_subplot(gs[0, 2])
@@ -275,7 +309,7 @@ def build_dataset_profile(df: pd.DataFrame) -> str:
         f"### Dataset profile\n\n"
         f"**Rows:** {len(df):,}  \n"
         f"**Columns:** {len(df.columns):,}  \n"
-        f"**Classes:** {', '.join(CLASS_ORDER)}"
+        f"**Source label:** `{df['noise_label'].iloc[0] if 'noise_label' in df.columns and not df.empty else 'n/a'}`"
     )
 
 
@@ -285,20 +319,26 @@ def refresh_explorer(dataset_key: str, split_name: str) -> Tuple[gr.update, pd.D
     splits = df["split"].dropna().unique().tolist() if "split" in df.columns else ["train"]
     if not splits:
         splits = ["train"]
+
     if split_name not in splits:
         split_name = splits[0]
+
     filtered = df[df["split"] == split_name] if "split" in df.columns else df
     display_df = filtered.head(12).copy()
+
     raw_qasm = display_df["qasm_raw"].iloc[0] if "qasm_raw" in display_df.columns and not display_df.empty else "// N/A"
     transpiled_qasm = display_df["qasm_transpiled"].iloc[0] if "qasm_transpiled" in display_df.columns and not display_df.empty else "// N/A"
+
     profile_box = build_dataset_profile(df)
     summary_box = (
         f"### Split summary\n\n"
         f"**Dataset:** `{dataset_key}`  \n"
         f"**Label:** `{REPO_CONFIG[dataset_key]['label']}`  \n"
+        f"**Path:** `{REPO_CONFIG[dataset_key]['path']}`  \n"
         f"**Available splits:** {', '.join(splits)}  \n"
         f"**Preview rows:** {len(display_df)}"
     )
+
     return (
         gr.update(choices=splits, value=split_name),
         display_df,
@@ -309,36 +349,50 @@ def refresh_explorer(dataset_key: str, split_name: str) -> Tuple[gr.update, pd.D
     )
 
 
-def sync_feature_picker(dataset_key: str) -> gr.update:
-    """Refresh the feature list from the selected dataset."""
-    df = load_single_dataset(dataset_key)
+def sync_feature_picker(dataset_keys: List[str]) -> gr.update:
+    """Refresh the feature list from the selected datasets."""
+    if not dataset_keys:
+        return gr.update(choices=[], value=[])
+
+    df = load_combined_dataset(dataset_keys)
     features = get_available_feature_columns(df)
     defaults = default_feature_selection(features)
     return gr.update(choices=features, value=defaults)
 
 
 def train_classifier(
-    dataset_key: str,
+    dataset_keys: List[str],
     feature_columns: List[str],
     test_size: float,
     n_estimators: int,
     max_depth: float,
     random_state: float,
 ) -> Tuple[Optional[plt.Figure], str]:
-    """Train a four-class classifier with better handling of class imbalance."""
+    """Train a binary classifier for hardware-aware detection."""
+    if not dataset_keys:
+        return None, "### ❌ Please select at least one dataset."
+
     if not feature_columns:
         return None, "### ❌ Please select at least one feature."
 
-    df = load_single_dataset(dataset_key)
-    required_cols = feature_columns + ["noise_label"]
-    train_df = df.dropna(subset=required_cols).copy()
-    train_df = train_df[train_df["noise_label"].isin(CLASS_ORDER)]
+    df = load_combined_dataset(dataset_keys).copy()
+    df["target_label"] = np.where(df["source_dataset"] == "hardware_aware", "hardware_aware", "other")
+
+    if "target_label" not in df.columns:
+        return None, "### ❌ Target label could not be created."
+
+    train_df = df.dropna(subset=["target_label"]).copy()
 
     if len(train_df) < 20:
         return None, "### ❌ Not enough rows after filtering missing values."
 
-    X = train_df[feature_columns]
-    y = train_df["noise_label"]
+    X = train_df[feature_columns].copy()
+    X = X.dropna(axis=1, how="all")
+    if X.shape[1] == 0:
+        return None, "### ❌ All selected features are empty in the chosen datasets."
+
+    feature_columns = X.columns.tolist()
+    y = train_df["target_label"]
 
     seed = int(random_state)
     depth = int(max_depth) if max_depth and int(max_depth) > 0 else None
@@ -346,17 +400,23 @@ def train_classifier(
 
     try:
         X_train, X_test, y_train, y_test = train_test_split(
-            X, y, test_size=test_size, random_state=seed, stratify=y
+            X,
+            y,
+            test_size=test_size,
+            random_state=seed,
+            stratify=y,
         )
     except ValueError:
         X_train, X_test, y_train, y_test = train_test_split(
-            X, y, test_size=test_size, random_state=seed
+            X,
+            y,
+            test_size=test_size,
+            random_state=seed,
         )
 
     model = Pipeline(
         steps=[
             ("imputer", SimpleImputer(strategy="median")),
-            ("scaler", StandardScaler()),
             (
                 "classifier",
                 HistGradientBoostingClassifier(
@@ -390,7 +450,7 @@ def train_classifier(
     )
     importances = perm.importances_mean
 
-    fig = make_classification_figure(y_test.to_numpy(), y_pred, CLASS_ORDER, list(feature_columns), importances)
+    fig = make_classification_figure(y_test.to_numpy(), y_pred, CLASS_ORDER, feature_columns, importances)
 
     report = classification_report(
         y_test,
@@ -401,7 +461,7 @@ def train_classifier(
     results = (
         "### Classification results\n\n"
         f"**Rows used:** {len(train_df):,}  \n"
-        f"**Dataset:** `{dataset_key}`  \n"
+        f"**Datasets used:** {', '.join(dataset_keys)}  \n"
         f"**Test size:** {test_size:.0%}  \n"
         f"**Accuracy:** {accuracy:.4f}  \n"
         f"**Macro F1:** {macro_f1:.4f}  \n"
@@ -447,10 +507,10 @@ with gr.Blocks(title=APP_TITLE) as demo:
                 transpiled_qasm = gr.Code(label="Transpiled QASM", language=None)
 
         with gr.TabItem("🧠 Classification"):
-            class_dataset_dropdown = gr.Dropdown(
-                list(REPO_CONFIG.keys()),
-                value="clean",
-                label="Dataset",
+            class_dataset_picker = gr.CheckboxGroup(
+                label="Datasets",
+                choices=list(REPO_CONFIG.keys()),
+                value=list(REPO_CONFIG.keys()),
             )
             feature_picker = gr.CheckboxGroup(label="Input features", choices=[])
             test_size = gr.Slider(0.1, 0.4, value=0.2, step=0.05, label="Test split")
@@ -483,11 +543,11 @@ with gr.Blocks(title=APP_TITLE) as demo:
         [split_dropdown, explorer_df, raw_qasm, transpiled_qasm, profile_box, summary_box],
     )
 
-    class_dataset_dropdown.change(sync_feature_picker, [class_dataset_dropdown], [feature_picker])
+    class_dataset_picker.change(sync_feature_picker, [class_dataset_picker], [feature_picker])
 
     run_btn.click(
         train_classifier,
-        [class_dataset_dropdown, feature_picker, test_size, n_estimators, max_depth, seed],
+        [class_dataset_picker, feature_picker, test_size, n_estimators, max_depth, seed],
         [plot, metrics],
     )
 
@@ -496,7 +556,7 @@ with gr.Blocks(title=APP_TITLE) as demo:
         [dataset_dropdown, split_dropdown],
         [split_dropdown, explorer_df, raw_qasm, transpiled_qasm, profile_box, summary_box],
     )
-    demo.load(sync_feature_picker, [class_dataset_dropdown], [feature_picker])
+    demo.load(sync_feature_picker, [class_dataset_picker], [feature_picker])
 
 
 if __name__ == "__main__":