Run real model: load 53MB int8 TorchScript, drop SVD fallback, slim deps

app.py

@@ -1,32 +1,28 @@
 #!/usr/bin/env python3
 """
-Flexynesis Tissue VAE - Web Application (v3, demo-safe)
-======================================================
-Upload a gene expression matrix -> get UBERON tissue classification + embeddings.
+Flexynesis Tissue VAE — Web Application (compressed, real-model inference)
+=========================================================================
+Upload a bulk RNA-seq gene-expression matrix → UBERON tissue-of-origin
+predictions + 121-dim latent embeddings.
 
-v3 (May 2026): Trained on 118,263 tissue-curated samples from TCGA, GTEx, ARCHS4
-(cell lines excluded, classes balanced). 42 UBERON tissues, 94.9% balanced accuracy.
-
-When the full model weights (vae_tissue.final_model.pth, ~8 GB) are absent
-(Hugging Face free tier), the app runs in demo mode: uploaded samples are
-projected with TruncatedSVD and classified by kNN against the pre-computed
-118K reference embeddings. When the .pth IS present (local), the trained VAE
-encoder is used for exact inference.
+Runs the trained supervised VAE via a 53 MB int8 TorchScript bundle
+(vae_tissue_int8.torchscript.pt) — exact inference, no flexynesis dependency,
+fits the Hugging Face free tier. Trained on 118,263 tissue-curated samples
+(TCGA, GTEx, ARCHS4), 42 UBERON tissues, 94.9% balanced accuracy.
 
 Author: Amit Pande, MDC Berlin/BIMSB
 """
+import json
+from pathlib import Path
+from collections import Counter
 
-import streamlit as st
-import pandas as pd
 import numpy as np
+import pandas as pd
+import streamlit as st
+import torch
 import joblib
-from pathlib import Path
-from sklearn.neighbors import KNeighborsClassifier
-from collections import Counter
 
 MODEL_DIR = Path(".")
-K = 5
-LATENT_DIM = 121
 
 st.set_page_config(page_title="Flexynesis Tissue VAE", page_icon="🧬", layout="wide")
 st.title("🧬 Flexynesis Tissue VAE")
@@ -36,103 +32,39 @@ st.markdown(
     "across **42 UBERON tissue categories** (94.9% balanced accuracy, 121-dim latent space)."
 )
 
+
 @st.cache_resource
-def load_all():
-    art        = joblib.load(MODEL_DIR / "vae_tissue.artifacts.joblib")
-    gene_list  = list(art['feature_lists']['gex'])
-    scaler     = art['transforms']['gex']
-
-    train_emb  = pd.read_csv(MODEL_DIR / "embeddings_train.csv", index_col=0)
-    test_emb   = pd.read_csv(MODEL_DIR / "embeddings_test.csv",  index_col=0)
-    train_clin = pd.read_csv(MODEL_DIR / "train_clin.csv", index_col=0)
-    test_clin  = pd.read_csv(MODEL_DIR / "test_clin.csv",  index_col=0)
-
-    all_emb  = pd.concat([train_emb, test_emb])
-    all_clin = pd.concat([train_clin, test_clin])
-    idx      = all_emb.index.intersection(all_clin.index)
-    ref_emb  = all_emb.loc[idx]
-    ref_clin = all_clin.loc[idx]
-    mask = (ref_clin['uberon_tissue'].notna() &
-            ~ref_clin['uberon_tissue'].isin(['unknown', 'other', 'unmapped', 'nan', '']))
-    ref_emb  = ref_emb[mask]
-    ref_clin = ref_clin[mask]
-
-    knn = KNeighborsClassifier(n_neighbors=K, metric='cosine', n_jobs=-1)
-    knn.fit(ref_emb.values, ref_clin['uberon_tissue'].values)
-
-    model = None
-    pth = MODEL_DIR / "vae_tissue.final_model.pth"
-    if pth.exists():
-        try:
-            import torch
-            model = torch.load(pth, map_location='cpu', weights_only=False)
-            model.eval()
-        except Exception as ex:
-            st.sidebar.warning(f"Weights present but not loaded ({ex}); using kNN demo mode.")
-            model = None
-
-    return model, gene_list, scaler, knn, ref_emb, ref_clin
+def load_model():
+    model = torch.jit.load(str(MODEL_DIR / "vae_tissue_int8.torchscript.pt"))
+    model.eval()
+    art = joblib.load(MODEL_DIR / "vae_tissue.artifacts.joblib")
+    gene_list = list(art["feature_lists"]["gex"])
+    scaler = art["transforms"]["gex"]
+    label_mapping = {int(k): v for k, v in
+                     json.loads((MODEL_DIR / "label_mapping.json").read_text()).items()}
+    return model, gene_list, scaler, label_mapping
+
 
 try:
