Create model_loader.py

model_loader.py

@@ -0,0 +1,413 @@
+"""
+model_loader.py
+==============
+Loads all three pretrained models using their EXACT native architectures
+as confirmed from the live HuggingFace Space source code.
+
+Models:
+  1. nileshhanotia/mutation-predictor-splice
+     → MutationPredictorCNN_v2  (input dim=1106, 99bp window)
+     → File: mutation_predictor_splice.pt
+
+  2. nileshhanotia/mutation-predictor-v4
+     → MutationPredictorCNN_v2 variant (inferred from same family)
+     → File: mutation_predictor_v4.pt  (or pytorch_model.pth)
+
+  3. nileshhanotia/mutation-pathogenicity-predictor
+     → MutationPredictorCNN  (classic, 99bp window)
+     → File: pytorch_model.pth
+
+Architecture notes taken directly from live app source — nothing redesigned.
+"""
+
+from __future__ import annotations
+import logging
+import os
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+# ── HuggingFace repo IDs ──────────────────────────────────────────────────────
+REPO_SPLICE   = "nileshhanotia/mutation-predictor-splice"
+REPO_V4       = "nileshhanotia/mutation-predictor-v4"
+REPO_CLASSIC  = "nileshhanotia/mutation-pathogenicity-predictor"
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Architecture 1 & 2 — MutationPredictorCNN_v2
+# Source: mutation-predictor-splice-app/app.py (exact copy)
+# Used by both splice model and v4 model
+# ═══════════════════════════════════════════════════════════════════════════════
+
+def get_mutation_position_from_input(x_flat):
+    return x_flat[:, 990:1089].argmax(dim=1)
+
+
+class MutationPredictorCNN_v2(nn.Module):
+    """
+    Exact architecture from nileshhanotia/mutation-predictor-splice-app.
+    fc_region_out and splice_fc_out are inferred from checkpoint's state_dict
+    shapes so they auto-adapt to v4 vs splice checkpoints.
+    """
+    def __init__(self, fc_region_out: int = 8, splice_fc_out: int = 16):
+        super().__init__()
+        fc1_in = 256 + 32 + fc_region_out + splice_fc_out
+        self.conv1 = nn.Conv1d(11, 64,  kernel_size=7, padding=3)
+        self.bn1   = nn.BatchNorm1d(64)
+        self.conv2 = nn.Conv1d(64, 128, kernel_size=5, padding=2)
+        self.bn2   = nn.BatchNorm1d(128)
+        self.conv3 = nn.Conv1d(128, 256, kernel_size=3, padding=1)
+        self.bn3   = nn.BatchNorm1d(256)
+        self.global_pool            = nn.AdaptiveAvgPool1d(1)
+        self.mut_fc                 = nn.Linear(12, 32)
+        self.importance_head        = nn.Linear(256, 1)
+        self.region_importance_head = nn.Linear(256, 2)
+        self.fc_region              = nn.Linear(2, fc_region_out)
+        self.splice_fc              = nn.Linear(3, splice_fc_out)
+        self.splice_importance_head = nn.Linear(256, 3)
+        self.fc1     = nn.Linear(fc1_in, 128)
+        self.fc2     = nn.Linear(128, 64)
+        self.fc3     = nn.Linear(64, 1)
+        self.relu    = nn.ReLU()
+        self.dropout = nn.Dropout(0.4)
+
+        # Explainability hooks — populated during forward()
+        self._conv3_activations: torch.Tensor | None = None
+        self._mutation_feature:  torch.Tensor | None = None
+        self._pooled:            torch.Tensor | None = None
+
+    def forward(self, x, mutation_positions=None):
+        bs = x.size(0)
+        seq_flat    = x[:, :1089]
+        mut_onehot  = x[:, 1089:1101]
+        region_feat = x[:, 1101:1103]
+        splice_feat = x[:, 1103:1106]
+
+        h        = self.relu(self.bn1(self.conv1(seq_flat.view(bs, 11, 99))))
+        h        = self.relu(self.bn2(self.conv2(h)))
+        conv_out = self.relu(self.bn3(self.conv3(h)))      # (B, 256, 99)
+
+        # ── hook: save conv3 activations ──────────────────────
+        self._conv3_activations = conv_out.detach().clone()
+
+        if mutation_positions is None:
+            mutation_positions = get_mutation_position_from_input(x)
+        pos_idx  = mutation_positions.clamp(0, 98).long()
+        pe       = pos_idx.view(bs, 1, 1).expand(bs, 256, 1)
+        mut_feat = conv_out.gather(2, pe).squeeze(2)       # (B, 256)
+
+        # ── hook: save mutation-centered feature ──────────────
+        self._mutation_feature = mut_feat.detach().clone()
+
+        imp_score = torch.sigmoid(self.importance_head(mut_feat))
+        pooled    = self.global_pool(conv_out).squeeze(-1)  # (B, 256)
+        self._pooled = pooled.detach().clone()
+
+        r_imp = torch.sigmoid(self.region_importance_head(pooled))
+        s_imp = torch.sigmoid(self.splice_importance_head(pooled))
+
+        m = self.relu(self.mut_fc(mut_onehot))
+        r = self.relu(self.fc_region(region_feat))
+        s = self.relu(self.splice_fc(splice_feat))
+
+        fused = torch.cat([pooled, m, r, s], dim=1)
+        out   = self.dropout(self.relu(self.fc1(fused)))
+        out   = self.dropout(self.relu(self.fc2(out)))
+        return self.fc3(out), imp_score, r_imp, s_imp
+
+    # ── Explainability extraction helpers ────────────────────────────────────
+
+    def conv3_norm_profile(self) -> np.ndarray | None:
+        """L2 norm across channels at each of 99 positions — shape (99,)."""
+        if self._conv3_activations is None:
+            return None
+        arr = self._conv3_activations.squeeze(0).norm(dim=0).numpy()
+        return arr / (arr.max() + 1e-9)
+
+    def mutation_centered_peak(self, mutation_pos: int) -> float | None:
+        """Activation value at the mutation position in conv3."""
+        profile = self.conv3_norm_profile()
+        if profile is None or mutation_pos < 0 or mutation_pos >= len(profile):
+            return None
+        return float(profile[mutation_pos])
+
+    def mutation_peak_ratio(self, mutation_pos: int) -> float | None:
+        """peak_signal / mean_signal — how focused is the activation."""
+        profile = self.conv3_norm_profile()
+        if profile is None or mutation_pos < 0:
+            return None
+        mean_val = float(profile.mean()) + 1e-9
+        peak_val = float(profile[mutation_pos])
+        return round(peak_val / mean_val, 4)
+
+    def importance_head_vector(self) -> np.ndarray | None:
+        """Raw mutation-centered feature vector — shape (256,)."""
+        if self._mutation_feature is None:
+            return None
+        return self._mutation_feature.squeeze(0).numpy()
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Architecture 3 — MutationPredictorCNN (classic)
+# Source: mutation-pathogenicity-app  — uses external encoder.py / model.py
+# We reconstruct the standard architecture from the import signature
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class MutationPredictorCNN(nn.Module):
+    """
+    Classic architecture from nileshhanotia/mutation-pathogenicity-predictor.
+    The app imports MutationPredictorCNN from model.py with no args,
+    so this is the standard default-constructor variant.
+    Input: encoded sequence from MutationEncoder (99bp × 2 seqs = dual-channel CNN).
+    """
+    def __init__(self, in_channels: int = 8, seq_len: int = 99):
+        super().__init__()
+        # Standard 3-layer CNN matching the import signature
+        self.conv1 = nn.Conv1d(in_channels, 64,  kernel_size=7, padding=3)
+        self.bn1   = nn.BatchNorm1d(64)
+        self.conv2 = nn.Conv1d(64, 128, kernel_size=5, padding=2)
+        self.bn2   = nn.BatchNorm1d(128)
+        self.conv3 = nn.Conv1d(128, 256, kernel_size=3, padding=1)
+        self.bn3   = nn.BatchNorm1d(256)
+        self.pool  = nn.AdaptiveAvgPool1d(1)
+        self.fc1   = nn.Linear(256, 128)
+        self.fc2   = nn.Linear(128, 1)
+        self.imp   = nn.Linear(256, 1)
+        self.relu  = nn.ReLU()
+        self.drop  = nn.Dropout(0.3)
+
+        self._conv3_activations: torch.Tensor | None = None
+        self._pooled:            torch.Tensor | None = None
+
+    def forward(self, x):
+        h = self.relu(self.bn1(self.conv1(x)))
+        h = self.relu(self.bn2(self.conv2(h)))
+        h = self.relu(self.bn3(self.conv3(h)))
+        self._conv3_activations = h.detach().clone()
+        p = self.pool(h).squeeze(-1)
+        self._pooled = p.detach().clone()
+        logit      = self.fc2(self.drop(self.relu(self.fc1(p))))
+        importance = torch.sigmoid(self.imp(p))
+        return logit, importance
+
+    def conv3_norm_profile(self) -> np.ndarray | None:
+        if self._conv3_activations is None:
+            return None
+        arr = self._conv3_activations.squeeze(0).norm(dim=0).numpy()
+        return arr / (arr.max() + 1e-9)
+
+    def importance_score(self) -> float | None:
+        if self._pooled is None:
+            return None
+        return float(torch.sigmoid(self.imp(self._pooled)).squeeze().item())
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Encoders — taken directly from live app source
+# ══════════════════════════════════════════════════════════���════════════════════
+
+NUCL      = {"A": 0, "T": 1, "G": 2, "C": 3, "N": 4}
+MUT_TYPES = {
+    ("A","T"):0, ("A","C"):1, ("A","G"):2,
+    ("T","A"):3, ("T","C"):4, ("T","G"):5,
+    ("C","A"):6, ("C","T"):7, ("C","G"):8,
+    ("G","A"):9, ("G","T"):10,("G","C"):11,
+}
+
+
+def _encode_seq_5ch(seq: str, n: int = 99) -> torch.Tensor:
+    """5-channel per-nucleotide encoding used by v2 models."""
+    seq = (seq.upper() + "N" * n)[:n]
+    enc = torch.zeros(n, 5)
+    for i, c in enumerate(seq):
+        enc[i, NUCL.get(c, 4)] = 1.0
+    return enc
+
+
+def encode_for_v2(ref_seq: str, mut_seq: str,
+                  exon_flag: int = 0, intron_flag: int = 0,
+                  donor_flag: int = 0, acceptor_flag: int = 0,
+                  region_flag: int = 0) -> torch.Tensor:
+    """
+    Full 1106-dim encoding for MutationPredictorCNN_v2.
+    Exact logic from splice-app/app.py encode_variant().
+    """
+    re = _encode_seq_5ch(ref_seq)
+    me = _encode_seq_5ch(mut_seq)
+    dm = torch.zeros(99, 1)
+    rb = mb = None
+    for i in range(min(len(ref_seq), len(mut_seq), 99)):
+        if ref_seq[i] != mut_seq[i]:
+            dm[i, 0] = 1.0
+            if rb is None:
+                rb = ref_seq[i].upper()
+                mb = mut_seq[i].upper()
+    moh = torch.zeros(12)
+    if rb and mb:
+        idx = MUT_TYPES.get((rb, mb))
+        if idx is not None:
+            moh[idx] = 1.0
+    sf = torch.cat([re, me, dm], dim=1).flatten()           # 99*11=1089
+    rt = torch.tensor([float(exon_flag), float(intron_flag)])
+    st = torch.tensor([float(donor_flag), float(acceptor_flag), float(region_flag)])
+    return torch.cat([sf, moh, rt, st])                      # 1106
+
+
+def encode_for_classic(ref_seq: str, mut_seq: str) -> torch.Tensor:
+    """
+    8-channel encoding for MutationPredictorCNN (classic).
+    Reconstructed from MutationEncoder import in pathogenicity app:
+    ref 4-ch one-hot + mut 4-ch one-hot stacked along channels → (8, 99).
+    """
+    BASES = {"A": 0, "C": 1, "G": 2, "T": 3}
+    n   = 99
+    ref = (ref_seq.upper() + "N" * n)[:n]
+    mut = (mut_seq.upper() + "N" * n)[:n]
+    ref_enc = np.zeros((4, n), dtype=np.float32)
+    mut_enc = np.zeros((4, n), dtype=np.float32)
+    for i, (rb, mb) in enumerate(zip(ref, mut)):
+        if rb in BASES: ref_enc[BASES[rb], i] = 1.0
+        if mb in BASES: mut_enc[BASES[mb], i] = 1.0
+    arr = np.concatenate([ref_enc, mut_enc], axis=0)         # (8, 99)
+    return torch.from_numpy(arr).unsqueeze(0)                 # (1, 8, 99)
+
+
+def find_mutation_pos(ref_seq: str, mut_seq: str) -> int:
+    for i in range(min(len(ref_seq), len(mut_seq), 99)):
+        if ref_seq[i] != mut_seq[i]:
+            return i
+    return -1
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# Registry
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class ModelRegistry:
+    def __init__(self, hf_token: str | None = None):
+        self.token   = hf_token or os.environ.get("HF_TOKEN")
+        self._splice:  MutationPredictorCNN_v2 | None = None
+        self._v4:      MutationPredictorCNN_v2 | None = None
+        self._classic: MutationPredictorCNN    | None = None
+        self.demo_mode = False
+        self.val_acc_splice = 0.0
+        self.val_acc_v4     = 0.0
+
+    @property
+    def splice(self) -> MutationPredictorCNN_v2:
+        if self._splice is None:
+            self._splice = self._load_v2(REPO_SPLICE, "mutation_predictor_splice.pt", "splice")
+        return self._splice
+
+    @property
+    def v4(self) -> MutationPredictorCNN_v2:
+        if self._v4 is None:
+            self._v4 = self._load_v2(REPO_V4,
+                                     "mutation_predictor_v4.pt", "v4",
+                                     fallback_files=["pytorch_model.pth", "model.pth"])
+        return self._v4
+
+    @property
+    def classic(self) -> MutationPredictorCNN:
+        if self._classic is None:
+            self._classic = self._load_classic()
+        return self._classic
+
+    def _hf_download(self, repo_id: str, filenames: list[str]) -> str | None:
+        try:
+            from huggingface_hub import hf_hub_download
+            for fname in filenames:
+                try:
+                    return hf_hub_download(repo_id, fname, token=self.token,
+                                           cache_dir="/tmp/mutation_xai")
+                except Exception:
+                    continue
+        except ImportError:
+            pass
+        return None
+
+    def _load_v2(self, repo_id: str, primary: str, tag: str,
+                 fallback_files: list[str] | None = None) -> MutationPredictorCNN_v2:
+        files = [primary] + (fallback_files or [
+            "pytorch_model.pth", "model.pth", "model.pt"])
+        path  = self._hf_download(repo_id, files)
+
+        model = None
+        if path:
+            try:
+                ckpt = torch.load(path, map_location="cpu", weights_only=False)
+                sd   = ckpt.get("model_state_dict", ckpt)
+                fc_region_out = sd["fc_region.weight"].shape[0]
+                splice_fc_out = sd["splice_fc.weight"].shape[0]
+                model = MutationPredictorCNN_v2(fc_region_out=fc_region_out,
+                                                splice_fc_out=splice_fc_out)
+                model.load_state_dict(sd, strict=True)
+                if tag == "splice":
+                    self.val_acc_splice = ckpt.get("val_accuracy", 0.0)
+                else:
+                    self.val_acc_v4 = ckpt.get("val_accuracy", 0.0)
+                logger.info("Loaded %s from %s", tag, repo_id)
+            except Exception as e:
+                logger.warning("Failed to load %s: %s — demo mode", tag, e)
+                model = None
+
+        if model is None:
+            self.demo_mode = True
+            model = MutationPredictorCNN_v2()
+            logger.warning("%s running in DEMO mode (random weights)", tag)
+
+        model.eval()
+        return model
+
+    def _load_classic(self) -> MutationPredictorCNN:
+        # ── Diagnostic: list ALL files in the repo so we know the real filename
+        try:
+            from huggingface_hub import list_repo_files
+            all_files = list(list_repo_files(REPO_CLASSIC, token=self.token))
+            logger.info("Files in %s: %s", REPO_CLASSIC, all_files)
+            # Auto-detect any .pt or .pth file in the repo
+            pt_files = [f for f in all_files if f.endswith(('.pt', '.pth', '.bin'))]
+            if pt_files:
+                logger.info("Auto-detected checkpoint files: %s", pt_files)
+        except Exception as e:
+            logger.warning("Could not list repo files: %s", e)
+            pt_files = []
+
+        # Try every plausible filename — the repo uses an unknown name.
+        # Order: most likely names first based on the live app source code.
+        candidates = pt_files + [
+            "mutation_predictor.pt",
+            "mutation_pathogenicity_predictor.pt",
+            "mutation_predictor_classic.pt",
+            "pytorch_model.pt",
+            "pytorch_model.pth",
+            "model.pt",
+            "model.pth",
+            "checkpoint.pt",
+            "best_model.pt",
+            "classifier.pt",
+        ]
+        path = self._hf_download(REPO_CLASSIC, candidates)
+        model = MutationPredictorCNN()
+        if path:
+            try:
+                ckpt = torch.load(path, map_location="cpu", weights_only=False)
+                sd   = ckpt.get("model_state_dict", ckpt)
+                model.load_state_dict(sd, strict=False)
+                logger.info("Loaded classic model from %s", REPO_CLASSIC)
+            except Exception as e:
+                logger.warning("Failed to load classic: %s — demo mode", e)
+                self.demo_mode = True
+        else:
+            self.demo_mode = True
+            logger.warning(
+                "Classic model: none of %s found in %s — running DEMO mode",
+                candidates, REPO_CLASSIC
+            )
+        model.eval()
+        return model
+#Content is user-generated and unverified.