Spaces:

zadid6pretam
/

BSTabDiff-Studio

Running on Zero

Update app.py

2c68822 verified 8 days ago

23.5 kB

"""BSTabDiff Studio: uploaded numeric CSV features + optional text/numeric label.
Auto-detect common class targets, encode them for BSTabDiff and decode on export.
Paper-default 10,000 epoch diffusion prior; dataset-aware latent block count M.
No GO-BS or Optuna. Shared-hosting limits apply.
"""
from __future__ import annotations

import html
import os
import tempfile
import time
from pathlib import Path

# ZeroGPU requires this import before any library that might touch CUDA.
import spaces

import numpy as np
import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import gradio as gr
import torch
from bstabdiff import BSTabDiff
from sklearn.decomposition import PCA
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler, LabelEncoder

MAX_FILE_MB = 50
MAX_ROWS = 10_000
MAX_FEATURES = 50_000
MAX_CELLS = 2_000_000         # n × m input budget
MAX_SYNTHETIC = 3_000
MAX_OUTPUT_CELLS = 2_000_000  # generated rows × m output budget
PAPER_EPOCHS = 10_000         # NeurIPS Appendix A16 default, not constructor default
PRIOR_BATCH = 128
PRIOR_LR = 1e-3
EMA_DECAY = 0.999
NO_TARGET = "None (unconditional)"
MAX_PREVIEW_COLUMNS = 16
MAX_PLOT_CHOICES = 150

CSS = """
.gradio-container {max-width: 1380px !important;}
#studio-hero {padding: 28px 30px; margin: 3px 0 18px; border-radius: 24px;
 background: linear-gradient(116deg,#182b6e 0%,#6037aa 55%,#c2488d 100%);
 color:white; box-shadow: 0 12px 28px rgba(57,40,131,.17);}
#studio-hero h1 {font-size: clamp(2rem,4vw,3.4rem); margin:0 0 6px; color:#fff !important; font-weight:800;}
#studio-hero p {font-size: 1rem; margin: 0; opacity:.95; color:#fff !important;}
#studio-hero .sub, #studio-hero .eyebrow {color:#fff !important;}
#studio-hero .sub {margin-top:12px;font-size:.85rem;opacity:.95;}
.studio-section {border-radius:18px !important;}
#generate-btn {background: linear-gradient(90deg,#5236be,#cf4f96) !important;
 color:#fff !important; border:0 !important; font-weight:800 !important;}
#generate-btn:hover {filter:brightness(1.08);}
#studio-metrics {line-height:1.8;}
#studio-metrics .metric-row {display:grid;grid-template-columns:repeat(4,minmax(0,1fr));gap:8px;margin-top:10px;}
#studio-metrics .metric {background:var(--block-background-fill);border:1px solid var(--block-border-color);
 padding:12px;border-radius:12px;text-align:center;}
#studio-metrics .metric b{display:block;font-size:1.4rem;color:#6653c8;}
#studio-metrics .metric span{font-size:.77rem;opacity:.8;}
@media(max-width:700px){#studio-hero{padding:21px 17px;}#studio-metrics .metric-row{grid-template-columns:1fr 1fr;}}
"""


def read_csv(path: str) -> pd.DataFrame:
    if not path:
        raise gr.Error("Upload a CSV file first.")
    file = Path(path)
    if file.suffix.lower() != ".csv":
        raise gr.Error("Only .csv files are accepted.")
    if file.stat().st_size > MAX_FILE_MB * 1024**2:
        raise gr.Error(f"Please upload a CSV smaller than {MAX_FILE_MB} MB.")
    try:
        df = pd.read_csv(path, low_memory=False)
    except Exception as exc:
        raise gr.Error(f"Could not read CSV: {exc}") from exc
    if df.empty or len(df.columns) < 2:
        raise gr.Error("The CSV must contain rows and at least two columns.")
    if len(df) < 20:
        raise gr.Error("Use at least 20 rows. A single or tiny sample is not suitable for fitting a generator.")
    if len(df) > MAX_ROWS:
        raise gr.Error(f"For this public demo, upload at most {MAX_ROWS:,} rows.")
    if len(df.columns) - 1 > MAX_FEATURES:
        raise gr.Error(f"For this public demo, upload at most {MAX_FEATURES:,} feature columns.")
    if df.columns.duplicated().any():
        raise gr.Error("Column names must be unique.")
    return df


def numeric_columns(df):
    return [c for c in df.columns if pd.api.types.is_numeric_dtype(df[c])]


def shape_card(df, target=NO_TARGET):
    n = len(df)
    m = len(df.columns) - (target != NO_TARGET and target in df.columns)
    rho = m / n
    kappa = n * m
    M = select_default_m(n, m)
    x = df.drop(columns=[target]) if target != NO_TARGET and target in df else df
    missing = x.isna().sum().sum() / max(x.size, 1)
    regime = "HDLSS (m > n)" if m > n else "n ≥ m"
    return f"""<div id='studio-metrics'><div class='metric-row'>
      <div class='metric'><b>{n:,}</b><span>Rows (n)</span></div>
      <div class='metric'><b>{m:,}</b><span>Features (m)</span></div>
      <div class='metric'><b>{rho:.2f}</b><span>Ratio ρ = m/n</span></div>
      <div class='metric'><b>{M:,}</b><span>Auto-selected M</span></div>
      </div><p><b>Data scale κ = nm:</b> {kappa:,} · <b>Missing feature cells:</b> {missing:.1%}
      · <b>Regime:</b> {regime}</p></div>"""


def detect_target(df: pd.DataFrame) -> str:
    """Conservatively auto-select a *named* classification target.
    Do not infer a target solely because a column happens to contain strings:
    a categorical feature could otherwise be mistaken for the label.
    """
    preferred = ("label", "target", "class", "class_label", "classlabel", "outcome", "y")
    for name in preferred:
        for col in df.columns:
            key = str(col).strip().lower().replace(" ", "_").replace("-", "_")
            if key != name:
                continue
            values = df[col]
            counts = values.value_counts(dropna=False)
            if 2 <= len(counts) <= 15 and not values.isna().any() and counts.min() >= 2:
                return str(col)
    return NO_TARGET


