Update model_loader.py

model_loader.py

@@ -1,443 +1,558 @@
 """
-model_loader.py — PeVe v1.3
-===========================
-ARCHITECTURE CHANGE (correct approach):
-  Models are NOT loaded as weights. They are already running as live
-  Gradio Spaces on HuggingFace. We query them via their HTTP /predict
-  (or /run/predict) Gradio API endpoints.
-
-  Space URLs:
-    splice  : https://nileshhanotia-mutation-predictor-splice-app.hf.space
-    context : https://nileshhanotia-mutation-predictor-v4-demo.hf.space
-    protein : https://nileshhanotia-mutation-pathogenicity-app.hf.space
-
-  The model_loader exposes three functions:
-    query_splice_model(ref_seq, mut_seq, chrom, pos, exon_flag, intron_flag)
-    query_context_model(chrom, pos, ref, alt)
-    query_protein_model(feature_vector: list[float])
-
-  Each returns a dict with:
-    {
-      "loaded": True/False,
-      "error_message": None or str,
-      ... model-specific output keys ...
-    }
+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 time
+import pickle
 import traceback
+import warnings
+from pathlib import Path
 from typing import Any
 
-import requests
-
-# ── Space base URLs ────────────────────────────────────────────────────────────
-# These are the standard HF Spaces inference URLs (Gradio 3/4 compatible)
-_SPLICE_SPACE   = os.environ.get(
-    "SPLICE_SPACE_URL",
-    "https://nileshhanotia-mutation-predictor-splice-app.hf.space"
-)
-_CONTEXT_SPACE  = os.environ.get(
-    "CONTEXT_SPACE_URL",
-    "https://nileshhanotia-mutation-predictor-v4-demo.hf.space"
-)
-_PROTEIN_SPACE  = os.environ.get(
-    "PROTEIN_SPACE_URL",
-    "https://nileshhanotia-mutation-pathogenicity-app.hf.space"
-)
+import numpy as np
 
+# ── HF token ──────────────────────────────────────────────────────────────────
 _HF_TOKEN: str | None = (
     os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
 )
 
-_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", "30"))   # seconds per call
-_MAX_RETRIES = 2
+# ── Model repo IDs ─────────────────────────────────────────────────────────────
+_REPO_SPLICE   = "nileshhanotia/mutation-predictor-splice"
+_REPO_CONTEXT  = "nileshhanotia/mutation-predictor-v4"
+_REPO_PROTEIN  = "nileshhanotia/mutation-pathogenicity-predictor"
 
-# ── Status dicts (updated on every call) ──────────────────────────────────────
-splice_model_status:  dict = {"loaded": False, "error_message": "Not yet called"}
-context_model_status: dict = {"loaded": False, "error_message": "Not yet called"}
-protein_model_status: dict = {"loaded": False, "error_message": "Not yet called"}
+# ── 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}
 
-# ── Low-level HTTP helper ──────────────────────────────────────────────────────
 
-def _gradio_predict(space_url: str, payload: dict, status: dict,
-                    label: str) -> dict | None:
-    """
-    POST to a Gradio /run/predict endpoint.
-    Gradio 3 payload shape: {"data": [...]}
-    Returns the parsed response dict, or None on failure.
-    Retries up to _MAX_RETRIES times with exponential back-off.
-    Updates `status` in place.
-    """
-    endpoint = f"{space_url.rstrip('/')}/run/predict"
-    headers = {"Content-Type": "application/json"}
-    if _HF_TOKEN:
-        headers["Authorization"] = f"Bearer {_HF_TOKEN}"
+# ══════════════════════════════════════════════════════════════════════════════
+# Model Architecture definitions
+# (must match checkpoint shapes exactly)
+# ══════════════════════════════════════════════════════════════════════════════
 
-    last_err = ""
-    for attempt in range(1, _MAX_RETRIES + 1):
-        try:
-            print(f"[PeVe] {label} → POST {endpoint} (attempt {attempt})")
-            resp = requests.post(
-                endpoint,
-                json=payload,
-                headers=headers,
-                timeout=_TIMEOUT,
-            )
-            resp.raise_for_status()
-            data = resp.json()
-            print(f"[PeVe] {label} ← HTTP 200, keys={list(data.keys())}")
-            status["loaded"] = True
-            status["error_message"] = None
-            return data
-
-        except requests.exceptions.HTTPError as e:
-            last_err = f"HTTP {resp.status_code}: {resp.text[:300]}"
-            print(f"[PeVe] {label} attempt {attempt} HTTPError: {last_err}")
-        except requests.exceptions.Timeout:
-            last_err = f"Timeout after {_TIMEOUT}s"
-            print(f"[PeVe] {label} attempt {attempt} Timeout")
-        except Exception:
-            last_err = traceback.format_exc()
-            print(f"[PeVe] {label} attempt {attempt} error:\n{last_err}")
-
-        if attempt < _MAX_RETRIES:
-            time.sleep(2 ** attempt)   # 2s, 4s …
-
-    status["loaded"] = False
-    status["error_message"] = last_err
-    return None
+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)
 
-def _wake_space(space_url: str, label: str) -> bool:
+    app.py passes tensors of shape (1, 401, 8).  We need an adapter.
     """
