""" model_loader.py — PeVe Unified Space Model Loader ================================================== Loads all three PeVe models with full error handling and per-model status reporting. Loading logic adapted from: - mutation-predictor-splice-app app.py (splice CNN model) - mutation-predictor-v4 app.py (context 401bp CNN model) - mutation-explainable-v6 model_v6.pkl (protein XGBoost model) """ from __future__ import annotations import os import pickle import traceback import warnings from pathlib import Path from typing import Any, Optional, Tuple # ── HF token (set HF_TOKEN env var for private repos) ──────────────────────── HF_TOKEN: Optional[str] = os.environ.get("HF_TOKEN", None) # ── Repository IDs ──────────────────────────────────────────────────────────── SPLICE_REPO = "nileshhanotia/mutation-predictor-splice" CONTEXT_REPO = "nileshhanotia/mutation-predictor-v4" PROTEIN_REPO = "nileshhanotia/mutation-explainable-v6" # ── Cached model objects ────────────────────────────────────────────────────── _splice_model: Optional[Any] = None _context_model: Optional[Any] = None _protein_model: Optional[Any] = None # ── Per-model status store ──────────────────────────────────────────────────── _status: dict = { "splice": {"loaded": False, "error_message": "Not yet loaded"}, "context": {"loaded": False, "error_message": "Not yet loaded"}, "protein": {"loaded": False, "error_message": "Not yet loaded"}, } # ══════════════════════════════════════════════════════════════════════════════ # SPLICE MODEL ARCHITECTURE # Adapted from: mutation-predictor-splice-app/app.py # ══════════════════════════════════════════════════════════════════════════════ def _build_splice_architecture(state_dict: dict): """Reconstruct MutationPredictorCNN_v2 from its state dict shape.""" import torch import torch.nn as nn class MutationPredictorCNN_v2(nn.Module): 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 = x[:, 990:1089].argmax(dim=1) 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 fc_region_out = state_dict["fc_region.weight"].shape[0] splice_fc_out = state_dict["splice_fc.weight"].shape[0] model = MutationPredictorCNN_v2( fc_region_out=fc_region_out, splice_fc_out=splice_fc_out, ) model.load_state_dict(state_dict) model.eval() return model # ══════════════════════════════════════════════════════════════════════════════ # PUBLIC LOADER FUNCTIONS # ══════════════════════════════════════════════════════════════════════════════ def load_splice_model() -> Tuple[Optional[Any], dict]: """ Load the splice-aware CNN model from: nileshhanotia/mutation-predictor-splice → mutation_predictor_splice.pt Loading logic adapted from: mutation-predictor-splice-app/app.py Returns: (model | None, status_dict) """ global _splice_model print("[PeVe] Loading splice model ...") try: import torch from huggingface_hub import hf_hub_download ckpt_path = hf_hub_download( repo_id=SPLICE_REPO, filename="mutation_predictor_splice.pt", token=HF_TOKEN, ) ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) sd = ckpt["model_state_dict"] model = _build_splice_architecture(sd) val_acc = ckpt.get("val_accuracy", float("nan")) print(f"[PeVe] Splice model ready val_acc={val_acc:.4f}") _splice_model = model _status["splice"] = {"loaded": True, "error_message": ""} return model, _status["splice"] except Exception: tb = traceback.format_exc() print(f"[PeVe] ERROR loading splice model:\n{tb}") _status["splice"] = {"loaded": False, "error_message": tb} return None, _status["splice"] def load_context_model() -> Tuple[Optional[Any], dict]: """ Load the 401 bp sequence CNN (context) model from: nileshhanotia/mutation-predictor-v4 Tries several common checkpoint filenames then falls back to a snapshot directory scan. Loading logic adapted from: mutation-predictor-v4 app.py Returns: (model | None, status_dict) """ global _context_model print("[PeVe] Loading context model ...") try: import torch from huggingface_hub import hf_hub_download, snapshot_download model = None loaded = False # ── Attempt 1: named checkpoint candidates ──────────────────────── for candidate in [ "mutation_predictor_v4.pt", "model.pt", "model.pth", "pytorch_model.bin", ]: try: path = hf_hub_download( repo_id=CONTEXT_REPO, filename=candidate, token=HF_TOKEN, ) obj = torch.load(path, map_location="cpu", weights_only=False) if isinstance(obj, dict): if "model_state_dict" in obj: # State-dict only — store raw dict; upstream code must # reconstruct the architecture and call load_state_dict. warnings.warn( f"[PeVe] Context checkpoint '{candidate}' contains a " "state_dict but no architecture class is defined here. " "Add the MutationPredictorCNN_v4 class to model_loader.py " "and call load_state_dict() to fully restore the model." ) model = obj elif "model" in obj: model = obj["model"] else: model = obj else: model = obj # nn.Module stored directly if hasattr(model, "eval"): model.eval() print(f"[PeVe] Context model loaded from '{candidate}'") loaded = True break except Exception: pass # try next candidate # ── Attempt 2: snapshot directory scan ─────────────────────────── if not loaded: local = snapshot_download(repo_id=CONTEXT_REPO, token=HF_TOKEN) candidates = ( list(Path(local).glob("*.pt")) + list(Path(local).glob("*.pth")) + list(Path(local).glob("*.bin")) ) if not candidates: raise FileNotFoundError( f"No PyTorch weight file found in {CONTEXT_REPO} snapshot." ) obj = torch.load(candidates[0], map_location="cpu", weights_only=False) model = obj.get("model", obj) if isinstance(obj, dict) else obj if hasattr(model, "eval"): model.eval() print(f"[PeVe] Context model loaded via snapshot → {candidates[0].name}") _context_model = model _status["context"] = {"loaded": True, "error_message": ""} return model, _status["context"] except Exception: tb = traceback.format_exc() print(f"[PeVe] ERROR loading context model:\n{tb}") _status["context"] = {"loaded": False, "error_message": tb} return None, _status["context"] def load_protein_model() -> Tuple[Optional[Any], dict]: """ Load the XGBoost protein-level model from: nileshhanotia/mutation-explainable-v6 → model_v6.pkl Uses Python pickle exactly as the original app does. Does NOT use XGBoost Booster.load_model() — the file is a pickled sklearn-compatible XGBClassifier / Pipeline object. Loading logic adapted from: mutation-explainable-v6 model_v6.pkl / app.py Returns: (model | None, status_dict) """ global _protein_model print("[PeVe] Loading protein (XGBoost) model ...") try: from huggingface_hub import hf_hub_download pkl_path = hf_hub_download( repo_id=PROTEIN_REPO, filename="model_v6.pkl", token=HF_TOKEN, ) with open(pkl_path, "rb") as fh: model = pickle.load(fh) print(f"[PeVe] Protein model ready type={type(model).__name__}") _protein_model = model _status["protein"] = {"loaded": True, "error_message": ""} return model, _status["protein"] except Exception: tb = traceback.format_exc() print(f"[PeVe] ERROR loading protein model:\n{tb}") _status["protein"] = {"loaded": False, "error_message": tb} return None, _status["protein"] # ══════════════════════════════════════════════════════════════════════════════ # STATUS + CONVENIENCE ACCESSORS # ══════════════════════════════════════════════════════════════════════════════ def get_model_status() -> dict: """ Trigger loading of all three models (lazy, cached) and return: { "splice": {"loaded": bool, "error_message": str}, "context": {"loaded": bool, "error_message": str}, "protein": {"loaded": bool, "error_message": str}, } Subsequent calls are free — models are cached in module-level globals. If a model fails to load, 'loaded' is False and 'error_message' contains the full Python traceback string. """ if _splice_model is None and not _status["splice"]["loaded"]: load_splice_model() if _context_model is None and not _status["context"]["loaded"]: load_context_model() if _protein_model is None and not _status["protein"]["loaded"]: load_protein_model() return dict(_status) def get_splice_model() -> Optional[Any]: """Return cached splice model (load on first call).""" if _splice_model is None: load_splice_model() return _splice_model def get_context_model() -> Optional[Any]: """Return cached context model (load on first call).""" if _context_model is None: load_context_model() return _context_model def get_protein_model() -> Optional[Any]: """Return cached protein model (load on first call).""" if _protein_model is None: load_protein_model() return _protein_model # ══════════════════════════════════════════════════════════════════════════════ # SELF-TEST # ══════════════════════════════════════════════════════════════════════════════ if __name__ == "__main__": import json print("=" * 62) print(" PeVe model_loader.py — self-test") print("=" * 62) status = get_model_status() print("\n── Final status ──────────────────────────────────────────") for key, info in status.items(): icon = "✅" if info["loaded"] else "❌" print(f" {icon} {key:10s} loaded={info['loaded']}") if not info["loaded"]: lines = info["error_message"].strip().splitlines() for line in lines[:6]: print(f" {line}") print("\n── Status dict (JSON) ────────────────────────────────────") printable = { k: { "loaded": v["loaded"], "error_message": ( v["error_message"][:300] + " …[truncated]" if len(v["error_message"]) > 300 else v["error_message"] ), } for k, v in status.items() } print(json.dumps(printable, indent=2))