-    model, gene_list, scaler, knn_ref, ref_emb, ref_clin = load_all()
-    mode_str = "Full VAE encoder + kNN" if model is not None else "kNN on pre-computed embeddings"
+    model, gene_list, scaler, label_mapping = load_model()
     st.success(
-        f"✅ {len(gene_list):,} genes · "
-        f"{len(ref_emb):,} reference samples · "
-        f"{ref_clin['uberon_tissue'].nunique()} tissues · "
-        f"Mode: **{mode_str}**"
+        f"✅ Model loaded · {len(gene_list):,} genes · "
+        f"{len(label_mapping)} tissue classes · exact VAE inference (int8 TorchScript)"
     )
-    if model is None:
-        st.info(
-            "ℹ️ Running in **demo mode**. The 118K reference embeddings are used as-is; "
-            "uploaded samples are projected with TruncatedSVD and classified by kNN. "
-            "Full VAE encoding of new samples requires the model weights "
-            "(vae_tissue.final_model.pth), deposited at Zenodo and too large for the free tier."
-        )
 except Exception as e:
-    st.error(f"Load error: {e}")
-    st.info("Place reference files in repo root (embeddings_*.csv, *_clin.csv, vae_tissue.artifacts.joblib).")
+    st.error(f"Model load error: {e}")
+    st.info("Ensure vae_tissue_int8.torchscript.pt, vae_tissue.artifacts.joblib, "
+            "and label_mapping.json are in the repository root.")
     st.stop()
 
+
 def orient_matrix(df):
-    genes_in_cols = len(set(df.columns) & set(gene_list))
-    genes_in_rows = len(set(df.index)   & set(gene_list))
-    if genes_in_rows > genes_in_cols:
+    in_cols = len(set(df.columns) & set(gene_list))
+    in_rows = len(set(df.index) & set(gene_list))
+    if in_rows > in_cols:
         df = df.T
     return df
 
-def embed_and_classify(X_scaled):
-    if model is not None:
-        import torch
-        X_t = torch.tensor(X_scaled, dtype=torch.float32)
-        with torch.no_grad():
-            h      = model.encoders[0](X_t)
-            mu     = model.FC_mean(h[0])
-            logits = model.MLPs['uberon_tissue'](mu)
-            lmap   = model.dataset.label_mappings['uberon_tissue']
-            n2i    = {v: k for k, v in lmap.items()}
-            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()
-        emb         = mu.numpy()
-        pred_labels = [lmap[int(i)] for i in pred_idx]
-        max_probs   = probs.numpy().max(axis=1)
-    else:
-        from sklearn.decomposition import TruncatedSVD
-        n_comp      = min(LATENT_DIM, X_scaled.shape[0] - 1, X_scaled.shape[1])
-        n_comp      = max(n_comp, 1)
-        svd         = TruncatedSVD(n_components=n_comp, random_state=42)
-        emb_reduced = svd.fit_transform(X_scaled)
-        if emb_reduced.shape[1] < LATENT_DIM:
-            pad = np.zeros((emb_reduced.shape[0], LATENT_DIM - emb_reduced.shape[1]))
-            emb = np.concatenate([emb_reduced, pad], axis=1)
-        else:
-            emb = emb_reduced[:, :LATENT_DIM]
-        pred_labels = knn_ref.predict(emb).tolist()
-        max_probs   = np.full(len(pred_labels), np.nan)
-    return emb, pred_labels, max_probs
 
 st.markdown("---")
 col1, col2 = st.columns([2, 1])