-    HF Spaces sleep after inactivity. A GET to / wakes them.
-    Returns True if the Space responds with 200.
+    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):
     """
-    try:
-        r = requests.get(space_url, timeout=15,
-                         headers={"Authorization": f"Bearer {_HF_TOKEN}"} if _HF_TOKEN else {})
-        print(f"[PeVe] wake {label}: HTTP {r.status_code}")
-        return r.status_code == 200
-    except Exception as e:
-        print(f"[PeVe] wake {label} failed: {e}")
-        return False
-
-
-# ── Splice Space ───────────────────────────────────────────────────────────────
-# Space: mutation-predictor-splice-app
-# Model: MutationPredictorCNN_v2
-# Gradio fn: predict_variant_with_causal_analysis(ref_seq, mut_seq, chrom, pos,
-#                                                  exon_flag, intron_flag)
-# Inputs: 6 values → data: [str, str, str, int, int, int]
-# Outputs: [report_str, act_fig, splice_fig, percentile_str,
-#           cf_text, cf_fig, grad_fig, ablation_fig, ablation_text]
-
-def query_splice_model(
-    ref_seq: str,
-    mut_seq: str,
-    chrom: str = "17",
-    pos: int = 0,
-    exon_flag: int = 1,
-    intron_flag: int = 0,
-) -> dict:
+    MutationPredictorCNN_v4 — 4-input forward (seq, mut, region, splice).
+    app.py passes a single (B, 401, 8) tensor, so forward() adapts.
     """
