Stability Patch: Improved model loading and error handling

app.py

@@ -118,25 +118,35 @@ html, body, [class*="css"] {
 
 # ── Model + pipeline loading ──────────────────────────────────────────────────
 
-@st.cache_resource(show_spinner="Loading curriculum & model from HuggingFace...")
+@st.cache_resource(show_spinner="Connecting to Logic Engine...")
 def load_pipelines():
     from plrs.model.model_loader import load_model_from_hub
+    import os
+    
+    # Check for token in environment (HF Spaces allow setting secrets)
+    token = os.environ.get("HF_TOKEN")
     
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     maps = ROOT / "data" / "knowledge_maps"
 
-    # Load model (tries decay, vanilla, then base)
-    model, model_type = load_model_from_hub(device=str(device))
+    # Load model with potential token for private/restricted access
+    try:
+        model, model_type = load_model_from_hub(device=str(device), token=token)
+    except Exception as e:
+        model, model_type = None, f"Error: {str(e)}"
 
     pipelines = {}
-    for domain, fname in [("math", "math_dag.json"), ("cs", "cs_dag.json")]:
-        path = maps / fname
-        if path.exists():
-            curriculum = load_dag(path)
-            pipeline   = PLRSPipeline(curriculum)
-            if model:
-                pipeline._model = model
-            pipelines[domain] = pipeline
+    try:
+        for domain, fname in [("math", "math_dag.json"), ("cs", "cs_dag.json")]:
+            path = maps / fname
+            if path.exists():
+                curriculum = load_dag(path)
+                pipeline   = PLRSPipeline(curriculum)
+                if model:
+                    pipeline._model = model
+                pipelines[domain] = pipeline
+    except Exception as e:
+        st.error(f"Curriculum load error: {e}")
 
     return pipelines, model is not None, model_type
 

plrs/model/model_loader.py

@@ -24,30 +24,25 @@ import torch
 HF_REPO = "Clementio/PLRS"
 
 
-def load_model_from_hub(device: str = "cpu"):
+def load_model_from_hub(device: str = "cpu", token: str | None = None):
     """
     Load SAKT model weights from HuggingFace Hub.
-
-    Tries files in priority order:
-      1. sakt_decay_best.pt    (v0.2.0 — decay attention)
-      2. sakt_vanilla_best.pt  (v0.2.0 — vanilla transformer)
-      3. sakt_model.pt         (v0.1.0 — synthetic baseline)
-
-    Returns (model, model_type_str) or (None, "unavailable").
     """
     try:
         from huggingface_hub import hf_hub_download
     except ImportError:
         return None, "huggingface_hub not installed"
 
+    # Try files in priority order
     for filename, model_type in [
         ("models/sakt_decay_best.pt",   "SAKTWithDecay"),
         ("models/sakt_vanilla_best.pt", "SAKTModel"),
         ("models/sakt_model.pt",        "SAKTModel"),
+        ("sakt_model.pt",               "SAKTModel"), # Backwards compatibility
     ]:
         try:
-            path = hf_hub_download(repo_id=HF_REPO, filename=filename)
-            model = _load_weights(path, model_type, device)
+            path = hf_hub_download(repo_id=HF_REPO, filename=filename, token=token)
+            model = _load_weights(path, model_type, device, token=token)
             if model is not None:
                 return model, model_type
         except Exception:
@@ -56,8 +51,9 @@ def load_model_from_hub(device: str = "cpu"):
     return None, "unavailable"
 
 
-def _load_weights(path: str, preferred_type: str, device: str):
+def _load_weights(path: str, preferred_type: str, device: str, token: str | None = None):
     """Load model weights from a .pt file, handling both old and new formats."""
+    from huggingface_hub import hf_hub_download
     try:
         payload = torch.load(path, map_location=device, weights_only=False)
     except Exception:
@@ -65,27 +61,27 @@ def _load_weights(path: str, preferred_type: str, device: str):
 
     # ── New format (v0.2.0): {"state_dict": ..., "model_type": ..., "config": ...}
     if isinstance(payload, dict) and "state_dict" in payload:
-        cfg        = payload.get("config", {})
+        cfg = payload.get("config", {})
         model_type = payload.get("model_type", preferred_type)
 
         if model_type == "SAKTWithDecay":
             from plrs.model.sakt_decay import SAKTWithDecay
             model = SAKTWithDecay(
                 num_skills=cfg.get("num_skills", 5737),
-                embed_dim=cfg.get("embed_dim", 64),
+                embed_dim=cfg.get("embed_dim", 128),
                 num_heads=cfg.get("num_heads", 8),
                 dropout=cfg.get("dropout", 0.2),
-                max_seq_len=cfg.get("max_seq_len", 100),
+                max_seq_len=cfg.get("max_seq_len", 200),
                 decay_init=cfg.get("decay_init", 1.0),
             )
         else:
             from plrs.model.sakt import SAKTModel
             model = SAKTModel(
                 num_skills=cfg.get("num_skills", 5737),
-                embed_dim=cfg.get("embed_dim", 64),
+                embed_dim=cfg.get("embed_dim", 128),
                 num_heads=cfg.get("num_heads", 8),
                 dropout=cfg.get("dropout", 0.2),
-                max_seq_len=cfg.get("max_seq_len", 100),
+                max_seq_len=cfg.get("max_seq_len", 200),
             )
 
         try:
@@ -96,21 +92,26 @@ def _load_weights(path: str, preferred_type: str, device: str):
         except Exception:
             return None
 
-    # ── Old format (v0.1.0 FYP): raw state_dict + separate config.json
+    # ── Old format (v0.1.0 FYP): raw state_dict + fetch config.json from Hub
     try:
-        config_path = Path(path).parent / "config.json"
-        if config_path.exists():
-            config = json.loads(config_path.read_text())
-        else:
-            config = {"num_skills": 5736, "embed_dim": 64}
+        # Try to download config.json from the Hub root
+        try:
+            cfg_path = hf_hub_download(repo_id=HF_REPO, filename="config.json", token=token)
+            with open(cfg_path) as f:
+                config = json.load(f)
+        except Exception:
+            config = {"num_skills": 5737, "embed_dim": 128, "num_heads": 8, "num_layers": 2, "max_seq_len": 200, "dropout": 0.2}
 
         from plrs.model.sakt import SAKTModel
         model = SAKTModel(
-            num_skills=config.get("num_skills", 5736),
-            embed_dim=config.get("embed_dim", 64),
+            num_skills=config.get("num_skills", 5737),
+            embed_dim=config.get("embed_dim", 128),
+            num_heads=config.get("num_heads", 8),
+            max_seq_len=config.get("max_seq_len", 200),
         )
         model.load_state_dict(payload, strict=False)
         model.eval()
+        model.to(device)
         return model
     except Exception:
         return None