@@ -150,64 +82,50 @@ with col2:
     """)
 
 if uploaded:
-    sep = '\t' if uploaded.name.endswith(('.tsv', '.txt')) else ','
-    df  = pd.read_csv(uploaded, index_col=0, sep=sep)
+    sep = "\t" if uploaded.name.endswith((".tsv", ".txt")) else ","
+    df = pd.read_csv(uploaded, index_col=0, sep=sep)
     st.write(f"**Uploaded:** {df.shape[0]:,} × {df.shape[1]:,}")
     st.dataframe(df.iloc[:5, :5], use_container_width=True)
 
     if st.button("🚀 Classify Tissues", type="primary"):
-        with st.spinner("Processing..."):
-            df      = orient_matrix(df)
+        with st.spinner("Running the VAE..."):
+            df = orient_matrix(df)
             overlap = len(set(df.columns) & set(gene_list))
-            st.write(f"Gene overlap: **{overlap:,}/{len(gene_list):,}** ({100*overlap/len(gene_list):.1f}%)")
+            st.write(f"Gene overlap: **{overlap:,}/{len(gene_list):,}** "
+                     f"({100*overlap/len(gene_list):.1f}%)")
             if overlap < 1000:
                 st.warning("Low gene overlap — results may be unreliable.")
 
-            aligned         = pd.DataFrame(0.0, index=df.index, columns=gene_list)
-            common          = [g for g in gene_list if g in df.columns]
+            aligned = pd.DataFrame(0.0, index=df.index, columns=gene_list)
+            common = [g for g in gene_list if g in df.columns]
             aligned[common] = df[common].values
-            aligned         = aligned.fillna(0)
-            X_scaled        = scaler.transform(aligned.values)
-
-            emb, pred_labels, max_probs = embed_and_classify(X_scaled)
-
-            distances, indices = knn_ref.kneighbors(emb)
-            breakdowns = []
-            for i in range(len(emb)):
-                src = Counter(ref_clin.iloc[indices[i]]['source'].values)
-                breakdowns.append('; '.join(f"{s}:{c}" for s, c in src.most_common()))
+            aligned = aligned.fillna(0)
+            X = torch.tensor(scaler.transform(aligned.values), dtype=torch.float32)
+
+            with torch.no_grad():
+                logits = model(X)
+            nan_idx = {v: k for k, v in label_mapping.items()}.get("nan")
+            if nan_idx is not None:
+                logits[:, nan_idx] = -1e9
+            probs = torch.softmax(logits, dim=1)
+            pred_idx = logits.argmax(dim=1)
+            pred_labels = [label_mapping[int(i)] for i in pred_idx]
+            conf = probs.max(dim=1).values.numpy()
 
         st.markdown("---")
         st.subheader("📊 Results")
-        conf_col = [f"{p:.1%}" if not np.isnan(p) else "kNN" for p in max_probs]
-        results  = pd.DataFrame({
-            'Sample':     df.index,
-            'Tissue':     pred_labels,
-            'Confidence': conf_col,
-            'kNN Dist':   distances.mean(axis=1).round(4),
-            'Sources':    breakdowns,
+        results = pd.DataFrame({
+            "Sample": df.index,
+            "Tissue": pred_labels,
+            "Confidence": [f"{c:.1%}" for c in conf],
         })
         st.dataframe(results, use_container_width=True, height=400)
 
-        col1, col2 = st.columns(2)
-        with col1:
-            st.subheader("Tissue Distribution")
-            st.bar_chart(pd.Series(pred_labels).value_counts())
-        with col2:
-            st.subheader("Confidence Distribution")
-            conf_vals = [p if not np.isnan(p) else 0 for p in max_probs]
-            st.bar_chart(pd.DataFrame({'Confidence': conf_vals}, index=df.index))
+        st.subheader("Tissue Distribution")
+        st.bar_chart(pd.Series(pred_labels).value_counts())
 
-        st.markdown("---")
-        c1, c2 = st.columns(2)
-        with c1:
-            emb_df = pd.DataFrame(emb, index=df.index,
-                                  columns=[f"z{i}" for i in range(emb.shape[1])])
-            st.download_button("📥 Embeddings (CSV)", emb_df.to_csv(),
-                               "flexynesis_embeddings.csv", "text/csv")
-        with c2:
-            st.download_button("📥 Classifications (CSV)", results.to_csv(index=False),
-                               "flexynesis_classifications.csv", "text/csv")
+        st.download_button("📥 Predictions (CSV)", results.to_csv(index=False),
+                           "flexynesis_predictions.csv", "text/csv")
 
 st.markdown("---")
 st.caption(

embeddings_train.csv

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:917d5c6677790a8725edaf78a589fa0c8cc675019548b2add49778f6edaeb2cc
-size 155547606

label_mapping.json

@@ -0,0 +1,44 @@
+{
+  "0": "adipose",
+  "1": "adrenal_gland",
+  "2": "biliary_tract",
+  "3": "bladder",
+  "4": "blood",
+  "5": "blood_vessel",
+  "6": "bone_marrow",
+  "7": "brain",
+  "8": "breast",
+  "9": "cervix",
+  "10": "colon",
+  "11": "esophagus",
+  "12": "eye",
+  "13": "fibroblast",
+  "14": "head_and_neck",
+  "15": "heart",
+  "16": "kidney",
+  "17": "liver",
+  "18": "lung",
+  "19": "lymphoid",
+  "20": "muscle",
+  "21": "nerve",
+  "22": "other",
+  "23": "ovary",
+  "24": "pancreas",
+  "25": "pituitary",
+  "26": "placenta",
+  "27": "pleura",
+  "28": "prostate",
+  "29": "salivary_gland",
+  "30": "skin",
+  "31": "small_intestine",
+  "32": "soft_tissue",
+  "33": "spinal_cord",
+  "34": "spleen",
+  "35": "stem_cell",
+  "36": "stomach",
+  "37": "testis",
+  "38": "thymus",
+  "39": "thyroid",
+  "40": "uterus",
+  "41": "vagina"
+}

requirements.txt

@@ -4,5 +4,3 @@ pandas>=1.5.0
 numpy>=1.24.0
 scikit-learn>=1.3.0
 joblib>=1.3.0
-h5py>=3.10.0
-flexynesis>=0.3.0

embeddings_test.csv → vae_tissue_int8.torchscript.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b8353267526ba6fe1b189650a11eb9918282fa9537515f9893a74195263ef688
-size 37219853
+oid sha256:8b7b0835aad6daab7947f21f1207a8dc05400135cf51f0ddb4b95d2edfe6a4e4
+size 52895629