Upload app.py with huggingface_hub

app.py

@@ -1,6 +1,9 @@
 
-import huggingface_hub as hf_hub
 import inspect
+import sys
+from pathlib import Path
+
+import huggingface_hub as hf_hub
 
 # Gradio <-> hub compatibility shim: newer huggingface_hub removed HfFolder.
 if not hasattr(hf_hub, "HfFolder"):
@@ -22,8 +25,20 @@ import pandas as pd
 from sklearn.manifold import TSNE
 from sklearn.decomposition import PCA
 from sklearn.preprocessing import StandardScaler
+from sklearn.metrics import confusion_matrix, f1_score
 import matplotlib.pyplot as plt
 
+# Repo root for local imports
+REPO_ROOT = Path(__file__).resolve().parent.parent
+if str(REPO_ROOT) not in sys.path:
+    sys.path.append(str(REPO_ROOT))
+
+from mixture.train_embedding_router import MoEPredictor  # type: ignore
+
+# ------------------------------------------------------------------------------
+# Data loading (t-SNE + evaluation)
+# ------------------------------------------------------------------------------
+
 # Load data
 def load_data():
     print("Loading data...")
@@ -41,9 +56,9 @@ def load_data():
         })
     df = pd.DataFrame(records)
     print(f"Loaded {len(df)} samples.")
-    return df
+    return df, data
 
-df = load_data()
+df, raw_samples = load_data()
 
 # Get unique values for filters
 tech_choices = sorted(list(df['tech'].unique()))