-    Query the splice Space and return structured outputs.
-
-    Returns
-    -------
-    {
-        "loaded": bool,
-        "error_message": str | None,
-        "splice_prob": float,          # sigmoid of logit from report text
-        "splice_signal_strength": float,
-        "counterfactual_delta": float,
-        "raw_report": str,
-    }
+    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:
     """
-    _wake_space(_SPLICE_SPACE, "splice")
+    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)
 
-    payload = {
-        "data": [ref_seq, mut_seq, str(chrom), int(pos),
-                 int(exon_flag), int(intron_flag)],
-        "fn_index": 0,       # predict_variant_with_causal_analysis is fn 0
-    }
+        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
 
-    result = _gradio_predict(_SPLICE_SPACE, payload, splice_model_status, "splice")
-    if result is None:
-        return {
-            "loaded": False,
-            "error_message": splice_model_status["error_message"],
-            "splice_prob": 0.0,
-            "splice_signal_strength": 0.0,
-            "counterfactual_delta": 0.0,
-            "raw_report": "",
-        }
-
-    # Gradio returns {"data": [output0, output1, ...]}
-    outputs = result.get("data", [])
-    report_str  = outputs[0] if len(outputs) > 0 else ""
-    cf_text     = outputs[4] if len(outputs) > 4 else ""
-    ablation_text = outputs[8] if len(outputs) > 8 else ""
-
-    # Parse probability from the text report (format: "Probability : 0.XXXX")
-    splice_prob = _parse_float_from_text(report_str, "Probability")
-    cf_delta    = _parse_float_from_text(cf_text, "Causal importance score")
+    # Allow shap.TreeExplainer to fail gracefully — app.py has try/except
+    def get_booster(self):
+        raise NotImplementedError("CNN wrapped as XGB — SHAP not available")
 
-    return {
-        "loaded": True,
-        "error_message": None,
-        "splice_prob": splice_prob,
-        "splice_signal_strength": splice_prob,   # direct proxy
-        "counterfactual_delta": cf_delta,
-        "raw_report": report_str,
-    }
 
+# ══════════════════════════════════════════════════════════════════════════════
+# Internal loaders
+# ══════════════════════════════════════════════════════════════════════════════
 
-# ── Context Space ──────────────────────────────────────────────────────────────
-# Space: mutation-predictor-v4-demo
-# Model: MutationPredictorCNN_v4  (BUT currently runs predict_dummy → random float)
-# Gradio fn: predict_dummy(chr, pos, ref, alt)
-# Inputs: [str, int, str, str]
-# Outputs: [markdown_str]
-#
-# NOTE: This Space currently returns a dummy probability. The model IS loaded
-# (torch.load on "mutation_predictor_splice_v4.pt") but predict_dummy bypasses it.
-# Until the Space is updated to call the real model, we extract the dummy float
-# and surface it with a warning.
-
-def query_context_model(
-    chrom: str,
-    pos: int,
-    ref: str,
-    alt: str,
-) -> dict:
-    """
-    Query the context/v4 Space.
-
-    Returns
-    -------
-    {
-        "loaded": bool,
-        "error_message": str | None,
-        "context_pathogenic_prob": float,
-        "activation_norm": float,
-        "activation_peak_position": int,
-        "importance_score": float,
-        "warning": str | None,
-    }
-    """
-    _wake_space(_CONTEXT_SPACE, "context")
+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
 
-    payload = {
-        "data": [str(chrom), int(pos), str(ref), str(alt)],
-        "fn_index": 0,
-    }
 
-    result = _gradio_predict(_CONTEXT_SPACE, payload, context_model_status, "context")
-    if result is None:
-        return {
-            "loaded": False,
-            "error_message": context_model_status["error_message"],
-            "context_pathogenic_prob": 0.0,
-            "activation_norm": 0.0,
-            "activation_peak_position": 200,
-            "importance_score": 0.0,
-            "warning": None,
-        }
+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}")
 
-    outputs = result.get("data", [])
-    text_out = outputs[0] if outputs else ""
+        ckpt = torch.load(str(candidates[0]), map_location="cpu", weights_only=False)
+        sd   = ckpt.get("model_state_dict", ckpt)
 
-    # Parse "Predicted probability: X.XXX" from markdown output
-    prob = _parse_float_from_text(text_out, "Predicted probability")
+        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
 
-    return {
-        "loaded": True,
-        "error_message": None,
-        "context_pathogenic_prob": prob,
-        "activation_norm": prob,          # best proxy available
-        "activation_peak_position": 200,  # not exposed by this Space yet
-        "importance_score": prob,
-        "warning": (
-            "context Space is running predict_dummy (placeholder). "
-            "Probability is random. Update Space to call real model."
-            if "placeholder" in text_out.lower() else 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
 
 
-# ── Protein / Pathogenicity Space ──────────────────────────────────────────────
-# Space: mutation-pathogenicity-app
-# Model: MutationPredictorCNN (model.py — CNN with input dim 1101)
-# We need to know the Gradio fn signature. Since we couldn't fetch app.py,
-# we attempt /info first to discover the API, then call /run/predict.
-
-def query_protein_model(
-    ref_seq: str,
-    mut_seq: str,
-    chrom: str = "17",
-    pos: int = 0,
-    exon_flag: int = 1,
-    intron_flag: int = 0,
-) -> dict:
-    """
-    Query the protein/pathogenicity Space.
-
-    Returns
-    -------
-    {
-        "loaded": bool,
-        "error_message": str | None,
-        "biochemical_risk_score": float,
-        "feature_pathogenic_prob": float,
-        "shap_feature_contributions": dict,
-    }
-    """
-    _wake_space(_PROTEIN_SPACE, "protein")
-
-    # Try to discover API shape first
-    api_info = _get_space_api_info(_PROTEIN_SPACE)
-    fn_index = 0
-
-    # MutationPredictorCNN takes input dim 1101 — if the Space accepts
-    # raw sequences, send them; fall back to numeric if it expects a vector.
-    payload = {
-        "data": [ref_seq, mut_seq, str(chrom), int(pos),
-                 int(exon_flag), int(intron_flag)],
-        "fn_index": fn_index,
-    }
+def _load_context() -> tuple:
+    import torch
+    print(f"[PeVe] Loading context model from {_REPO_CONTEXT}")
+    try:
+        local = _download_repo(_REPO_CONTEXT)
 
-    result = _gradio_predict(_PROTEIN_SPACE, payload, protein_model_status, "protein")
-    if result is None:
-        # Second attempt: maybe the Space expects only chr/pos/ref/alt
-        payload2 = {
-            "data": [str(chrom), int(pos), ref_seq[0] if ref_seq else "N",
-                     mut_seq[0] if mut_seq else "N"],
-            "fn_index": fn_index,
-        }
-        result = _gradio_predict(_PROTEIN_SPACE, payload2, protein_model_status,
-                                 "protein-fallback")
-
-    if result is None:
-        return {
-            "loaded": False,
-            "error_message": protein_model_status["error_message"],
-            "biochemical_risk_score": 0.0,
-            "feature_pathogenic_prob": 0.0,
-            "shap_feature_contributions": {},
-        }
-
-    outputs = result.get("data", [])
-    text_out = str(outputs[0]) if outputs else ""
-
-    prob = _parse_float_from_text(text_out, "probability") or \
-           _parse_float_from_text(text_out, "Probability") or \
-           _parse_float_from_text(text_out, "score")
+        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}")
 
-    return {
-        "loaded": True,
-        "error_message": None,
-        "biochemical_risk_score": prob,
-        "feature_pathogenic_prob": prob,
-        "shap_feature_contributions": {},   # not exposed by this Space yet
-    }
+        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
 
-# ── Utility helpers ────────────────────────────────────────────────────────────
+    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 _parse_float_from_text(text: str, keyword: str) -> float:
-    """
-    Extract a float after 'keyword' in text.
-    e.g. "Probability : 0.8743" → 0.8743
-    """
-    import re
-    pattern = re.compile(
-        rf"{re.escape(keyword)}\s*[:\|=]?\s*([0-9]+\.?[0-9]*)",
-        re.IGNORECASE,
-    )
-    m = pattern.search(text)
-    if m:
-        try:
-            return float(m.group(1))
-        except ValueError:
-            pass
-    return 0.0
 
+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"]
 
-def _get_space_api_info(space_url: str) -> dict:
-    """Fetch /info endpoint to understand Space API (Gradio 3.x+)."""
     try:
-        r = requests.get(
-            f"{space_url.rstrip('/')}/info",
-            timeout=10,
-            headers={"Authorization": f"Bearer {_HF_TOKEN}"} if _HF_TOKEN else {},
-        )
-        if r.status_code == 200:
-            return r.json()
+        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:
-        pass
-    return {}
+        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
 
-# ── Status accessors ───────────────────────────────────────────────────────────
 
 def get_model_status() -> dict:
-    """Return last-call status for all three models."""
     return {
-        "splice":  splice_model_status,
-        "context": context_model_status,
-        "protein": protein_model_status,
+        "splice":  dict(splice_model_status),
+        "context": dict(context_model_status),
+        "protein": dict(protein_model_status),
     }
 
 
-# ── Test block ─────────────────────────────────────────────────────────────────
+# ══════════════════════════════════════════════════════════════════════════════
+# Test block
+# ══════════════════════════════════════════════════════════════════════════════
 
 def test_model_loading() -> dict:
-    """
-    Smoke-test all three Spaces with a real BRCA1 variant window.
-    chr17:43092176 G>T (the variant from your JSON example).
-    """
+    import torch
     print("[PeVe] ── test_model_loading() ──")
 
-    # 99-bp window centred on BRCA1 position 43092176 (synthetic for test)
-    ref_seq = ("ATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCG"
-               "GCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGATCGAT")[:99]
-    mut_seq = (ref_seq[:49] + "T" + ref_seq[50:])[:99]   # G→T at centre
-
-    r_splice  = query_splice_model(ref_seq, mut_seq, "17", 43092176, 1, 0)
-    r_context = query_context_model("17", 43092176, "G", "T")
-    r_protein = query_protein_model(ref_seq, mut_seq, "17", 43092176, 1, 0)
-
-    result = {
-        "model_status": {
-            "splice":  {"loaded": r_splice["loaded"],
-                        "error_message": r_splice["error_message"]},
-            "context": {"loaded": r_context["loaded"],
-                        "error_message": r_context["error_message"]},
-            "protein": {"loaded": r_protein["loaded"],
-                        "error_message": r_protein["error_message"]},
-        },
-        "sample_outputs": {
-            "splice_prob":              r_splice.get("splice_prob"),
-            "counterfactual_delta":     r_splice.get("counterfactual_delta"),
-            "context_pathogenic_prob":  r_context.get("context_pathogenic_prob"),
-            "context_warning":          r_context.get("warning"),
-            "biochemical_risk_score":   r_protein.get("biochemical_risk_score"),
-        },
-        "all_loaded": all([r_splice["loaded"],
-                           r_context["loaded"],
-                           r_protein["loaded"]]),
-    }
+    sm, _  = get_splice_model()
+    cm, _  = get_context_model()
+    pm     = get_protein_model()
 
-    print(f"[PeVe] Result: {result}")
-    return result
+    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}"
 
-if __name__ == "__main__":
+    # 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
-    report = test_model_loading()
-    print(json.dumps(report, indent=2))
-    
+    print(json.dumps(final, indent=2, default=str))
+    return final
+
+
+if __name__ == "__main__":
+    test_model_loading()