fix path

README.md

@@ -435,8 +435,8 @@ huggingface-cli download ChatterjeeLab/PeptiVerse \
     --local-dir . \
     --local-dir-use-symlinks False
 ```
-### TODOs
-Bug loading transformer half-life model now, will fix soon.
+### Trouble installing cuML
+For error related to cuda library, reinstall the `torch` after installing `cuML`.
 
 ## Citation
 

inference.py

@@ -1,4 +1,3 @@
-# peptiverse_infer.py
 from __future__ import annotations
 
 import csv, re, json
@@ -14,7 +13,8 @@ import xgboost as xgb
 
 from transformers import EsmModel, EsmTokenizer, AutoModelForMaskedLM
 from tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
-
+from lightning.pytorch import seed_everything
+seed_everything(1986)
 
 # -----------------------------
 # Manifest
@@ -138,7 +138,7 @@ def load_artifact(model_dir: Path, device: torch.device) -> Tuple[str, Any, Path
 
     if art.suffix == ".json":
         booster = xgb.Booster()
-        print(str(art))
+        #print(str(art))
         booster.load_model(str(art))
         return "xgb", booster, art
 
@@ -226,6 +226,41 @@ def _infer_in_dim_from_sd(sd: dict, model_name: str) -> int:
         return int(sd["proj.weight"].shape[1])
     raise ValueError(model_name)
 
+def _infer_num_layers_from_sd(sd: dict, prefix: str = "enc.layers.") -> int:
+    # enc.layers.0.*, enc.layers.1.*, ...
+    idxs = set()
+    for k in sd.keys():
+        if k.startswith(prefix):
+            rest = k[len(prefix):]
+            m = re.match(r"(\d+)\.", rest)
+            if m:
+                idxs.add(int(m.group(1)))
+    return (max(idxs) + 1) if idxs else 1
+
+def _infer_transformer_arch_from_sd(sd: dict) -> Tuple[int, int, int]:
+    """
+    Returns (d_model, layers, ff) inferred from weights.
+    - d_model from proj.weight (shape: [d_model, in_dim])
+    - layers from count of enc.layers.*
+    - ff from enc.layers.0.linear1.weight (shape: [ff, d_model])
+    """
+    if "proj.weight" not in sd:
+        raise KeyError("Missing proj.weight in state_dict; cannot infer transformer d_model.")
+    d_model = int(sd["proj.weight"].shape[0])
+    layers = _infer_num_layers_from_sd(sd, prefix="enc.layers.")
+    if "enc.layers.0.linear1.weight" in sd:
+        ff = int(sd["enc.layers.0.linear1.weight"].shape[0])
+    else:
+        ff = 4 * d_model
+    return d_model, layers, ff
+
+def _pick_nhead(d_model: int) -> int:
+    # prefer common head counts; must divide d_model
+    for h in (8, 6, 4, 3, 2, 1):
+        if d_model % h == 0:
+            return h
+    return 1
+
 def build_torch_model_from_ckpt(model_name: str, ckpt: dict, device: torch.device) -> nn.Module:
     params = ckpt["best_params"]
     sd = ckpt["state_dict"]
@@ -238,25 +273,30 @@ def build_torch_model_from_ckpt(model_name: str, ckpt: dict, device: torch.devic
         model = CNNHead(in_ch=in_dim, c=int(params["channels"]), k=int(params["kernel"]),
                         layers=int(params["layers"]), dropout=dropout)
     elif model_name == "transformer":
-        d_model = (
-            params.get("d_model")
-            or params.get("hidden")
-            or params.get("hidden_dim")
-        )
+        # if transfer-learning ckpt omits arch params, infer from state_dict. special case for transformer_wt_log
+        d_model = params.get("d_model") or params.get("hidden") or params.get("hidden_dim")
+
         if d_model is None:
-            raise KeyError(
-                f"Transformer checkpoint missing d_model/hidden. "
-                f"Available keys: {list(params.keys())}"
+            d_model_i, layers_i, ff_i = _infer_transformer_arch_from_sd(sd)
+            nhead_i = _pick_nhead(d_model_i)
+            model = TransformerHead(
+                in_dim=in_dim,
+                d_model=int(d_model_i),
+                nhead=int(params.get("nhead", nhead_i)),
+                layers=int(params.get("layers", layers_i)),
+                ff=int(params.get("ff", ff_i)),
+                dropout=float(params.get("dropout", dropout)),
+            )
+        else:
+            d_model = int(d_model)
+            model = TransformerHead(
+                in_dim=in_dim,
+                d_model=d_model,
+                nhead=int(params.get("nhead", _pick_nhead(d_model))),
+                layers=int(params.get("layers", 2)),
+                ff=int(params.get("ff", 4 * d_model)),
+                dropout=dropout
             )
-
-        model = TransformerHead(
-            in_dim=in_dim,
-            d_model=int(d_model),
-            nhead=int(params["nhead"]),
-            layers=int(params["layers"]),
-            ff=int(params.get("ff", 4 * int(d_model))),
-            dropout=dropout
-        )
     else:
         raise ValueError(f"Unknown NN model_name={model_name}")
 
@@ -678,6 +718,12 @@ class PeptiVersePredictor:
             if d.exists():
                 return d
 
+        # special handling for halflife transformer wt log folder
+        if prop_key == "halflife" and mode == "wt" and model_name == "transformer":
+            d = base / "transformer_wt_log"
+            if d.exists():
+                return d
+        
         if prop_key == "halflife" and model_name == "xgb":
             d = base / ("xgb_wt_log" if mode == "wt" else "xgb_smiles")
             if d.exists():
@@ -920,7 +966,7 @@ class PeptiVersePredictor:
 if __name__ == "__main__":
     predictor = PeptiVersePredictor(
        manifest_path="best_models.txt",
-       classifier_weight_root="./Classifier_Weight"
+       classifier_weight_root="./"
      )
     print(predictor.predict_property("hemolysis", "wt", "GIGAVLKVLTTGLPALISWIKRKRQQ"))
     print(predictor.predict_binding_affinity("wt", target_seq="...", binder_str="..."))