Update app.py

app.py

@@ -1,18 +1,34 @@
 import gradio as gr
 import numpy as np
 import pandas as pd
+import torch
 from huggingface_hub import hf_hub_download
 from transformers import AutoConfig, AutoModel, AutoTokenizer
-import torch
 
 print("Loading BulkRNABert...")
-config = AutoConfig.from_pretrained("InstaDeepAI/BulkRNABert", trust_remote_code=True)
+
+# --- Load model and config (official API) ---
+config = AutoConfig.from_pretrained(
+    "InstaDeepAI/BulkRNABert",
+    trust_remote_code=True,
+)
 config.embeddings_layers_to_save = (4,)
-tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/BulkRNABert", trust_remote_code=True)
-model = AutoModel.from_pretrained("InstaDeepAI/BulkRNABert", config=config, trust_remote_code=True)
-print("✓ Model loaded")
 
-# Cache the required gene list once at startup
+tokenizer = AutoTokenizer.from_pretrained(
+    "InstaDeepAI/BulkRNABert",
+    trust_remote_code=True,
+)
+
+model = AutoModel.from_pretrained(
+    "InstaDeepAI/BulkRNABert",
+    config=config,
+    trust_remote_code=True,
+)
+model.eval()
+
+print(f"✓ Model loaded. Required genes: {config.n_genes}")
+
+# --- Download reference gene list once at startup ---
 csv_path = hf_hub_download(
     repo_id="InstaDeepAI/BulkRNABert",
     filename="data/tcga_sample.csv",
@@ -20,27 +36,67 @@ csv_path = hf_hub_download(
 )
 ref_df = pd.read_csv(csv_path)
 ref_genes = [c for c in ref_df.columns if c != "identifier"]
+print(f"✓ Reference genes: {len(ref_genes)}")
 
-def get_embeddings(file):
-    df = pd.read_csv(file.name, sep="\t", index_col=0).T
 
-    aligned = pd.DataFrame(0, index=df.index, columns=ref_genes)
-    for g in ref_genes:
+def align_genes(df, required_genes):
+    """Align dataframe to the required gene set, fill missing with 0."""
+    aligned = pd.DataFrame(0, index=df.index, columns=required_genes)
+    for g in required_genes:
         if g in df.columns:
             aligned[g] = df[g]
+    return aligned
+
+
+def get_embeddings(df, model, tokenizer, config):
+    """Generate BulkRNABert embeddings following the official pipeline."""
+    expr = df.to_numpy()
+    expr = np.log10(1 + expr)  # Official transform
+
+    assert expr.shape[1] == config.n_genes, \
+        f"Expected {config.n_genes} genes, got {expr.shape[1]}"
 
-    expr = np.log10(1 + aligned.to_numpy())
     ids = tokenizer.batch_encode_plus(expr, return_tensors="pt")["input_ids"]
 
     with torch.no_grad():
-        emb = model(ids)["embeddings_4"].mean(axis=1).numpy()
+        out = model(ids)
+
+    # Official: extract layer-4 embeddings, average across genes
+    return out["embeddings_4"].mean(axis=1).numpy()
+
+
+def process_file(file):
+    """Handle uploaded TSV file and return embeddings."""
+    if file is None:
+        return {"error": "No file uploaded"}
 
-    return {"shape": list(emb.shape), "first_5_values": emb[:5].tolist()}
+    try:
+        # Read uploaded file
+        df = pd.read_csv(file.name, sep="\t", index_col=0).T
 
+        # Align to required genes
+        aligned = align_genes(df, ref_genes)
+
+        # Generate embeddings
+        emb = get_embeddings(aligned, model, tokenizer, config)
+
+        return {
+            "status": "success",
+            "n_samples": int(emb.shape[0]),
+            "embedding_dim": int(emb.shape[1]),
+            "first_sample_preview": emb[0][:10].tolist(),
+        }
+    except Exception as e:
+        return {"error": str(e)}
+
+
+# --- Build Gradio interface ---
 iface = gr.Interface(
-    fn=get_embeddings,
-    inputs=gr.File(label="Upload your TSV expression file"),
-    outputs=gr.JSON(label="Embeddings"),
+    fn=process_file,
+    inputs=gr.File(label="Upload TSV expression file (genes × samples)"),
+    outputs=gr.JSON(label="Embeddings Result"),
     title="BulkRNABert Embeddings",
+    description="Upload a TSV file with genes as rows and samples as columns.",
 )
+
 iface.launch()