def inspect_csv(path):
    if not path:
        return (gr.update(choices=[NO_TARGET], value=NO_TARGET),
                gr.update(choices=[], value=None), "Upload a dataset to get started.", pd.DataFrame())
    df = read_csv(path)
    numeric = numeric_columns(df)
    inferred_target = detect_target(df)
    choices = [NO_TARGET] + [str(c) for c in df.columns]
    feature_choices = [str(c) for c in numeric if str(c) != inferred_target][:MAX_PLOT_CHOICES]
    summary = shape_card(df, inferred_target)
    nonnumeric_features = [str(c) for c in df.columns
                           if str(c) != inferred_target and not pd.api.types.is_numeric_dtype(df[c])]
    if inferred_target != NO_TARGET:
        target_msg = ("<strong>Auto-selected class target:</strong> <code>" +
                      html.escape(inferred_target) +
                      "</code>. Change it in the dropdown if needed. ")
    else:
        target_msg = ("<strong>Unconditional generation selected.</strong> "
                      "Choose a class target from the dropdown if you have one. ")
    if nonnumeric_features:
        target_msg += ("<strong>Other nonnumeric feature columns:</strong> " +
                       html.escape(", ".join(nonnumeric_features[:6])) +
                       ". Encode or remove these, or select the correct target column.")
    else:
        target_msg += "All remaining feature columns are numeric."
    summary += "<p>" + target_msg + "</p>"
    return (gr.update(choices=choices, value=inferred_target),
            gr.update(choices=feature_choices,
                      value=feature_choices[0] if feature_choices else None),
            summary, df.iloc[:8, :MAX_PREVIEW_COLUMNS])


def update_target_summary(path, target):
    """Keep feature counts accurate when a visitor changes target selection."""
    if not path:
        return "Upload a dataset to get started."
    df = read_csv(path)
    if target != NO_TARGET and target not in df.columns:
        return shape_card(df)
    nonnumeric = [str(col) for col in df.columns
                  if str(col) != target and not pd.api.types.is_numeric_dtype(df[col])]
    msg = ("<strong>Class target:</strong> " + html.escape(target) + ". ")
    if nonnumeric:
        msg += "Encode or remove nonnumeric features: " + html.escape(", ".join(nonnumeric[:6]))
    else:
        msg += "All remaining feature columns are numeric."
    return shape_card(df, target) + "<p>" + msg + "</p>"


def select_default_m(n, m):
    # Dataset-aware no-Optuna initialization from the BSTabDiff project README.
    if m <= 64:
        return m
    rho, kappa = m/n, n*m
    if rho < 10:
        return min(64, m)
    return min(32 if kappa < 250_000 else
               64 if kappa < 1_000_000 else
               128 if kappa < 3_000_000 else 192, m)


def validate_features(df, target):
    if target != NO_TARGET and target not in df.columns:
        raise gr.Error("Select an existing target column or None.")
    feature_df = df.drop(columns=[target]) if target != NO_TARGET else df.copy()
    if feature_df.shape[1] < 2:
        raise gr.Error("At least two numerical feature columns are required.")
    invalid = [str(c) for c in feature_df.columns if not pd.api.types.is_numeric_dtype(feature_df[c])]
    if invalid:
        names = ", ".join(invalid[:4])
        raise gr.Error(f"Feature columns must be numerical for this first demo. Encode or remove: {names}.")
    if feature_df.shape[1] > MAX_FEATURES:
        raise gr.Error(f"Please use at most {MAX_FEATURES:,} feature columns.")
    if len(df) * feature_df.shape[1] > MAX_CELLS:
        raise gr.Error(f"Shared-hardware limit: n × m must be at most {MAX_CELLS:,}. Try a smaller CSV.")
    X = feature_df.to_numpy(dtype=np.float32)
    if np.isinf(X).any():
        raise gr.Error("Infinite values are not supported; replace ±inf with missing values first.")
    if not np.isfinite(X).any():
        raise gr.Error("The CSV contains no finite numerical feature values.")
    all_missing = np.all(np.isnan(X), axis=0)
    if all_missing.any():
        raise gr.Error(f"Remove entirely missing feature column: {feature_df.columns[np.flatnonzero(all_missing)[0]]}")
    y = None
    if target != NO_TARGET:
        y = df[target]
        if y.isna().any():
            raise gr.Error("The target column must not contain missing class labels.")
        if y.nunique() < 2 or y.nunique() > 15:
            raise gr.Error("Select a classification target with 2–15 classes, or choose None.")
        if y.value_counts().min() < 2:
            raise gr.Error("Every class must have at least two rows for this demo.")
    return feature_df, X, y


def chart_layout(fig, title):
    fig.update_layout(template="plotly_white", title=dict(text=title, x=.02),
                      paper_bgcolor="rgba(0,0,0,0)", plot_bgcolor="rgba(0,0,0,0)",
                      margin=dict(l=28,r=18,t=55,b=35), font=dict(size=12),
                      legend=dict(orientation="h", y=-.2), height=390)
    return fig


def distribution_plot(real, synth, column):
    if not column or column not in real.columns:
        return go.Figure()
    fig = go.Figure()
    for name, frame, color in [("Real", real, "#536ed7"), ("Synthetic", synth, "#f06f9c")]:
        values = pd.to_numeric(frame[column], errors="coerce").replace([np.inf,-np.inf],np.nan).dropna()
        fig.add_trace(go.Histogram(x=values, name=name, marker_color=color,
                                   opacity=.58, nbinsx=30, histnorm="probability density"))
    fig.update_layout(barmode="overlay", xaxis_title=str(column), yaxis_title="Density")
    return chart_layout(fig, f"Feature distribution · {column}")


