model_loader_old.py

"""
model_loader.py — PeVe v1.4
============================

What app.py needs (read from app.py source):

  from model_loader import get_splice_model, get_context_model, get_protein_model

  model, tokenizer = get_splice_model()
    → model accepts: torch.Tensor of shape (1, 401, 8) or (1, 401, 8) + flags
    → used inside torch.no_grad()
    → returns tensor or tuple[tensor, ...]

  model, tokenizer = get_context_model()
    → same calling convention as splice

  model = get_protein_model()
    → used with: xgb.DMatrix(X, feature_names=[...])
    → model.predict(dmat)  → float array
    → also passed to shap.TreeExplainer(model)

Model sources (from the Space app.py files):
  splice  → nileshhanotia/mutation-predictor-splice
              file: mutation_predictor_splice.pt
              arch: MutationPredictorCNN_v2  (input flat 1106)
              NOTE: app.py passes (1, 401, 8) tensor — loader must reshape

  context → nileshhanotia/mutation-predictor-v4
              file: mutation_predictor_splice_v4.pt
              arch: MutationPredictorCNN_v4 (4-tensor forward: seq,mut,region,splice)
              NOTE: app.py passes (1, 401, 8) tensor — loader wraps forward

  protein → nileshhanotia/mutation-pathogenicity-predictor
              file: *.json / *.ubj / *.model / *.pkl
              NOTE: app.py uses xgb.DMatrix + shap — MUST be XGBoost
              If file is actually a CNN checkpoint, we wrap it as an
              XGBoost-compatible object so app.py code paths still work.
"""
from __future__ import annotations

import os
import pickle
import traceback
import warnings
from pathlib import Path
from typing import Any

import numpy as np

# ── HF token ──────────────────────────────────────────────────────────────────
_HF_TOKEN: str | None = (
    os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
)

# ── Model repo IDs ─────────────────────────────────────────────────────────────
_REPO_SPLICE   = "nileshhanotia/mutation-predictor-splice"
_REPO_CONTEXT  = "nileshhanotia/mutation-predictor-v4"
_REPO_PROTEIN  = "nileshhanotia/mutation-pathogenicity-predictor"

# ── Global cache ───────────────────────────────────────────────────────────────
_splice_model   = None
_splice_tok     = None
_context_model  = None
_context_tok    = None
_protein_model  = None

# ── Structured status dicts ────────────────────────────────────────────────────
splice_model_status:  dict = {"loaded": False, "error_message": None}
context_model_status: dict = {"loaded": False, "error_message": None}
protein_model_status: dict = {"loaded": False, "error_message": None}


# ══════════════════════════════════════════════════════════════════════════════
# Model Architecture definitions
# (must match checkpoint shapes exactly)
# ══════════════════════════════════════════════════════════════════════════════