@@ -160,40 +175,137 @@ def plot_tsne(tech_filter, snr_filter, mod_filter, mob_filter, representation, c
     trace_info = f"traces: {len(filtered_df[color_by].unique())}"
     return fig, f"{status_msg} | filtered samples: {len(filtered_df)} | {coord_info} | {trace_info}"
 
+# ------------------------------------------------------------------------------
+# Evaluation utilities (confusion matrix, F1) using the MoE checkpoint
+# ------------------------------------------------------------------------------
+
+_predictor: MoEPredictor | None = None
+
+
+def load_predictor() -> MoEPredictor:
+    global _predictor
+    if _predictor is not None:
+        return _predictor
+
+    # Prefer local checkpoint if present; otherwise pull from Hub
+    candidates = [
+        REPO_ROOT / "mixture" / "runs" / "embedding_router" / "moe_checkpoint.pth",
+        REPO_ROOT / "moe_checkpoint.pth",
+    ]
+    ckpt_path = None
+    for cand in candidates:
+        if cand.exists():
+            ckpt_path = cand
+            break
+    if ckpt_path is None:
+        ckpt_path = Path(
+            hf_hub.hf_hub_download(repo_id="wi-lab/lwm-spectro", filename="moe_checkpoint.pth")
+        )
+
+    _predictor = MoEPredictor.from_checkpoint(ckpt_path)
+    return _predictor
+
+
+def _to_tensor(spec) -> torch.Tensor:
+    t = spec
+    if not isinstance(t, torch.Tensor):
+        t = torch.as_tensor(t)
+    if t.dim() == 2:
+        t = t.unsqueeze(0)
+    return t
+
+
+def compute_eval(task: str):
+    """Compute confusion matrix + macro F1 for the small demo set."""
+    predictor = load_predictor()
+    y_true, y_pred = [], []
+
+    for sample in raw_samples:
+        spec = _to_tensor(sample["data"])
+        res = predictor.predict(spec, return_routing=True)
+
+        if task == "comm":
+            routing = res.get("routing") or []
+            pred = routing[0]["comm"] if routing else "Unknown"
+            true = sample["tech"]
+        else:  # snr_mobility
+            pred = res.get("label", res["predicted_class"])
+            true = (sample["snr"], sample["mob"])
+        y_true.append(true)
+        y_pred.append(pred)
+
+    labels = sorted(list({*y_true, *y_pred}))
+    cm = confusion_matrix(y_true, y_pred, labels=labels)
+    f1 = f1_score(y_true, y_pred, labels=labels, average="macro", zero_division=0)
+    acc = (np.array(y_true) == np.array(y_pred)).mean()
+    return cm, labels, f1, acc
+
+
+def plot_confusion(cm: np.ndarray, labels):
+    fig, ax = plt.subplots(figsize=(6, 5))
+    im = ax.imshow(cm, cmap="Blues")
+    ax.figure.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+    ax.set_xticks(np.arange(len(labels)), labels=labels, rotation=45, ha="right")
+    ax.set_yticks(np.arange(len(labels)), labels=labels)
+    ax.set_xlabel("Predicted")
+    ax.set_ylabel("True")
+    for i in range(cm.shape[0]):
+        for j in range(cm.shape[1]):
+            ax.text(j, i, int(cm[i, j]), ha="center", va="center", color="black")
+    fig.tight_layout()
+    return fig
+
+
+def run_eval(task):
+    cm, labels, f1, acc = compute_eval(task)
+    fig = plot_confusion(cm, labels)
+    summary = f"Task: {task} | Accuracy: {acc:.4f} | Macro F1: {f1:.4f}"
+    return fig, summary
+
+
 # UI
 with gr.Blocks(title="LWM-Spectro Demo") as demo:
-    gr.Markdown("# 🔬 LWM-Spectro Interactive t-SNE Demo")
-    gr.Markdown("""
-    Visualise t-SNE just like the local `task1/plot_tsne.py` script:
-    standardised inputs, per-sample normalisation, and SNR/mod/tech/mob colour options.
-    """)
-    
-    with gr.Row():
-        with gr.Column(scale=1, min_width=300):
-            gr.Markdown("### Filters")
-            tech_filter = gr.CheckboxGroup(choices=tech_choices, value=tech_choices[:1], label="Technology (default: single tech)")
-            snr_filter = gr.Dropdown(choices=snr_choices, value=None, multiselect=True, label="SNR (Empty = All)")
-            mod_filter = gr.Dropdown(choices=mod_choices, value=None, multiselect=True, label="Modulation (Empty = All)")
-            mob_filter = gr.Dropdown(choices=mob_choices, value=None, multiselect=True, label="Mobility (Empty = All)")
-            
-            gr.Markdown("### Visualization Settings")
-            representation = gr.Radio(choices=["LWM Embedding", "Raw Spectrogram"], value="LWM Embedding", label="Representation")
-            color_by = gr.Dropdown(choices=["tech", "snr", "mod", "mob"], value="snr", label="Color By")
-            
-            with gr.Accordion("Advanced t-SNE Settings", open=False):
-                perplexity = gr.Slider(minimum=5, maximum=50, value=10, step=1, label="Perplexity")
-                n_iter = gr.Slider(minimum=250, maximum=2000, value=1000, step=50, label="Iterations")
-                
-            btn = gr.Button("Update Plot", variant="primary")
-            status = gr.Textbox(label="Status", interactive=False)
-            
-        with gr.Column(scale=3):
-            plot = gr.Plot(label="t-SNE Visualization")
-            
-    btn.click(plot_tsne, inputs=[tech_filter, snr_filter, mod_filter, mob_filter, representation, color_by, perplexity, n_iter], outputs=[plot, status])
-    
-    # Initial load
-    demo.load(plot_tsne, inputs=[tech_filter, snr_filter, mod_filter, mob_filter, representation, color_by, perplexity, n_iter], outputs=[plot, status])
+    gr.Markdown("# 🔬 LWM-Spectro Interactive Demo")
+    gr.Markdown("Compare embeddings vs raw for t-SNE, and view quick metrics from the latest MoE checkpoint.")
+
+    with gr.Tab("t-SNE"):
+        with gr.Row():
+            with gr.Column(scale=1, min_width=300):
+                gr.Markdown("### Filters")
+                tech_filter = gr.CheckboxGroup(choices=tech_choices, value=tech_choices[:1], label="Technology (default: single tech)")
+                snr_filter = gr.Dropdown(choices=snr_choices, value=None, multiselect=True, label="SNR (Empty = All)")
+                mod_filter = gr.Dropdown(choices=mod_choices, value=None, multiselect=True, label="Modulation (Empty = All)")
+                mob_filter = gr.Dropdown(choices=mob_choices, value=None, multiselect=True, label="Mobility (Empty = All)")
+
+                gr.Markdown("### Visualization Settings")
+                representation = gr.Radio(choices=["LWM Embedding", "Raw Spectrogram"], value="LWM Embedding", label="Representation")
+                color_by = gr.Dropdown(choices=["tech", "snr", "mod", "mob"], value="snr", label="Color By")
+
+                with gr.Accordion("Advanced t-SNE Settings", open=False):
+                    perplexity = gr.Slider(minimum=5, maximum=50, value=10, step=1, label="Perplexity")
+                    n_iter = gr.Slider(minimum=250, maximum=2000, value=1000, step=50, label="Iterations")
+
+                btn = gr.Button("Update Plot", variant="primary")
+                status = gr.Textbox(label="Status", interactive=False)
+
+            with gr.Column(scale=3):
+                plot = gr.Plot(label="t-SNE Visualization")
+
+        btn.click(plot_tsne, inputs=[tech_filter, snr_filter, mod_filter, mob_filter, representation, color_by, perplexity, n_iter], outputs=[plot, status])
+
+        # Initial load
+        demo.load(plot_tsne, inputs=[tech_filter, snr_filter, mod_filter, mob_filter, representation, color_by, perplexity, n_iter], outputs=[plot, status])
+
+    with gr.Tab("Evaluation (MoE)"):
+        gr.Markdown("Uses the latest MoE checkpoint to score the bundled demo set. Communication uses router gating; SNR/Mobility uses the classifier head.")
+        task_choice = gr.Radio(choices=["comm", "snr_mobility"], value="snr_mobility", label="Task")
+        eval_btn = gr.Button("Run Evaluation", variant="primary")
+        cm_plot = gr.Plot(label="Confusion Matrix")
+        eval_summary = gr.Textbox(label="Metrics", interactive=False)
+
+        eval_btn.click(run_eval, inputs=[task_choice], outputs=[cm_plot, eval_summary])
+        # Run once on load for convenience
+        demo.load(run_eval, inputs=[task_choice], outputs=[cm_plot, eval_summary])
 
 if __name__ == "__main__":
     demo.launch()