def pca_plot(real, synth):
    # A joint coordinate system fitted on real training data only.
    # Limit feature count to avoid a prohibitively large SVD for 50K features.
    cols = real.var(skipna=True).fillna(0).nlargest(min(200, real.shape[1])).index
    nr, ns = min(500, len(real)), min(500, len(synth))
    R = real.loc[:, cols].head(nr).to_numpy(dtype=float)
    S = synth.loc[:, cols].head(ns).to_numpy(dtype=float)
    imp = SimpleImputer(strategy="median", keep_empty_features=True)
    R1 = imp.fit_transform(R)
    S1 = imp.transform(S)
    scale = StandardScaler()
    R2 = scale.fit_transform(R1)
    S2 = scale.transform(S1)
    pca = PCA(n_components=2, svd_solver="randomized", random_state=42)
    rp = pca.fit_transform(R2)
    sp = pca.transform(S2)
    fig = go.Figure()
    for name, pts, color in [("Real",rp,"#536ed7"),("Synthetic",sp,"#f06f9c")]:
        fig.add_trace(go.Scatter(x=pts[:,0],y=pts[:,1],mode="markers",name=name,
                                 marker=dict(size=7,color=color,opacity=.63)))
    fig.update_layout(xaxis_title="PC 1",yaxis_title="PC 2")
    return chart_layout(fig,"Shared PCA projection · training-data diagnostic")


def correlation_plot(real, synth):
    variance = real.var(skipna=True).fillna(0).sort_values(ascending=False)
    cols = list(variance.head(min(18,len(variance))).index)
    if len(cols) < 2:
        return go.Figure()
    rc = real[cols].corr().fillna(0)
    sc = synth[cols].corr().fillna(0)
    fig = make_subplots(rows=1,cols=2, subplot_titles=("Real", "Synthetic"),horizontal_spacing=.12)
    for idx, mat in enumerate((rc,sc),start=1):
        fig.add_trace(go.Heatmap(z=mat.values,x=cols,y=cols,zmin=-1,zmax=1,
                                 colorscale="RdBu", reversescale=True,
                                 showscale=idx==2,colorbar=dict(title="r") if idx==2 else None),
                      row=1,col=idx)
    fig.update_layout(height=490)
    fig.update_xaxes(tickangle=-45, tickfont=dict(size=9))
    fig.update_yaxes(tickfont=dict(size=9))
    chart_layout(fig,"Feature correlations · highest-variance columns")
    fig.update_layout(height=490,margin=dict(l=28,r=25,t=65,b=65))
    return fig


@spaces.GPU(duration=180)
def generate(path, target, n_synth, seed, feature):
    df = read_csv(path)
    features, X, y = validate_features(df,target)
    n_synth = int(n_synth)
    seed = int(seed)
    if n_synth < 10 or n_synth > MAX_SYNTHETIC:
        raise gr.Error(f"Choose between 10 and {MAX_SYNTHETIC} synthetic rows.")
    if not 0 <= seed <= 2**31-1:
        raise gr.Error("Use a nonnegative random seed smaller than 2,147,483,648.")
    n, m = X.shape
    if n_synth * m > MAX_OUTPUT_CELLS:
        allowed = MAX_OUTPUT_CELLS // m
        raise gr.Error(
            f"Output would contain {n_synth * m:,} cells. "
            f"With {m:,} features, generate at most {allowed:,} rows "
            f"(maximum {MAX_OUTPUT_CELLS:,} output cells)."
        )
    M = select_default_m(n, m)
    # Map arbitrary user class names (e.g. control/treated) to integer IDs
    # for the generator, then restore the original labels in the exported CSV.
    encoder = None
    y_fit = None
    if y is not None:
        encoder = LabelEncoder()
        y_fit = encoder.fit_transform(y.to_numpy())
    started = time.perf_counter()
    torch.set_num_threads(min(2,os.cpu_count() or 2))
    try:
        # Fixed paper-default diffusion settings, no Optuna, no ordering.
        model = BSTabDiff(
            M=M,
            prior_type="diffusion",
            device="cuda" if torch.cuda.is_available() else "cpu",
            random_state=seed,
            prior_epochs=PAPER_EPOCHS,
            prior_batch=PRIOR_BATCH,
            prior_lr=PRIOR_LR,
            use_ema=True,
            ema_decay=EMA_DECAY,
            verbose_every=0,
            save_dir=None,
            use_gobs=False,
            use_gobs_fc=False,
        )
        model.fit(features, y_fit)
        X_syn,R_syn,y_syn = model.sample(n_samples=n_synth,return_mask=True)
    except Exception as exc:
        raise gr.Error(f"BSTabDiff training or sampling failed: {type(exc).__name__}: {exc}") from exc
    X_syn = np.asarray(X_syn,dtype=np.float32)
    if X_syn.shape != (n_synth,m):
        raise gr.Error(f"Unexpected generated shape {X_syn.shape}; expected {(n_synth,m)}.")
    if R_syn is not None:
        mask = np.asarray(R_syn)
        if mask.shape == X_syn.shape:
            X_syn = X_syn.copy()
            X_syn[mask == 0] = np.nan
    syn = pd.DataFrame(X_syn,columns=features.columns)
    if y is not None:
        if y_syn is None:
            raise gr.Error("The generator did not return labels for the selected class target.")
        sampled_codes = np.asarray(y_syn).reshape(-1)
        if len(sampled_codes) != n_synth:
            raise gr.Error("Generated class label count does not match synthetic rows.")
        try:
            sampled_codes = sampled_codes.astype(int)
        except (TypeError, ValueError, OverflowError) as exc:
            raise gr.Error("Generated class labels could not be decoded.") from exc
        if np.any(sampled_codes < 0) or np.any(sampled_codes >= len(encoder.classes_)):
            raise gr.Error("Generator returned an unknown class ID.")
        syn[target] = encoder.inverse_transform(sampled_codes)
    real = features.reset_index(drop=True)
    syn_features = syn[features.columns]
    plot_columns = list(features.columns[:MAX_PLOT_CHOICES])
    column = feature if feature in plot_columns else plot_columns[0]
    path_csv = os.path.join(tempfile.mkdtemp(prefix="bstabdiff_studio_"),"bstabdiff_synthetic.csv")
    syn.to_csv(path_csv,index=False)
    seconds = time.perf_counter()-started
    summary = (f"### ✅ Generated **{n_synth:,}** synthetic samples\n"
               f"**Input:** {n:,} rows × {m:,} features · "
               f"**ρ=m/n:** {m/n:.3f} · **κ=nm:** {n*m:,} · "
               f"**Auto-selected M:** {M}\n\n"
               f"**Diffusion epochs:** {PAPER_EPOCHS:,} · "
               f"**Batch:** {PRIOR_BATCH} · **LR:** {PRIOR_LR:g} · "
               f"**EMA decay:** {EMA_DECAY} · **Seed:** {seed} · "
               f"**Device:** {model.device_} · **Elapsed:** {seconds:.1f}s\n\n"
               "*Paper-default configuration with dataset-aware M, no Optuna or GO-BS. "
               "Plots compare against uploaded **training data**; this is not "
               "a held-out utility or privacy evaluation.*")
    # The input frame may contain a class target; output a privacy-minded 10-row preview only.
    state=(real,syn_features)
    return (summary,syn.head(10),distribution_plot(*state,column),
            pca_plot(*state),correlation_plot(*state),path_csv,state,
            gr.update(choices=plot_columns,value=column))


