Update app.py

app.py

@@ -51,7 +51,6 @@ def load_all():
     knn = KNeighborsClassifier(n_neighbors=K, metric='cosine', n_jobs=-1)
     knn.fit(ref_emb.values, ref_clin['uberon_tissue'].values)
 
-    # Try loading full VAE model (optional — needed for new sample encoding)
     model = None
     pth = MODEL_DIR / "vae_tissue.final_model.pth"
     if pth.exists():
@@ -107,21 +106,23 @@ def classify(df):
             ni     = n2i.get('nan', None)
             if ni is not None:
                 logits[:, ni] = -1e9
-            probs      = torch.softmax(logits, dim=1)
-            pred_idx   = logits.argmax(dim=1).numpy()
+            probs    = torch.softmax(logits, dim=1)
+            pred_idx = logits.argmax(dim=1).numpy()
         emb         = mu.numpy()
         pred_labels = [lmap[int(i)] for i in pred_idx]
         max_probs   = probs.numpy().max(axis=1)
-   else:
+    else:
         from sklearn.decomposition import TruncatedSVD
-        n_comp = min(121, X_scaled.shape[0] - 1, X_scaled.shape[1])
-        svd = TruncatedSVD(n_components=n_comp, random_state=42)
+        n_comp      = min(121, X_scaled.shape[0] - 1, X_scaled.shape[1])
+        svd         = TruncatedSVD(n_components=n_comp, random_state=42)
         emb_reduced = svd.fit_transform(X_scaled)
-        # Pad to exactly 121 dims
-        pad = np.zeros((emb_reduced.shape[0], 121 - emb_reduced.shape[1]))
-        emb = np.concatenate([emb_reduced, pad], axis=1)
+        pad         = np.zeros((emb_reduced.shape[0], 121 - emb_reduced.shape[1]))
+        emb         = np.concatenate([emb_reduced, pad], axis=1)
         pred_labels = knn_ref.predict(emb).tolist()
         max_probs   = np.ones(len(pred_labels)) * float('nan')
+
+    return emb, pred_labels, max_probs
+
 # ── Upload ──
 st.markdown("---")
 col1, col2 = st.columns([2, 1])
@@ -153,7 +154,6 @@ if uploaded:
                 st.warning("Low gene overlap — results may be unreliable.")
 
             emb, pred_labels, max_probs = classify(df)
-
             distances, indices = knn_ref.kneighbors(emb)
             breakdowns = []
             for i in range(len(emb)):
@@ -197,4 +197,4 @@ st.caption(
     "Flexynesis Tissue VAE · Akalin Lab, MDC Berlin/BIMSB · "
     "75,619 samples · 43 UBERON tissues · 90.7% balanced accuracy · "
     "github.com/BIMSBbioinfo/flexynesis"
-)
+)