def _build_splice_arch(sd: dict):
    """
    MutationPredictorCNN_v2 — infer fc_region_out and splice_fc_out from
    the checkpoint's weight shapes, exactly as the Space app does.

    Forward signature in the Space:
        logit, imp_score, r_imp, s_imp = model(flat_tensor, mutation_positions)

    app.py passes tensors of shape (1, 401, 8).  We need an adapter.
    """
    import torch
    import torch.nn as nn
    import torch.nn.functional as F

    fc_region_out = sd["fc_region.weight"].shape[0]
    splice_fc_out = sd["splice_fc.weight"].shape[0]

    class MutationPredictorCNN_v2(nn.Module):
        def __init__(self):
            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_flat(self, x_flat, mutation_positions=None):
            """Original forward for flat 1106-dim input."""
            bs = x_flat.size(0)
            seq_flat   = x_flat[:, :1089]
            mut_onehot = x_flat[:, 1089:1101]
            region_feat= x_flat[:, 1101:1103]
            splice_feat= x_flat[:, 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_flat[:, 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

        def forward(self, x, mutation_positions=None):
            """
            Accept whatever shape app.py sends:
              (B, 401, 8) — raw encoded window from app.py
              (B, 1106)   — flat input (native format)
            """
            if x.dim() == 3:
                # app.py sends (B, 401, 8) — flatten and zero-pad to 1106
                bs = x.size(0)
                flat = x.reshape(bs, -1)                        # → (B, 3208)
                # Take first 1089 dims (99*11), pad mut/region/splice to zeros
                seq_part = flat[:, :1089]
                pad = torch.zeros(bs, 1106 - 1089, device=x.device)
                flat_padded = torch.cat([seq_part, pad], dim=1)  # → (B, 1106)
                return self._forward_flat(flat_padded, mutation_positions)
            return self._forward_flat(x, mutation_positions)

    model = MutationPredictorCNN_v2()
    model.load_state_dict(sd)
    return model


def _build_context_arch(sd: dict):
    """
    MutationPredictorCNN_v4 — 4-input forward (seq, mut, region, splice).
    app.py passes a single (B, 401, 8) tensor, so forward() adapts.
    """
    import torch
    import torch.nn as nn

    class MutationPredictorCNN_v4(nn.Module):
        def __init__(self):
            super().__init__()
            self.conv1 = nn.Conv1d(11, 64,  7, padding=3)
            self.conv2 = nn.Conv1d(64, 128, 5, padding=2)
            self.conv3 = nn.Conv1d(128,256, 3, padding=1)
            self.pool  = nn.AdaptiveAvgPool1d(1)
            self.mut_fc    = nn.Linear(12, 32)
            self.region_fc = nn.Linear(2,   8)
            self.splice_fc = nn.Linear(3,  16)
            self.fc1 = nn.Linear(312, 128)
            self.fc2 = nn.Linear(128,  64)
            self.fc3 = nn.Linear(64,    1)
            self.relu    = nn.ReLU()
            self.dropout = nn.Dropout(0.3)

        def _forward_native(self, seq, mut, region, splice):
            x = self.relu(self.conv1(seq))
            x = self.relu(self.conv2(x))
            x = self.relu(self.conv3(x))
            x = self.pool(x).squeeze(-1)
            m = self.relu(self.mut_fc(mut))
            r = self.relu(self.region_fc(region))
            s = self.relu(self.splice_fc(splice))
            x = torch.cat([x, m, r, s], dim=1)
            x = self.dropout(self.relu(self.fc1(x)))
            x = self.relu(self.fc2(x))
            return self.fc3(x)

        def forward(self, x, *args):
            """
            Accept:
              (B, 401, 8)  → reshape to 99*11 window, zero-pad aux inputs
              (seq, mut, region, splice) tensors — native
            """
            import torch
            if isinstance(x, torch.Tensor) and x.dim() == 3:
                bs = x.size(0)
                flat = x.reshape(bs, -1)
                seq_flat = flat[:, :1089].view(bs, 11, 99)
                mut    = torch.zeros(bs, 12,  device=x.device)
                region = torch.zeros(bs,  2,  device=x.device)
                splice = torch.zeros(bs,  3,  device=x.device)
                return self._forward_native(seq_flat, mut, region, splice)
            # flat 1089+12+2+3 = 1106 case
            if isinstance(x, torch.Tensor) and x.dim() == 2:
                bs = x.size(0)
                seq_flat = x[:, :1089].view(bs, 11, 99)
                mut    = x[:, 1089:1101]
                region = x[:, 1101:1103]
                splice = x[:, 1103:1106]
                return self._forward_native(seq_flat, mut, region, splice)
            # called with 4 separate tensors
            return self._forward_native(x, *args)

    model = MutationPredictorCNN_v4()
    model.load_state_dict(sd)
    return model


class _CNNasXGB:
    """
    Wrapper that makes a PyTorch CNN look like an XGBoost Booster
    to satisfy the app.py code paths:
        dmat  = xgb.DMatrix(X, feature_names=[...])
        pred  = model.predict(dmat)   ← needs .predict()
        shap.TreeExplainer(model)     ← will fail gracefully; app has try/except
    """
    def __init__(self, torch_model, feature_names: list[str]):
        import torch
        self._model = torch_model
        self._model.eval()
        self._features = feature_names
        self._device   = torch.device("cpu")

    def predict(self, dmat_or_array) -> np.ndarray:
        import torch
        try:
            # xgb.DMatrix → get_data() returns scipy sparse or np array
            try:
                X = dmat_or_array.get_data().toarray()
            except Exception:
                X = np.array(dmat_or_array)
        except Exception:
            X = np.zeros((1, len(self._features)), dtype=np.float32)

        t = torch.tensor(X, dtype=torch.float32)
        with torch.no_grad():
            out = self._model(t)
            if isinstance(out, (tuple, list)):
                out = out[0]
            probs = torch.sigmoid(out).cpu().numpy().flatten()
        return probs

    # Allow shap.TreeExplainer to fail gracefully — app.py has try/except
    def get_booster(self):
        raise NotImplementedError("CNN wrapped as XGB — SHAP not available")


# ══════════════════════════════════════════════════════════════════════════════
# Internal loaders
# ══════════════════════════════════════════════════════════════════════════════

def _download_repo(repo_id: str) -> Path:
    from huggingface_hub import snapshot_download
    local = snapshot_download(repo_id=repo_id, token=_HF_TOKEN)
    p = Path(local)
    files = [f.name for f in p.rglob("*") if f.is_file()]
    print(f"[PeVe] {repo_id} files: {files}")
    return p


def _load_splice() -> tuple:
    import torch
    print(f"[PeVe] Loading splice model from {_REPO_SPLICE}")
    try:
        local = _download_repo(_REPO_SPLICE)

        # Look for the checkpoint — priority: named file, then any .pt/.pth/.bin
        candidates = (
            list(local.glob("mutation_predictor_splice.pt"))
            + list(local.glob("*.pt"))
            + list(local.glob("*.pth"))
            + list(local.glob("*.bin"))
        )
        if not candidates:
            raise FileNotFoundError(f"No checkpoint in {_REPO_SPLICE}")

        ckpt = torch.load(str(candidates[0]), map_location="cpu", weights_only=False)
        sd   = ckpt.get("model_state_dict", ckpt)

        model = _build_splice_arch(sd)
        model.eval()
        val_acc = ckpt.get("val_accuracy", "n/a")
        print(f"[PeVe] ✓ splice loaded ({candidates[0].name}) val_acc={val_acc}")
        splice_model_status.update({"loaded": True, "error_message": None})
        return model, None

    except Exception:
        tb = traceback.format_exc()
        print(f"[PeVe] ✗ splice load failed:\n{tb}")
        splice_model_status.update({"loaded": False, "error_message": tb})
        return None, None


def _load_context() -> tuple:
    import torch
    print(f"[PeVe] Loading context model from {_REPO_CONTEXT}")
    try:
        local = _download_repo(_REPO_CONTEXT)

        candidates = (
            list(local.glob("mutation_predictor_splice_v4.pt"))
            + list(local.glob("*.pt"))
            + list(local.glob("*.pth"))
            + list(local.glob("*.bin"))
        )
        if not candidates:
            raise FileNotFoundError(f"No checkpoint in {_REPO_CONTEXT}")

        sd = torch.load(str(candidates[0]), map_location="cpu", weights_only=False)
        if isinstance(sd, dict) and "model_state_dict" in sd:
            sd = sd["model_state_dict"]

        model = _build_context_arch(sd)
        model.eval()
        print(f"[PeVe] ✓ context loaded ({candidates[0].name})")
        context_model_status.update({"loaded": True, "error_message": None})
        return model, None

    except Exception:
        tb = traceback.format_exc()
        print(f"[PeVe] ✗ context load failed:\n{tb}")
        context_model_status.update({"loaded": False, "error_message": tb})
        return None, None


def _load_protein():
    import xgboost as xgb
    import torch
    print(f"[PeVe] Loading protein model from {_REPO_PROTEIN}")

    _FEAT = ["gnomAD_AF", "Grantham", "Charge_change",
             "Hydrophobicity_diff", "Protein_pos_norm", "VEP_IMPACT"]

    try:
        local = _download_repo(_REPO_PROTEIN)

        # ── Try XGBoost formats first ──────────────────────────────────────
        for ext in ["*.json", "*.ubj", "*.model"]:
            for p in local.glob(ext):
                try:
                    m = xgb.Booster()
                    m.load_model(str(p))
                    print(f"[PeVe] ✓ protein loaded as XGBoost Booster ({p.name})")
                    protein_model_status.update({"loaded": True, "error_message": None})
                    return m
                except Exception as e:
                    print(f"[PeVe]   xgb.Booster failed for {p.name}: {e}")

        # ── Try pickle ────────────────────────────────────────────────────
        for p in local.glob("*.pkl"):
            try:
                with open(p, "rb") as f:
                    m = pickle.load(f)
                print(f"[PeVe] ✓ protein loaded via pickle ({p.name})")
                protein_model_status.update({"loaded": True, "error_message": None})
                return m
            except Exception as e:
                print(f"[PeVe]   pickle failed for {p.name}: {e}")

        # ── Fallback: PyTorch checkpoint — wrap as XGB-compatible ─────────
        for ext in ["*.pt", "*.pth", "*.bin"]:
            for p in local.glob(ext):
                try:
                    ckpt = torch.load(str(p), map_location="cpu", weights_only=False)
                    sd   = ckpt.get("model_state_dict", ckpt) if isinstance(ckpt, dict) else ckpt

                    if isinstance(sd, dict):
                        # Try MutationPredictorCNN (protein space model.py)
                        from torch import nn
                        class MutationPredictorCNN(nn.Module):
                            def __init__(self):
                                super().__init__()
                                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.adaptive_pool   = nn.AdaptiveAvgPool1d(1)
                                self.mut_fc          = nn.Linear(12, 32)
                                self.fc1 = nn.Linear(288, 128)
                                self.fc2 = nn.Linear(128, 64)
                                self.fc3 = nn.Linear(64, 1)
                                self.importance_head = nn.Linear(256, 1)

                            def forward(self, x):
                                import torch.nn.functional as F
                                bs = x.size(0)
                                mut_type  = x[:, 1089:1101]
                                x_seq     = x[:, :1089].view(bs, 11, 99)
                                x_conv    = F.relu(self.bn1(self.conv1(x_seq)))
                                x_conv    = nn.MaxPool1d(2, 2)(x_conv)
                                x_conv    = F.relu(self.bn2(self.conv2(x_conv)))
                                x_conv    = nn.MaxPool1d(2, 2)(x_conv)
                                x_conv    = F.relu(self.bn3(self.conv3(x_conv)))
                                x_conv    = nn.MaxPool1d(2, 2)(x_conv)
                                x_conv    = self.adaptive_pool(x_conv)
                                conv_feat = x_conv.view(bs, 256)
                                mut_feat  = F.relu(self.mut_fc(mut_type))
                                combined  = torch.cat([conv_feat, mut_feat], dim=1)
                                x_out     = F.relu(self.fc1(combined))
                                x_out     = F.relu(self.fc2(x_out))
                                cls       = torch.sigmoid(self.fc3(x_out))
                                imp       = torch.sigmoid(self.importance_head(conv_feat))
                                return cls, imp

                        torch_model = MutationPredictorCNN()
                        torch_model.load_state_dict(sd)
                        torch_model.eval()
                        wrapped = _CNNasXGB(torch_model, _FEAT)
                        print(f"[PeVe] ✓ protein loaded as CNN→XGB wrapper ({p.name})")
                        protein_model_status.update({"loaded": True, "error_message": None})
                        return wrapped
                    else:
                        # Raw nn.Module saved whole
                        wrapped = _CNNasXGB(sd, _FEAT)
                        protein_model_status.update({"loaded": True, "error_message": None})
                        return wrapped

                except Exception as e:
                    print(f"[PeVe]   torch fallback failed for {p.name}: {e}")

        raise FileNotFoundError("No loadable model file found in protein repo")

    except Exception:
        tb = traceback.format_exc()
        print(f"[PeVe] ✗ protein load failed:\n{tb}")
        protein_model_status.update({"loaded": False, "error_message": tb})
        return None


# ══════════════════════════════════════════════════════════════════════════════
# Public API  (names that app.py imports)
# ══════════════════════════════════════════════════════════════════════════════

def get_splice_model() -> tuple:
    """Returns (model, tokenizer). model(tensor) → (logit, imp, r_imp, s_imp)"""
    global _splice_model, _splice_tok
    if _splice_model is None:
        _splice_model, _splice_tok = _load_splice()
    return _splice_model, _splice_tok


def get_context_model() -> tuple:
    """Returns (model, tokenizer). model(tensor) → logit tensor"""
    global _context_model, _context_tok
    if _context_model is None:
        _context_model, _context_tok = _load_context()
    return _context_model, _context_tok


def get_protein_model():
    """Returns XGBoost Booster (or CNN wrapper). model.predict(dmat) → float array"""
    global _protein_model
    if _protein_model is None:
        _protein_model = _load_protein()
    return _protein_model


def get_model_status() -> dict:
    return {
        "splice":  dict(splice_model_status),
        "context": dict(context_model_status),
        "protein": dict(protein_model_status),
    }


════════════════════════════════════════════════════════════════════════
# Test block
# ══════════════════════════════════════════════════════════════════════════════

def test_model_loading() -> dict:
    import torch
    print("[PeVe] ── test_model_loading() ──")

    sm, _  = get_splice_model()
    cm, _  = get_context_model()
    pm     = get_protein_model()

    results = {}

    # Test splice
    try:
        if sm is not None:
            dummy = torch.zeros(1, 1106)
            out = sm(dummy)
            results["splice"] = f"✓ output shapes: {[o.shape for o in out]}"
        else:
            results["splice"] = "✗ model is None"
    except Exception as e:
        results["splice"] = f"✗ forward failed: {e}"

    # Test context
    try:
        if cm is not None:
            dummy = torch.zeros(1, 1106)
            out = cm(dummy)
            results["context"] = f"✓ output shape: {out.shape if hasattr(out,'shape') else type(out)}"
        else:
            results["context"] = "✗ model is None"
    except Exception as e:
        results["context"] = f"✗ forward failed: {e}"

    # Test protein
    try:
        if pm is not None:
            import xgboost as xgb
            feat = ["gnomAD_AF","Grantham","Charge_change",
                    "Hydrophobicity_diff","Protein_pos_norm","VEP_IMPACT"]
            X = np.array([[0.001, 100.0, 0.0, 0.5, 0.5, 2.0]], dtype=np.float32)
            dmat = xgb.DMatrix(X, feature_names=feat)
            pred = pm.predict(dmat)
            results["protein"] = f"✓ prediction: {pred}"
        else:
            results["protein"] = "✗ model is None"
    except Exception as e:
        results["protein"] = f"✗ predict failed: {e}"

    status = get_model_status()
    final = {
        "model_status": status,
        "forward_tests": results,
        "all_loaded": all(v["loaded"] for v in status.values()),
    }
    import json
    print(json.dumps(final, indent=2, default=str))
    return final


if __name__ == "__main__":
    test_model_loading()