def update_plot(feature, state):
    if not state or not feature:
        return go.Figure()
    real, synthetic = state
    return distribution_plot(real,synthetic,feature)


with gr.Blocks(title="BSTabDiff Studio — NeurIPS 2026",
               delete_cache=(3600,3600)) as demo:
    gr.HTML("""
      <div id="studio-hero">
        <div class="eyebrow" style="color:#fff;font-weight:700;letter-spacing:.13em;font-size:.74rem;opacity:.93">NEURIPS 2026 · RESEARCH DEMO</div>
        <h1>✦ BSTabDiff Studio</h1>
        <p>From your CSV to high-dimensional synthetic data.</p>
        <div class="sub" style="color:#fff">Upload → Generate → Visualize → Download · Default diffusion, no hyperparameter tuning.</div>
      </div>
    """)
    gr.Markdown("**How it works:** Supply a CSV containing numeric features and optionally a discrete class label. "
                "This Space trains BSTabDiff on your uploaded table and generates new synthetic rows. "
                "**Please do not upload confidential, patient-level, or personally identifying data.**")

    with gr.Row(equal_height=False):
        with gr.Column(scale=5,elem_classes=["studio-section"]):
            gr.Markdown("### ① Upload & configure")
            file_input=gr.File(label="Upload CSV",file_types=[".csv"],type="filepath")
            gr.Markdown(f"*Limits: {MAX_FILE_MB} MB, {MAX_ROWS:,} rows, "
                        f"{MAX_FEATURES:,} features, {MAX_CELLS:,} input cells; "
                        f"{MAX_OUTPUT_CELLS:,} synthetic output cells. "
                        f"Feature explorer lists the first {MAX_PLOT_CHOICES} numeric columns.*")
            target=gr.Dropdown(label="Optional class target",choices=[NO_TARGET],value=NO_TARGET,
                               info="Common column names like label/target/class are automatically selected. Text labels are supported; remaining features must be numeric.")
            with gr.Row():
                samples=gr.Slider(label="Synthetic rows",minimum=10,maximum=MAX_SYNTHETIC,
                                  step=10,value=100)
                seed_input=gr.Number(label="Random seed",value=42,precision=0,minimum=0,maximum=9999)
            gr.Markdown(f"**Paper-default generator:** {PAPER_EPOCHS:,} diffusion epochs, "
                        f"batch {PRIOR_BATCH}, LR {PRIOR_LR:g}, EMA {EMA_DECAY}; "
                        "M is chosen automatically from n, m, ρ and κ. "
                        "No GO-BS or Optuna. Large datasets may exceed ZeroGPU runtime limits.")
            run_btn=gr.Button("✦ Generate synthetic dataset",variant="primary",elem_id="generate-btn",size="lg")
        with gr.Column(scale=6,elem_classes=["studio-section"]):
            gr.Markdown("### ② Dataset snapshot")
            file_info=gr.HTML("<p>Upload a CSV to inspect its structure.</p>")
            preview=gr.Dataframe(label=f"Original data · first 8 rows / {MAX_PREVIEW_COLUMNS} columns",interactive=False,
                                 max_height=310)

    gr.Markdown("## ③ Explore your generated data")
    result=gr.Markdown("The synthetic-data dashboard will appear here after generation.")
    state=gr.State(value=None)
    with gr.Row():
        real_plot_feature=gr.Dropdown(label="Feature to inspect",choices=[],interactive=True)
        download=gr.File(label="Download synthetic CSV",interactive=False)
    with gr.Row():
        with gr.Column():
            with gr.Tab("📊 Feature distribution"):
                hist=gr.Plot(label="Real vs synthetic feature distribution")
            with gr.Tab("🧭 PCA projection"):
                pca=gr.Plot(label="Shared PCA projection")
        with gr.Column():
            with gr.Tab("🧩 Correlations"):
                corr=gr.Plot(label="Real and synthetic correlation matrices")
            with gr.Tab("🧬 Synthetic preview"):
                generated_table=gr.Dataframe(label="Synthetic data · first 10 rows",interactive=False)

    gr.HTML("""<div style="text-align:center;padding:16px;font-size:12px;opacity:.75">
      <a href="https://www.zadidhabib.com/bstabdiff.html" target="_blank">Project webpage</a> ·
      <a href="https://github.com/zadid6pretam/BSTabDiff" target="_blank">GitHub</a> ·
      <a href="https://pypi.org/project/bstabdiff/" target="_blank">PyPI</a> ·
      <a href="https://neurips.cc/virtual/2026/loc/atlanta/poster/151140" target="_blank" rel="noopener noreferrer">NeurIPS Poster</a> ·
      <a href="https://openreview.net/forum?id=ciPCMXX1aj" target="_blank">NeurIPS 2026 Paper</a>
      <p>Research demonstration only. Synthetic data are not necessarily private or faithful to the source distribution.</p>
      </div>""")

    file_input.change(inspect_csv,inputs=file_input,
                      outputs=[target,real_plot_feature,file_info,preview],show_progress="minimal")
    target.change(update_target_summary, inputs=[file_input, target], outputs=file_info,
                  show_progress="hidden")
    run_btn.click(generate,inputs=[file_input,target,samples,seed_input,real_plot_feature],
                  outputs=[result,generated_table,hist,pca,corr,download,state,real_plot_feature],
                  concurrency_limit=1,show_progress="full")
    real_plot_feature.change(update_plot,inputs=[real_plot_feature,state],outputs=hist,
                             show_progress="minimal")

