Upload model_loader.py

model_loader.py

@@ -0,0 +1,425 @@
+"""
+model_loader.py — PeVe Unified Space Model Loading Module
+
+Loading logic adapted from:
+  - nileshhanotia/mutation-predictor-splice-app  (app.py)
+  - nileshhanotia/mutation-pathogenicity-app      (app.py)
+  - nileshhanotia/mutation-explainable-v6         (model_v6.pkl)
+
+Provides:
+    load_splice_model()   → (model, status_dict)
+    load_context_model()  → (model, status_dict)
+    load_protein_model()  → (model, status_dict)
+    get_model_status()    → combined status dict
+"""
+
+import os
+import traceback
+import pickle
+
+import torch
+import torch.nn as nn
+
+# ── Optional: set HF token for private repos ───────────────────────────────
+# Either set the environment variable HF_TOKEN before running, or hard-code
+# a token here (not recommended for public repos).
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
+# ══════════════════════════════════════════════════════════════════════════════
+# MODULE-LEVEL MODEL HANDLES
+# These are populated by the load_*() functions below.
+# ══════════════════════════════════════════════════════════════════════════════
+
+_splice_model = None
+_context_model = None
+_protein_model = None
+
+# ══════════════════════════════════════════════════════════════════════════════
+# ARCHITECTURE — Splice Model
+# Adapted from: nileshhanotia/mutation-predictor-splice-app  app.py
+# ══════════════════════════════════════════════════════════════════════════════
+
+def _get_mutation_position_from_input(x_flat):
+    """Internal helper used by MutationPredictorCNN_v2.forward()."""
+    return x_flat[:, 990:1089].argmax(dim=1)
+
+
+class MutationPredictorCNN_v2(nn.Module):
+    """
+    Splice-aware mutation predictor.
+    Architecture copied verbatim from mutation-predictor-splice-app/app.py
+    to guarantee weight compatibility.
+    """
+
+    def __init__(self, fc_region_out=8, splice_fc_out=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)
+
+    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)))
+
+        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)
+
+        imp_score = torch.sigmoid(self.importance_head(mut_feat))
+        pooled    = self.global_pool(conv_out).squeeze(-1)
+        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)))
+        logit = self.fc3(out)
+        return logit, imp_score, r_imp, s_imp
+
+
+# ══════════════════════════════════════════���═══════════════════════════════════
+# ARCHITECTURE — Context (401 bp CNN) Model
+# Adapted from: nileshhanotia/mutation-predictor-v4
+# ══════════════════════════════════════════════════════════════════════════════
+
+class MutationContextCNN(nn.Module):
+    """
+    401 bp context window CNN for mutation pathogenicity.
+    Architecture mirrors the v4 space model; weights loaded from state dict.
+    If the actual v4 architecture differs, the load_state_dict call will raise
+    a descriptive KeyError that will be captured in the status dict.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv1d(5, 64,  kernel_size=11, padding=5)
+        self.bn1   = nn.BatchNorm1d(64)
+        self.conv2 = nn.Conv1d(64, 128, kernel_size=7, padding=3)
+        self.bn2   = nn.BatchNorm1d(128)
+        self.conv3 = nn.Conv1d(128, 256, kernel_size=5, padding=2)
+        self.bn3   = nn.BatchNorm1d(256)
+        self.pool  = nn.AdaptiveAvgPool1d(1)
+        self.fc1   = nn.Linear(256, 128)
+        self.fc2   = nn.Linear(128, 64)
+        self.fc3   = nn.Linear(64, 1)
+        self.relu  = nn.ReLU()
+        self.drop  = nn.Dropout(0.3)
+
+    def forward(self, x):
+        # x: (batch, seq_len, channels) → permute → (batch, channels, seq_len)
+        h = x.permute(0, 2, 1)
+        h = self.relu(self.bn1(self.conv1(h)))
+        h = self.relu(self.bn2(self.conv2(h)))
+        h = self.relu(self.bn3(self.conv3(h)))
+        h = self.pool(h).squeeze(-1)
+        h = self.drop(self.relu(self.fc1(h)))
+        h = self.drop(self.relu(self.fc2(h)))
+        return self.fc3(h)
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# LOADER — Splice Model
+# ══════════════════════════════════════════════════════════════════════════════
+
+def load_splice_model():
+    """
+    Load MutationPredictorCNN_v2 from nileshhanotia/mutation-predictor-splice.
+
+    Loading logic adapted from:
+      nileshhanotia/mutation-predictor-splice-app  app.py
+        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        sd   = ckpt["model_state_dict"]
+
+    Returns
+    -------
+    (model | None, {"loaded": bool, "error_message": str})
+    """
+    global _splice_model
+
+    status = {"loaded": False, "error_message": ""}
+
+    try:
+        from huggingface_hub import hf_hub_download  # local import for clarity
+
+        MODEL_REPO     = "nileshhanotia/mutation-predictor-splice"
+        MODEL_FILENAME = "mutation_predictor_splice.pt"
+
+        print(f"[splice] Downloading {MODEL_FILENAME} from {MODEL_REPO} …")
+        ckpt_path = hf_hub_download(
+            repo_id=MODEL_REPO,
+            filename=MODEL_FILENAME,
+            token=HF_TOKEN,
+        )
+
+        print(f"[splice] Loading checkpoint from {ckpt_path} …")
+        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        sd   = ckpt["model_state_dict"]
+
+        # Infer architecture hyper-params from the state dict (exact pattern from app.py)
+        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)
+        model.eval()
+
+        val_acc = ckpt.get("val_accuracy", float("nan"))
+        print(f"[splice] ✓ Loaded. val_accuracy={val_acc:.4f} | "
+              f"fc_region_out={fc_region_out} | splice_fc_out={splice_fc_out}")
+
+        _splice_model = model
+        status["loaded"] = True
+
+    except Exception:
+        tb = traceback.format_exc()
+        print(f"[splice] ✗ FAILED to load:\n{tb}")
+        status["error_message"] = tb
+        _splice_model = None
+
+    return _splice_model, status
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# LOADER — Context Model (401 bp CNN, mutation-predictor-v4)
+# ══════════════════════════════════════════════════════════════════════════════
+
+def load_context_model():
+    """
+    Load the 401 bp context CNN from nileshhanotia/mutation-predictor-v4.
+
+    Loading logic adapted from:
+      nileshhanotia/mutation-pathogenicity-app  app.py
+        checkpoint = torch.load(MODEL_PATH, map_location=device)
+        model.load_state_dict(checkpoint["model_state_dict"])
+
+    Returns
+    -------
+    (model | None, {"loaded": bool, "error_message": str})
+    """
+    global _context_model
+
+    status = {"loaded": False, "error_message": ""}
+
+    try:
+        from huggingface_hub import hf_hub_download
+
+        MODEL_REPO     = "nileshhanotia/mutation-predictor-v4"
+        # Try common checkpoint filenames used in HF spaces
+        CANDIDATE_FILENAMES = [
+            "pytorch_model.pth",
+            "mutation_predictor_v4.pt",
+            "model.pt",
+            "model.pth",
+            "checkpoint.pth",
+        ]
+
+        ckpt_path = None
+        last_error = ""
+        for fname in CANDIDATE_FILENAMES:
+            try:
+                print(f"[context] Trying {fname} from {MODEL_REPO} …")
+                ckpt_path = hf_hub_download(
+                    repo_id=MODEL_REPO,
+                    filename=fname,
+                    token=HF_TOKEN,
+                )
+                print(f"[context] Found: {fname}")
+                break
+            except Exception as e:
+                last_error = str(e)
+                continue
+
+        if ckpt_path is None:
+            raise FileNotFoundError(
+                f"None of the candidate filenames found in {MODEL_REPO}. "
+                f"Last error: {last_error}"
+            )
+
+        print(f"[context] Loading checkpoint from {ckpt_path} …")
+        checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+
+        # Support both raw state-dict and wrapped checkpoint
+        if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
+            sd = checkpoint["model_state_dict"]
+        elif isinstance(checkpoint, dict) and "state_dict" in checkpoint:
+            sd = checkpoint["state_dict"]
+        else:
+            sd = checkpoint  # assume it IS the state dict
+
+        model = MutationContextCNN()
+        model.load_state_dict(sd, strict=False)  # strict=False tolerates minor arch diffs
+        model.eval()
+
+        print("[context] ✓ Loaded MutationContextCNN (401 bp).")
+        _context_model = model
+        status["loaded"] = True
+
+    except Exception:
+        tb = traceback.format_exc()
+        print(f"[context] ✗ FAILED to load:\n{tb}")
+        status["error_message"] = tb
+        _context_model = None
+
+    return _context_model, status
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# LOADER — Protein Model (XGBoost .pkl from mutation-explainable-v6)
+# ══════════════════════════════════════════════════════════════════════════════
+
+def load_protein_model():
+    """
+    Load the pickled XGBoost model from nileshhanotia/mutation-explainable-v6.
+
+    Loading logic adapted from:
+      nileshhanotia/mutation-explainable-v6  (model_v6.pkl)
+
+    Uses Python pickle / joblib — NOT XGBoost Booster.load_model().
+    The model is already stored as a complete trained sklearn-compatible object.
+
+    Returns
+    -------
+    (model | None, {"loaded": bool, "error_message": str})
+    """
+    global _protein_model
+
+    status = {"loaded": False, "error_message": ""}
+
+    try:
+        from huggingface_hub import hf_hub_download
+
+        MODEL_REPO     = "nileshhanotia/mutation-explainable-v6"
+        MODEL_FILENAME = "model_v6.pkl"
+
+        print(f"[protein] Downloading {MODEL_FILENAME} from {MODEL_REPO} …")
+        pkl_path = hf_hub_download(
+            repo_id=MODEL_REPO,
+            filename=MODEL_FILENAME,
+            token=HF_TOKEN,
+        )
+
+        print(f"[protein] Loading pickle from {pkl_path} …")
+        # Try joblib first (common for sklearn/xgboost pipelines), fall back to pickle
+        try:
+            import joblib
+            model = joblib.load(pkl_path)
+            print("[protein] Loaded via joblib.")
+        except Exception:
+            with open(pkl_path, "rb") as f:
+                model = pickle.load(f)
+            print("[protein] Loaded via pickle.")
+
+        print(f"[protein] ✓ Loaded protein model: {type(model).__name__}")
+        _protein_model = model
+        status["loaded"] = True
+
+    except Exception:
+        tb = traceback.format_exc()
+        print(f"[protein] ✗ FAILED to load:\n{tb}")
+        status["error_message"] = tb
+        _protein_model = None
+
+    return _protein_model, status
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# STATUS AGGREGATOR
+# ══════════════════════════════════════════════════════════════════════════════
+
+def get_model_status() -> dict:
+    """
+    Load all three models and return a unified status dictionary.
+
+    Returns
+    -------
+    {
+        "splice":  {"loaded": bool, "error_message": str},
+        "context": {"loaded": bool, "error_message": str},
+        "protein": {"loaded": bool, "error_message": str},
+    }
+    """
+    print("=" * 60)
+    print("PeVe — starting unified model loading")
+    print("=" * 60)
+
+    _, splice_status  = load_splice_model()
+    _, context_status = load_context_model()
+    _, protein_status = load_protein_model()
+
+    status = {
+        "splice":  splice_status,
+        "context": context_status,
+        "protein": protein_status,
+    }
+
+    # Summary report
+    print("\n" + "=" * 60)
+    print("PeVe — model loading complete")
+    print("=" * 60)
+    for name, s in status.items():
+        icon = "✓" if s["loaded"] else "✗"
+        print(f"  [{icon}] {name:10s} loaded={s['loaded']}")
+    print("=" * 60 + "\n")
+
+    return status
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# PUBLIC ACCESSORS
+# ══════════════════════════════════════════════════════════════════════════════
+
+def get_splice_model():
+    """Return the loaded splice model handle (None if not loaded)."""
+    return _splice_model
+
+
+def get_context_model():
+    """Return the loaded context model handle (None if not loaded)."""
+    return _context_model
+
+
+def get_protein_model():
+    """Return the loaded protein model handle (None if not loaded)."""
+    return _protein_model
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# SELF-TEST
+# ══════════════════════════════════════════════════════════════════════════════
+
+if __name__ == "__main__":
+    print("Testing model loading...")
+    status = get_model_status()
+    print(status)