gradio

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

app.py

@@ -0,0 +1,260 @@
+import os
+from dataclasses import dataclass
+from typing import List, Tuple
+
+import gradio as gr
+import numpy as np
+import plotly.express as px
+import torch
+import umap
+from Bio import SeqIO
+from transformers import AutoModel, AutoTokenizer
+
+from model import MicrobiomeTransformer
+
+
+MAX_GENES = 800
+MAX_SEQ_LEN = 1024
+PROKBERT_MODEL_ID = os.getenv("PROKBERT_MODEL_ID", "neuralbioinfo/prokbert-mini-long")
+CHECKPOINT_PATH = os.getenv("CHECKPOINT_PATH", "large-notext.pt")
+BATCH_SIZE = int(os.getenv("EMBED_BATCH_SIZE", "32"))
+TRUST_REMOTE_CODE = "true"
+
+
+@dataclass
+class LoadedModels:
+    tokenizer: AutoTokenizer
+    prokbert: AutoModel
+    microbiome: MicrobiomeTransformer
+    device: torch.device
+
+
+_MODELS: LoadedModels | None = None
+
+
+def _load_models() -> LoadedModels:
+    global _MODELS
+    if _MODELS is not None:
+        return _MODELS
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    tokenizer = AutoTokenizer.from_pretrained(PROKBERT_MODEL_ID, trust_remote_code=TRUST_REMOTE_CODE)
+    prokbert = AutoModel.from_pretrained(PROKBERT_MODEL_ID, trust_remote_code=TRUST_REMOTE_CODE)
+    prokbert.to(device)
+    prokbert.eval()
+
+    checkpoint = torch.load(CHECKPOINT_PATH, map_location=device)
+    state_dict = checkpoint.get("model_state_dict", checkpoint)
+    microbiome = MicrobiomeTransformer(
+        input_dim_type1=384,
+        input_dim_type2=1536,
+        d_model=100,
+        nhead=5,
+        num_layers=5,
+        dim_feedforward=400,
+        dropout=0.1,
+        use_output_activation=False,
+    )
+    microbiome.load_state_dict(state_dict, strict=False)
+    microbiome.to(device)
+    microbiome.eval()
+
+    _MODELS = LoadedModels(
+        tokenizer=tokenizer,
+        prokbert=prokbert,
+        microbiome=microbiome,
+        device=device,
+    )
+    return _MODELS
+
+
+def _read_fasta(path: str) -> Tuple[List[str], List[str], int, int]:
+    ids: List[str] = []
+    seqs: List[str] = []
+    truncated = 0
+
+    for record in SeqIO.parse(path, "fasta"):
+        seq = str(record.seq).upper()
+        if len(seq) > MAX_SEQ_LEN:
+            seq = seq[:MAX_SEQ_LEN]
+            truncated += 1
+        ids.append(record.id)
+        seqs.append(seq)
+
+    original_n = len(ids)
+    if original_n == 0:
+        raise ValueError("No FASTA records found.")
+
+    if original_n > MAX_GENES:
+        ids = ids[:MAX_GENES]
+        seqs = seqs[:MAX_GENES]
+
+    return ids, seqs, original_n, truncated
+
+
+def _mean_pool(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+    mask = attention_mask.unsqueeze(-1).to(last_hidden_state.dtype)
+    summed = (last_hidden_state * mask).sum(dim=1)
+    counts = mask.sum(dim=1).clamp(min=1e-8)
+    return summed / counts
+
+
+def _embed_sequences(seqs: List[str], models: LoadedModels) -> np.ndarray:
+    pooled_batches: List[np.ndarray] = []
+
+    for i in range(0, len(seqs), BATCH_SIZE):
+        batch = seqs[i : i + BATCH_SIZE]
+        inputs = models.tokenizer(
+            batch,
+            return_tensors="pt",
+            truncation=True,
+            max_length=MAX_SEQ_LEN,
+            padding=True,
+        )
+        inputs = {k: v.to(models.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = models.prokbert(**inputs)
+            pooled = _mean_pool(outputs.last_hidden_state, inputs["attention_mask"])
+
+        pooled_batches.append(pooled.detach().cpu().numpy())
+
+    emb = np.vstack(pooled_batches)
+    if emb.shape[1] != 384:
+        raise ValueError(
+            f"Expected 384-d ProkBERT embeddings, got {emb.shape[1]} dimensions from {PROKBERT_MODEL_ID}."
+        )
+    return emb
+
+
+def _infer_logits_and_final_embeddings(input_embeddings: np.ndarray, models: LoadedModels) -> Tuple[np.ndarray, np.ndarray]:
+    x = torch.tensor(input_embeddings, dtype=torch.float32, device=models.device).unsqueeze(0)
+    n = x.shape[1]
+
+    empty_text = torch.zeros((1, 0, 1536), dtype=torch.float32, device=models.device)
+    mask = torch.ones((1, n), dtype=torch.bool, device=models.device)
+    type_indicators = torch.zeros((1, n), dtype=torch.long, device=models.device)
+
+    batch = {
+        "embeddings_type1": x,
+        "embeddings_type2": empty_text,
+        "mask": mask,
+        "type_indicators": type_indicators,
+    }
+
+    with torch.no_grad():
+        x_proj = models.microbiome.input_projection_type1(batch["embeddings_type1"])
+        final_hidden = models.microbiome.transformer(x_proj, src_key_padding_mask=~mask)
+        logits = models.microbiome.output_projection(final_hidden).squeeze(-1)
+
+    return (
+        logits.squeeze(0).detach().cpu().numpy(),
+        final_hidden.squeeze(0).detach().cpu().numpy(),
+    )
+
+
+def _umap_df(vectors: np.ndarray, labels: List[str], value_name: str):
+    n = vectors.shape[0]
+    if n < 2:
+        raise ValueError("Need at least 2 genes to compute UMAP.")
+
+    reducer = umap.UMAP(
+        n_components=2,
+        n_neighbors=min(15, n - 1),
+        min_dist=0.1,
+        metric="cosine",
+        random_state=42,
+    )
+    coords = reducer.fit_transform(vectors)
+    return {
+        "x": coords[:, 0],
+        "y": coords[:, 1],
+        "gene": labels,
+        value_name: np.linalg.norm(vectors, axis=1),
+    }
+
+
+def _plot_umap(vectors: np.ndarray, labels: List[str], title: str):
+    df = _umap_df(vectors, labels, "norm")
+    fig = px.scatter(
+        df,
+        x="x",
+        y="y",
+        hover_name="gene",
+        color="norm",
+        title=title,
+        color_continuous_scale="Viridis",
+    )
+    fig.update_traces(marker={"size": 9, "line": {"width": 0.5, "color": "black"}})
+    return fig
+
+
+def _plot_logits(logits: np.ndarray, labels: List[str]):
+    fig = px.histogram(
+        x=logits,
+        nbins=min(50, max(10, len(logits) // 4)),
+        title="Logit Distribution Over Input DNA Embeddings",
+    )
+    fig.update_layout(xaxis_title="Logit", yaxis_title="Count")
+    return fig
+
+
+def run_pipeline(fasta_file: str):
+    if fasta_file is None:
+        raise gr.Error("Upload a FASTA file first.")
+
+    models = _load_models()
+    labels, seqs, original_n, truncated = _read_fasta(fasta_file)
+
+    input_embeddings = _embed_sequences(seqs, models)
+    logits, final_embeddings = _infer_logits_and_final_embeddings(input_embeddings, models)
+
+    input_umap = _plot_umap(input_embeddings, labels, "UMAP of Input DNA Embeddings (ProkBERT Mean-Pooled)")
+    final_umap = _plot_umap(final_embeddings, labels, "UMAP of Final Embeddings (After large-notext Transformer)")
+    logits_hist = _plot_logits(logits, labels)
+
+    capped_n = len(labels)
+    info = (
+        f"Loaded {original_n} genes. "
+        f"Used {capped_n} (cap={MAX_GENES}). "
+        f"Truncated {truncated} sequence(s) to {MAX_SEQ_LEN} nt."
+    )
+
+    top_idx = np.argsort(logits)[::-1]
+    top_rows = [[labels[i], float(logits[i])] for i in top_idx[: min(50, len(labels))]]
+
+    return info, input_umap, final_umap, logits_hist, top_rows
+
+
+with gr.Blocks(title="Microbiome Space: ProkBERT -> large-notext") as demo:
+    gr.Markdown(
+        """
+# Microbiome Gene Scoring Explorer
+Upload a FASTA of genes, embed with `prokbert-mini-long` (mean pooling), score with `large-notext`, and inspect embedding geometry + logit distribution.
+
+Constraints:
+- Max genes per run: 800
+- Max gene length: 1024 nt (longer sequences are truncated)
+"""
+    )
+
+    with gr.Row():
+        fasta_in = gr.File(label="FASTA file", file_types=[".fa", ".fasta", ".fna", ".txt"], type="filepath")
+        run_btn = gr.Button("Run", variant="primary")
+
+    status = gr.Textbox(label="Run Summary")
+    input_umap_plot = gr.Plot(label="Input Embedding UMAP")
+    final_umap_plot = gr.Plot(label="Final Embedding UMAP")
+    logits_plot = gr.Plot(label="Logit Distribution")
+    top_table = gr.Dataframe(headers=["gene_id", "logit"], label="Top genes by logit")
+
+    run_btn.click(
+        fn=run_pipeline,
+        inputs=[fasta_in],
+        outputs=[status, input_umap_plot, final_umap_plot, logits_plot, top_table],
+    )
+
+
+if __name__ == "__main__":
+    demo.queue().launch()

requirements.txt

@@ -0,0 +1,8 @@
+gradio>=5.0.0
+torch>=2.1.0
+transformers>=4.44.0
+sentencepiece>=0.2.0
+biopython>=1.84
+umap-learn>=0.5.6
+plotly>=5.24.0
+numpy>=1.26.0