demo.queue(max_size=6)

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(primary_hue="indigo",secondary_hue="pink"), css=CSS)

	"""BSTabDiff Studio: uploaded numeric CSV features + optional text/numeric label.

	Auto-detect common class targets, encode them for BSTabDiff and decode on export.
	Paper-default 10,000 epoch diffusion prior; dataset-aware latent block count M.
	No GO-BS or Optuna. Shared-hosting limits apply.
	"""
	from __future__ import annotations

	import html
	import os
	import tempfile
	import time
	from pathlib import Path

	# ZeroGPU requires this import before any library that might touch CUDA.
	import spaces

	import numpy as np
	import pandas as pd
	import plotly.graph_objects as go
	from plotly.subplots import make_subplots
	import gradio as gr
	import torch
	from bstabdiff import BSTabDiff
	from sklearn.decomposition import PCA
	from sklearn.impute import SimpleImputer
	from sklearn.preprocessing import StandardScaler, LabelEncoder

	MAX_FILE_MB = 50
	MAX_ROWS = 10_000
	MAX_FEATURES = 50_000
	MAX_CELLS = 2_000_000 # n × m input budget
	MAX_SYNTHETIC = 3_000
	MAX_OUTPUT_CELLS = 2_000_000 # generated rows × m output budget
	PAPER_EPOCHS = 10_000 # NeurIPS Appendix A16 default, not constructor default
	PRIOR_BATCH = 128
	PRIOR_LR = 1e-3
	EMA_DECAY = 0.999
	NO_TARGET = "None (unconditional)"
	MAX_PREVIEW_COLUMNS = 16
	MAX_PLOT_CHOICES = 150

	CSS = """
	.gradio-container {max-width: 1380px !important;}
	#studio-hero {padding: 28px 30px; margin: 3px 0 18px; border-radius: 24px;
	background: linear-gradient(116deg,#182b6e 0%,#6037aa 55%,#c2488d 100%);
	color:white; box-shadow: 0 12px 28px rgba(57,40,131,.17);}
	#studio-hero h1 {font-size: clamp(2rem,4vw,3.4rem); margin:0 0 6px; color:#fff !important; font-weight:800;}
	#studio-hero p {font-size: 1rem; margin: 0; opacity:.95; color:#fff !important;}
	#studio-hero .sub, #studio-hero .eyebrow {color:#fff !important;}
	#studio-hero .sub {margin-top:12px;font-size:.85rem;opacity:.95;}
	.studio-section {border-radius:18px !important;}
	#generate-btn {background: linear-gradient(90deg,#5236be,#cf4f96) !important;
	color:#fff !important; border:0 !important; font-weight:800 !important;}
	#generate-btn:hover {filter:brightness(1.08);}
	#studio-metrics {line-height:1.8;}
	#studio-metrics .metric-row {display:grid;grid-template-columns:repeat(4,minmax(0,1fr));gap:8px;margin-top:10px;}
	#studio-metrics .metric {background:var(--block-background-fill);border:1px solid var(--block-border-color);
	padding:12px;border-radius:12px;text-align:center;}
	#studio-metrics .metric b{display:block;font-size:1.4rem;color:#6653c8;}
	#studio-metrics .metric span{font-size:.77rem;opacity:.8;}
	@media(max-width:700px){#studio-hero{padding:21px 17px;}#studio-metrics .metric-row{grid-template-columns:1fr 1fr;}}
	"""


	def read_csv(path: str) -> pd.DataFrame:
	if not path:
	raise gr.Error("Upload a CSV file first.")
	file = Path(path)
	if file.suffix.lower() != ".csv":
	raise gr.Error("Only .csv files are accepted.")
	if file.stat().st_size > MAX_FILE_MB * 1024**2:
	raise gr.Error(f"Please upload a CSV smaller than {MAX_FILE_MB} MB.")
	try:
	df = pd.read_csv(path, low_memory=False)
	except Exception as exc:
	raise gr.Error(f"Could not read CSV: {exc}") from exc
	if df.empty or len(df.columns) < 2:
	raise gr.Error("The CSV must contain rows and at least two columns.")
	if len(df) < 20:
	raise gr.Error("Use at least 20 rows. A single or tiny sample is not suitable for fitting a generator.")
	if len(df) > MAX_ROWS:
	raise gr.Error(f"For this public demo, upload at most {MAX_ROWS:,} rows.")
	if len(df.columns) - 1 > MAX_FEATURES:
	raise gr.Error(f"For this public demo, upload at most {MAX_FEATURES:,} feature columns.")
	if df.columns.duplicated().any():
	raise gr.Error("Column names must be unique.")
	return df


	def numeric_columns(df):
	return [c for c in df.columns if pd.api.types.is_numeric_dtype(df[c])]


	def shape_card(df, target=NO_TARGET):
	n = len(df)
	m = len(df.columns) - (target != NO_TARGET and target in df.columns)
	rho = m / n
	kappa = n * m
	M = select_default_m(n, m)
	x = df.drop(columns=[target]) if target != NO_TARGET and target in df else df
	missing = x.isna().sum().sum() / max(x.size, 1)
	regime = "HDLSS (m > n)" if m > n else "n ≥ m"
	return f"""<div id='studio-metrics'><div class='metric-row'>
	<div class='metric'><b>{n:,}</b><span>Rows (n)</span></div>
	<div class='metric'><b>{m:,}</b><span>Features (m)</span></div>
	<div class='metric'><b>{rho:.2f}</b><span>Ratio ρ = m/n</span></div>
	<div class='metric'><b>{M:,}</b><span>Auto-selected M</span></div>
	</div><p><b>Data scale κ = nm:</b> {kappa:,} · <b>Missing feature cells:</b> {missing:.1%}
	· <b>Regime:</b> {regime}</p></div>"""


	def detect_target(df: pd.DataFrame) -> str:
	"""Conservatively auto-select a named classification target.

	Do not infer a target solely because a column happens to contain strings:
	a categorical feature could otherwise be mistaken for the label.
	"""
	preferred = ("label", "target", "class", "class_label", "classlabel", "outcome", "y")
	for name in preferred:
	for col in df.columns:
	key = str(col).strip().lower().replace(" ", "_").replace("-", "_")
	if key != name:
	continue
	values = df[col]
	counts = values.value_counts(dropna=False)
	if 2 <= len(counts) <= 15 and not values.isna().any() and counts.min() >= 2:
	return str(col)
	return NO_TARGET


	def inspect_csv(path):
	if not path:
	return (gr.update(choices=[NO_TARGET], value=NO_TARGET),
	gr.update(choices=[], value=None), "Upload a dataset to get started.", pd.DataFrame())
	df = read_csv(path)
	numeric = numeric_columns(df)
	inferred_target = detect_target(df)
	choices = [NO_TARGET] + [str(c) for c in df.columns]
	feature_choices = [str(c) for c in numeric if str(c) != inferred_target][:MAX_PLOT_CHOICES]
	summary = shape_card(df, inferred_target)
	nonnumeric_features = [str(c) for c in df.columns
	if str(c) != inferred_target and not pd.api.types.is_numeric_dtype(df[c])]
	if inferred_target != NO_TARGET:
	target_msg = ("<strong>Auto-selected class target:</strong> <code>" +
	html.escape(inferred_target) +
	"</code>. Change it in the dropdown if needed. ")
	else:
	target_msg = ("<strong>Unconditional generation selected.</strong> "
	"Choose a class target from the dropdown if you have one. ")
	if nonnumeric_features:
	target_msg += ("<strong>Other nonnumeric feature columns:</strong> " +
	html.escape(", ".join(nonnumeric_features[:6])) +
	". Encode or remove these, or select the correct target column.")
	else:
	target_msg += "All remaining feature columns are numeric."
	summary += "<p>" + target_msg + "</p>"
	return (gr.update(choices=choices, value=inferred_target),
	gr.update(choices=feature_choices,
	value=feature_choices[0] if feature_choices else None),
	summary, df.iloc[:8, :MAX_PREVIEW_COLUMNS])


	def update_target_summary(path, target):
	"""Keep feature counts accurate when a visitor changes target selection."""
	if not path:
	return "Upload a dataset to get started."
	df = read_csv(path)
	if target != NO_TARGET and target not in df.columns:
	return shape_card(df)
	nonnumeric = [str(col) for col in df.columns
	if str(col) != target and not pd.api.types.is_numeric_dtype(df[col])]
	msg = ("<strong>Class target:</strong> " + html.escape(target) + ". ")
	if nonnumeric:
	msg += "Encode or remove nonnumeric features: " + html.escape(", ".join(nonnumeric[:6]))
	else:
	msg += "All remaining feature columns are numeric."
	return shape_card(df, target) + "<p>" + msg + "</p>"


	def select_default_m(n, m):
	# Dataset-aware no-Optuna initialization from the BSTabDiff project README.
	if m <= 64:
	return m
	rho, kappa = m/n, n*m
	if rho < 10:
	return min(64, m)
	return min(32 if kappa < 250_000 else
	64 if kappa < 1_000_000 else
	128 if kappa < 3_000_000 else 192, m)


	def validate_features(df, target):
	if target != NO_TARGET and target not in df.columns:
	raise gr.Error("Select an existing target column or None.")
	feature_df = df.drop(columns=[target]) if target != NO_TARGET else df.copy()
	if feature_df.shape[1] < 2:
	raise gr.Error("At least two numerical feature columns are required.")
	invalid = [str(c) for c in feature_df.columns if not pd.api.types.is_numeric_dtype(feature_df[c])]
	if invalid:
	names = ", ".join(invalid[:4])
	raise gr.Error(f"Feature columns must be numerical for this first demo. Encode or remove: {names}.")
	if feature_df.shape[1] > MAX_FEATURES:
	raise gr.Error(f"Please use at most {MAX_FEATURES:,} feature columns.")
	if len(df) * feature_df.shape[1] > MAX_CELLS:
	raise gr.Error(f"Shared-hardware limit: n × m must be at most {MAX_CELLS:,}. Try a smaller CSV.")
	X = feature_df.to_numpy(dtype=np.float32)
	if np.isinf(X).any():
	raise gr.Error("Infinite values are not supported; replace ±inf with missing values first.")
	if not np.isfinite(X).any():
	raise gr.Error("The CSV contains no finite numerical feature values.")
	all_missing = np.all(np.isnan(X), axis=0)
	if all_missing.any():
	raise gr.Error(f"Remove entirely missing feature column: {feature_df.columns[np.flatnonzero(all_missing)[0]]}")
	y = None
	if target != NO_TARGET:
	y = df[target]
	if y.isna().any():
	raise gr.Error("The target column must not contain missing class labels.")
	if y.nunique() < 2 or y.nunique() > 15:
	raise gr.Error("Select a classification target with 2–15 classes, or choose None.")
	if y.value_counts().min() < 2:
	raise gr.Error("Every class must have at least two rows for this demo.")
	return feature_df, X, y


	def chart_layout(fig, title):
	fig.update_layout(template="plotly_white", title=dict(text=title, x=.02),
	paper_bgcolor="rgba(0,0,0,0)", plot_bgcolor="rgba(0,0,0,0)",
	margin=dict(l=28,r=18,t=55,b=35), font=dict(size=12),
	legend=dict(orientation="h", y=-.2), height=390)
	return fig


	def distribution_plot(real, synth, column):
	if not column or column not in real.columns:
	return go.Figure()
	fig = go.Figure()
	for name, frame, color in [("Real", real, "#536ed7"), ("Synthetic", synth, "#f06f9c")]:
	values = pd.to_numeric(frame[column], errors="coerce").replace([np.inf,-np.inf],np.nan).dropna()
	fig.add_trace(go.Histogram(x=values, name=name, marker_color=color,
	opacity=.58, nbinsx=30, histnorm="probability density"))
	fig.update_layout(barmode="overlay", xaxis_title=str(column), yaxis_title="Density")
	return chart_layout(fig, f"Feature distribution · {column}")


	def pca_plot(real, synth):
	# A joint coordinate system fitted on real training data only.
	# Limit feature count to avoid a prohibitively large SVD for 50K features.
	cols = real.var(skipna=True).fillna(0).nlargest(min(200, real.shape[1])).index
	nr, ns = min(500, len(real)), min(500, len(synth))
	R = real.loc[:, cols].head(nr).to_numpy(dtype=float)
	S = synth.loc[:, cols].head(ns).to_numpy(dtype=float)
	imp = SimpleImputer(strategy="median", keep_empty_features=True)
	R1 = imp.fit_transform(R)
	S1 = imp.transform(S)
	scale = StandardScaler()
	R2 = scale.fit_transform(R1)
	S2 = scale.transform(S1)
	pca = PCA(n_components=2, svd_solver="randomized", random_state=42)
	rp = pca.fit_transform(R2)
	sp = pca.transform(S2)
	fig = go.Figure()
	for name, pts, color in [("Real",rp,"#536ed7"),("Synthetic",sp,"#f06f9c")]:
	fig.add_trace(go.Scatter(x=pts[:,0],y=pts[:,1],mode="markers",name=name,
	marker=dict(size=7,color=color,opacity=.63)))
	fig.update_layout(xaxis_title="PC 1",yaxis_title="PC 2")
	return chart_layout(fig,"Shared PCA projection · training-data diagnostic")


	def correlation_plot(real, synth):
	variance = real.var(skipna=True).fillna(0).sort_values(ascending=False)
	cols = list(variance.head(min(18,len(variance))).index)
	if len(cols) < 2:
	return go.Figure()
	rc = real[cols].corr().fillna(0)
	sc = synth[cols].corr().fillna(0)
	fig = make_subplots(rows=1,cols=2, subplot_titles=("Real", "Synthetic"),horizontal_spacing=.12)
	for idx, mat in enumerate((rc,sc),start=1):
	fig.add_trace(go.Heatmap(z=mat.values,x=cols,y=cols,zmin=-1,zmax=1,
	colorscale="RdBu", reversescale=True,
	showscale=idx==2,colorbar=dict(title="r") if idx==2 else None),
	row=1,col=idx)
	fig.update_layout(height=490)
	fig.update_xaxes(tickangle=-45, tickfont=dict(size=9))
	fig.update_yaxes(tickfont=dict(size=9))
	chart_layout(fig,"Feature correlations · highest-variance columns")
	fig.update_layout(height=490,margin=dict(l=28,r=25,t=65,b=65))
	return fig


	@spaces.GPU(duration=180)
	def generate(path, target, n_synth, seed, feature):
	df = read_csv(path)
	features, X, y = validate_features(df,target)
	n_synth = int(n_synth)
	seed = int(seed)
	if n_synth < 10 or n_synth > MAX_SYNTHETIC:
	raise gr.Error(f"Choose between 10 and {MAX_SYNTHETIC} synthetic rows.")
	if not 0 <= seed <= 2**31-1:
	raise gr.Error("Use a nonnegative random seed smaller than 2,147,483,648.")
	n, m = X.shape
	if n_synth * m > MAX_OUTPUT_CELLS:
	allowed = MAX_OUTPUT_CELLS // m
	raise gr.Error(
	f"Output would contain {n_synth * m:,} cells. "
	f"With {m:,} features, generate at most {allowed:,} rows "
	f"(maximum {MAX_OUTPUT_CELLS:,} output cells)."
	)
	M = select_default_m(n, m)
	# Map arbitrary user class names (e.g. control/treated) to integer IDs
	# for the generator, then restore the original labels in the exported CSV.
	encoder = None
	y_fit = None
	if y is not None:
	encoder = LabelEncoder()
	y_fit = encoder.fit_transform(y.to_numpy())
	started = time.perf_counter()
	torch.set_num_threads(min(2,os.cpu_count() or 2))
	try:
	# Fixed paper-default diffusion settings, no Optuna, no ordering.
	model = BSTabDiff(
	M=M,
	prior_type="diffusion",
	device="cuda" if torch.cuda.is_available() else "cpu",
	random_state=seed,
	prior_epochs=PAPER_EPOCHS,
	prior_batch=PRIOR_BATCH,
	prior_lr=PRIOR_LR,
	use_ema=True,
	ema_decay=EMA_DECAY,
	verbose_every=0,
	save_dir=None,
	use_gobs=False,
	use_gobs_fc=False,
	)
	model.fit(features, y_fit)
	X_syn,R_syn,y_syn = model.sample(n_samples=n_synth,return_mask=True)
	except Exception as exc:
	raise gr.Error(f"BSTabDiff training or sampling failed: {type(exc).__name__}: {exc}") from exc
	X_syn = np.asarray(X_syn,dtype=np.float32)
	if X_syn.shape != (n_synth,m):
	raise gr.Error(f"Unexpected generated shape {X_syn.shape}; expected {(n_synth,m)}.")
	if R_syn is not None:
	mask = np.asarray(R_syn)
	if mask.shape == X_syn.shape:
	X_syn = X_syn.copy()
	X_syn[mask == 0] = np.nan
	syn = pd.DataFrame(X_syn,columns=features.columns)
	if y is not None:
	if y_syn is None:
	raise gr.Error("The generator did not return labels for the selected class target.")
	sampled_codes = np.asarray(y_syn).reshape(-1)
	if len(sampled_codes) != n_synth:
	raise gr.Error("Generated class label count does not match synthetic rows.")
	try:
	sampled_codes = sampled_codes.astype(int)
	except (TypeError, ValueError, OverflowError) as exc:
	raise gr.Error("Generated class labels could not be decoded.") from exc
	if np.any(sampled_codes < 0) or np.any(sampled_codes >= len(encoder.classes_)):
	raise gr.Error("Generator returned an unknown class ID.")
	syn[target] = encoder.inverse_transform(sampled_codes)
	real = features.reset_index(drop=True)
	syn_features = syn[features.columns]
	plot_columns = list(features.columns[:MAX_PLOT_CHOICES])
	column = feature if feature in plot_columns else plot_columns[0]
	path_csv = os.path.join(tempfile.mkdtemp(prefix="bstabdiff_studio_"),"bstabdiff_synthetic.csv")
	syn.to_csv(path_csv,index=False)
	seconds = time.perf_counter()-started
	summary = (f"### ✅ Generated {n_synth:,} synthetic samples\n"
	f"Input: {n:,} rows × {m:,} features · "
	f"ρ=m/n: {m/n:.3f} · κ=nm: {n*m:,} · "
	f"Auto-selected M: {M}\n\n"
	f"Diffusion epochs: {PAPER_EPOCHS:,} · "
	f"Batch: {PRIOR_BATCH} · LR: {PRIOR_LR:g} · "
	f"EMA decay: {EMA_DECAY} · Seed: {seed} · "
	f"Device: {model.device_} · Elapsed: {seconds:.1f}s\n\n"
	"*Paper-default configuration with dataset-aware M, no Optuna or GO-BS. "
	"Plots compare against uploaded training data; this is not "
	"a held-out utility or privacy evaluation.*")
	# The input frame may contain a class target; output a privacy-minded 10-row preview only.
	state=(real,syn_features)
	return (summary,syn.head(10),distribution_plot(*state,column),
	pca_plot(state),correlation_plot(state),path_csv,state,
	gr.update(choices=plot_columns,value=column))


	def update_plot(feature, state):
	if not state or not feature:
	return go.Figure()
	real, synthetic = state
	return distribution_plot(real,synthetic,feature)


	with gr.Blocks(title="BSTabDiff Studio — NeurIPS 2026",
	delete_cache=(3600,3600)) as demo:
	gr.HTML("""
	<div id="studio-hero">
	<div class="eyebrow" style="color:#fff;font-weight:700;letter-spacing:.13em;font-size:.74rem;opacity:.93">NEURIPS 2026 · RESEARCH DEMO</div>
	<h1>✦ BSTabDiff Studio</h1>
	<p>From your CSV to high-dimensional synthetic data.</p>
	<div class="sub" style="color:#fff">Upload → Generate → Visualize → Download · Default diffusion, no hyperparameter tuning.</div>
	</div>
	""")
	gr.Markdown("How it works: Supply a CSV containing numeric features and optionally a discrete class label. "
	"This Space trains BSTabDiff on your uploaded table and generates new synthetic rows. "
	"Please do not upload confidential, patient-level, or personally identifying data.")

	with gr.Row(equal_height=False):
	with gr.Column(scale=5,elem_classes=["studio-section"]):
	gr.Markdown("### ① Upload & configure")
	file_input=gr.File(label="Upload CSV",file_types=[".csv"],type="filepath")
	gr.Markdown(f"*Limits: {MAX_FILE_MB} MB, {MAX_ROWS:,} rows, "
	f"{MAX_FEATURES:,} features, {MAX_CELLS:,} input cells; "
	f"{MAX_OUTPUT_CELLS:,} synthetic output cells. "
	f"Feature explorer lists the first {MAX_PLOT_CHOICES} numeric columns.*")
	target=gr.Dropdown(label="Optional class target",choices=[NO_TARGET],value=NO_TARGET,
	info="Common column names like label/target/class are automatically selected. Text labels are supported; remaining features must be numeric.")
	with gr.Row():
	samples=gr.Slider(label="Synthetic rows",minimum=10,maximum=MAX_SYNTHETIC,
	step=10,value=100)
	seed_input=gr.Number(label="Random seed",value=42,precision=0,minimum=0,maximum=9999)
	gr.Markdown(f"Paper-default generator: {PAPER_EPOCHS:,} diffusion epochs, "
	f"batch {PRIOR_BATCH}, LR {PRIOR_LR:g}, EMA {EMA_DECAY}; "
	"M is chosen automatically from n, m, ρ and κ. "
	"No GO-BS or Optuna. Large datasets may exceed ZeroGPU runtime limits.")
	run_btn=gr.Button("✦ Generate synthetic dataset",variant="primary",elem_id="generate-btn",size="lg")
	with gr.Column(scale=6,elem_classes=["studio-section"]):
	gr.Markdown("### ② Dataset snapshot")
	file_info=gr.HTML("<p>Upload a CSV to inspect its structure.</p>")
	preview=gr.Dataframe(label=f"Original data · first 8 rows / {MAX_PREVIEW_COLUMNS} columns",interactive=False,
	max_height=310)

	gr.Markdown("## ③ Explore your generated data")
	result=gr.Markdown("The synthetic-data dashboard will appear here after generation.")
	state=gr.State(value=None)
	with gr.Row():
	real_plot_feature=gr.Dropdown(label="Feature to inspect",choices=[],interactive=True)
	download=gr.File(label="Download synthetic CSV",interactive=False)
	with gr.Row():
	with gr.Column():
	with gr.Tab("📊 Feature distribution"):
	hist=gr.Plot(label="Real vs synthetic feature distribution")
	with gr.Tab("🧭 PCA projection"):
	pca=gr.Plot(label="Shared PCA projection")
	with gr.Column():
	with gr.Tab("🧩 Correlations"):
	corr=gr.Plot(label="Real and synthetic correlation matrices")
	with gr.Tab("🧬 Synthetic preview"):
	generated_table=gr.Dataframe(label="Synthetic data · first 10 rows",interactive=False)

	gr.HTML("""<div style="text-align:center;padding:16px;font-size:12px;opacity:.75">
	<a href="https://www.zadidhabib.com/bstabdiff.html" target="_blank">Project webpage</a> ·
	<a href="https://github.com/zadid6pretam/BSTabDiff" target="_blank">GitHub</a> ·
	<a href="https://pypi.org/project/bstabdiff/" target="_blank">PyPI</a> ·
	<a href="https://neurips.cc/virtual/2026/loc/atlanta/poster/151140" target="_blank" rel="noopener noreferrer">NeurIPS Poster</a> ·
	<a href="https://openreview.net/forum?id=ciPCMXX1aj" target="_blank">NeurIPS 2026 Paper</a>
	<p>Research demonstration only. Synthetic data are not necessarily private or faithful to the source distribution.</p>
	</div>""")

	file_input.change(inspect_csv,inputs=file_input,
	outputs=[target,real_plot_feature,file_info,preview],show_progress="minimal")
	target.change(update_target_summary, inputs=[file_input, target], outputs=file_info,
	show_progress="hidden")
	run_btn.click(generate,inputs=[file_input,target,samples,seed_input,real_plot_feature],
	outputs=[result,generated_table,hist,pca,corr,download,state,real_plot_feature],
	concurrency_limit=1,show_progress="full")
	real_plot_feature.change(update_plot,inputs=[real_plot_feature,state],outputs=hist,
	show_progress="minimal")

	demo.queue(max_size=6)

	if __name__ == "__main__":
	demo.launch(theme=gr.themes.Soft(primary_hue="indigo",secondary_hue="pink"), css=CSS)