tested light download

download_light.py

@@ -1,23 +1,23 @@
 #!/usr/bin/env python3
 from __future__ import annotations
 
-import os
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+from typing import List, Optional, Tuple
 
 from huggingface_hub import snapshot_download
 from inference import (
     PeptiVersePredictor,
     read_best_manifest_csv,
-    canon_model,
+    _parse_model_and_emb,
+    EMB_TAG_TO_FOLDER_SUFFIX,
 )
 
 # -----------------------------
 # Config
 # -----------------------------
-root = Path(__file__).resolve().parent  # current script folder
-MODEL_REPO = "ChatterjeeLab/PeptiVerse"  
-DEFAULT_ASSETS_DIR = Path(root)   # where downloaded models live
+root = Path(__file__).resolve().parent
+MODEL_REPO = "ChatterjeeLab/PeptiVerse"
+DEFAULT_ASSETS_DIR = Path(root)
 DEFAULT_MANIFEST = Path("./basic_models.txt")
 
 BANNED_MODELS = {"svm", "enet", "svm_gpu", "enet_gpu"}
@@ -26,25 +26,41 @@ BANNED_MODELS = {"svm", "enet", "svm_gpu", "enet_gpu"}
 def _norm_prop_disk(prop_key: str) -> str:
     return "half_life" if prop_key == "halflife" else prop_key
 
-def _resolve_expected_model_dir(prop_key: str, model_name: str, mode: str) -> str:
+
+def _resolve_expected_model_dir(
+    prop_key: str, model_name: str, emb_tag: Optional[str]
+) -> str:
+    """
+    Resolve the subfolder path inside training_classifiers/<property>/.
+
+    Args:
+        prop_key:   manifest property key, e.g. 'hemolysis', 'halflife'.
+        model_name: canonical model string, e.g. 'xgb', 'cnn', 'transformer'.
+        emb_tag:    embedding tag from manifest: 'wt', 'peptideclm', or 'chemberta'.
+                    None means WT (falls back to 'wt').
+    """
     disk_prop = _norm_prop_disk(prop_key)
     base = f"training_classifiers/{disk_prop}"
 
-    # binding affinity is special: its label is pooled/unpooled and folder uses wt_<mode>_<pooled|unpooled>
+    folder_suffix = EMB_TAG_TO_FOLDER_SUFFIX.get(emb_tag or "wt", emb_tag or "wt")
+
     if prop_key == "binding_affinity":
-        pooled_or_unpooled = model_name  # "pooled" or "unpooled"
-        return f"{base}/wt_{mode}_{pooled_or_unpooled}"
-
-    # halflife special folders
-    if prop_key == "halflife":
-        if model_name in {"xgb_wt_log", "xgb_smiles"}:
-            return f"{base}/{model_name}"
-        if mode == "wt" and model_name == "transformer":
+        return f"{base}/{model_name}"
+
+    # ------------------------------------------------------------------
+    # halflife WT: folders carry a _log suffix
+    # ------------------------------------------------------------------
+    if prop_key == "halflife" and (emb_tag is None or emb_tag == "wt"):
+        if model_name == "transformer":
             return f"{base}/transformer_wt_log"
-        if model_name == "xgb":
-            return f"{base}/{'xgb_wt_log' if mode == 'wt' else 'xgb_smiles'}"
+        if model_name in {"xgb", "xgb_reg"}:
+            return f"{base}/xgb_wt_log"
 
-    return f"{base}/{model_name}_{mode}"
+    # ------------------------------------------------------------------
+    # Default: <model>_<folder_suffix>
+    # e.g. cnn_chemberta, xgb_peptideclm, transformer_wt
+    # ------------------------------------------------------------------
+    return f"{base}/{model_name}_{folder_suffix}"
 
 
 def build_allow_patterns_from_manifest(manifest_path: Path) -> List[str]:
@@ -52,19 +68,19 @@ def build_allow_patterns_from_manifest(manifest_path: Path) -> List[str]:
 
     allow: List[str] = []
 
-    # For each property, fetch best artifacts for wt + smiles
     for prop_key, row in best.items():
-        for mode, label in [("wt", row.best_wt), ("smiles", row.best_smiles)]:
-            m = canon_model(label)
-            if m is None:
+        for col, parsed in [("wt", row.best_wt), ("smiles", row.best_smiles)]:
+            if parsed is None:
                 continue
 
-            if m in BANNED_MODELS:
-                m = "xgb"
+            model_name, emb_tag = parsed
 
-            model_dir = _resolve_expected_model_dir(prop_key, m, mode)
+            # Replace banned models with xgb
+            if model_name in BANNED_MODELS:
+                model_name = "xgb"
+
+            model_dir = _resolve_expected_model_dir(prop_key, model_name, emb_tag)
 
-            # fetch only "basic" artifacts, not everything in the folder
             allow += [
                 f"{model_dir}/best_model.json",
                 f"{model_dir}/best_model.pt",
@@ -72,6 +88,7 @@ def build_allow_patterns_from_manifest(manifest_path: Path) -> List[str]:
                 f"{model_dir}/best_model*.json",
             ]
 
+    # Deduplicate while preserving order
     seen = set()
     out = []
     for p in allow:
@@ -106,41 +123,51 @@ def download_assets(
 def main():
     import argparse
 
-    ap = argparse.ArgumentParser(description="Lightweight PeptiVerse inference with on-demand model download.")
-    ap.add_argument("--repo", default=MODEL_REPO, help="HF repo id containing weights/assets.")
-    ap.add_argument("--manifest", default=str(DEFAULT_MANIFEST), help="Path to best_models.txt")
-    ap.add_argument("--assets", default=str(DEFAULT_ASSETS_DIR), help="Where to store downloaded assets")
-    ap.add_argument("--device", default=None, help="cuda / cpu / cuda:0, etc")
-
-    ap.add_argument("--property", default="hemolysis", help="Property key (e.g. hemolysis, solubility, ...)")
-    ap.add_argument("--mode", default="wt", choices=["wt", "smiles"], help="Input type: wt=AA sequence, smiles=SMILES")
-    ap.add_argument("--input", default="GIGAVLKVLTTGLPALISWIKRKRQQ", help="Sequence or SMILES string")
-    ap.add_argument("--target_seq", default=None, help="Target WT sequence for binding_affinity")
-    ap.add_argument("--binder", default=None, help="Binder string (AA or SMILES) for binding_affinity")
+    ap = argparse.ArgumentParser(
+        description="Lightweight PeptiVerse inference with on-demand model download."
+    )
+    ap.add_argument("--repo",     default=MODEL_REPO,              help="HF repo id containing weights/assets.")
+    ap.add_argument("--manifest", default=str(DEFAULT_MANIFEST),   help="Path to best_models.txt")
+    ap.add_argument("--assets",   default=str(DEFAULT_ASSETS_DIR), help="Where to store downloaded assets")
+    ap.add_argument("--device",   default=None,                    help="cuda / cpu / cuda:0, etc")
+    ap.add_argument("--dry-run",  action="store_true",             help="Print allow-patterns without downloading")
+
+    ap.add_argument("--property",   default="hemolysis")
+    ap.add_argument("--mode",       default="wt", choices=["wt", "smiles"])
+    ap.add_argument("--input",      default="GIGAVLKVLTTGLPALISWIKRKRQQ")
+    ap.add_argument("--target_seq", default="GIGAVLKVLTTGLPALISWIKRKRQQ")
+    ap.add_argument("--binder",     default="GIGAVLKV")
     args = ap.parse_args()
 
     manifest_path = Path(args.manifest)
     if not manifest_path.exists():
         raise FileNotFoundError(f"Manifest not found: {manifest_path}")
 
-    assets_dir = download_assets(args.repo, manifest_path=manifest_path, out_dir=Path(args.assets))
+    if args.dry_run:
+        patterns = build_allow_patterns_from_manifest(manifest_path)
+        print(f"Would download {len(patterns)} patterns:")
+        for p in patterns:
+            print(" ", p)
+        return
+
+    assets_dir = download_assets(
+        args.repo, manifest_path=manifest_path, out_dir=Path(args.assets)
+    )
 
-    """ OPTIONAL TEST CODE
+    """TEST CODE
     predictor = PeptiVersePredictor(
-        manifest_path="basic_models.txt",  # use the downloaded copy to be consistent
+        manifest_path=manifest_path,
         classifier_weight_root=str(assets_dir),
         device=args.device,
     )
-
     if args.property == "binding_affinity":
         if not args.target_seq or not args.binder:
             raise ValueError("For binding_affinity, provide --target_seq and --binder.")
         out = predictor.predict_binding_affinity(args.mode, target_seq=args.target_seq, binder_str=args.binder)
     else:
         out = predictor.predict_property(args.property, args.mode, args.input)
-
     print(out)
     """
 
 if __name__ == "__main__":
-    main()
+    main()

inference.py

@@ -19,13 +19,7 @@ seed_everything(1986)
 
 EMB_TAG_TO_FOLDER_SUFFIX = {
     "wt":         "wt",
-    "peptideclm": "smiles",
-    "chemberta":  "chemberta",
-}
-
-EMB_TAG_TO_RUNTIME_MODE = {
-    "wt":         "wt",
-    "peptideclm": "smiles",
+    "peptideclm": "peptideclm",
     "chemberta":  "chemberta",
 }
 
@@ -778,7 +772,7 @@ class PeptiVersePredictor:
 
                 # infer emb_tag
                 if emb_tag is None:
-                    emb_tag = col
+                    emb_tag = "wt"
 
                 model_dir = self._resolve_dir(prop_key, model_name, emb_tag)
                 kind, obj, art = load_artifact(model_dir, self.device)
@@ -1054,7 +1048,7 @@ if __name__ == "__main__":
     root = Path(__file__).resolve().parent  # current script folder
 
     predictor = PeptiVersePredictor(
-        manifest_path=root / "best_models.txt",
+        manifest_path=root / "basic_models.txt",
         classifier_weight_root=root
     )
     print(predictor.training_root)
@@ -1065,17 +1059,21 @@ if __name__ == "__main__":
     smiles = "C(C)C[C@@H]1NC(=O)[C@@H]2CCCN2C(=O)[C@@H](CC(C)C)NC(=O)[C@@H](CC(C)C)N(C)C(=O)[C@H](C)NC(=O)[C@H](Cc2ccccc2)NC1=O"
 
     print(predictor.predict_property("hemolysis",   "wt",     seq))
-    print(predictor.predict_property("hemolysis",   "smiles",     smiles, uncertainty=True))
-    print(predictor.predict_property("nf",    "wt",     seq, uncertainty=True))
-    print(predictor.predict_property("nf",    "smiles",     smiles, uncertainty=True))
+    print(predictor.predict_property("hemolysis",   "smiles",     smiles, uncertainty=False))
+    print(predictor.predict_property("nf",    "wt",     seq, uncertainty=False))
+    print(predictor.predict_property("nf",    "smiles",     smiles, uncertainty=False))
     print(predictor.predict_binding_affinity("wt",  target_seq=seq, binder_str="GIGAVLKVLT"))
-    print(predictor.predict_binding_affinity("wt",  target_seq=seq, binder_str="GIGAVLKVLT", uncertainty=True))
+    print(predictor.predict_binding_affinity("wt",  target_seq=seq, binder_str="GIGAVLKVLT", uncertainty=False))
     seq1 = "GIGAVLKVLTTGLPALISWIKRKRQQ"
     seq2 = "ACDEFGHIKLMNPQRSTVWY"
     
-    r1 = predictor.predict_binding_affinity("wt",  target_seq=seq2, binder_str="GIGAVLKVLT", uncertainty=True)
-    r2 = predictor.predict_property("nf", "wt", seq1, uncertainty=True)
-    r3 = predictor.predict_property("nf", "wt", seq2, uncertainty=True)
+    r1 = predictor.predict_binding_affinity("wt",  target_seq=seq2, binder_str="GIGAVLKVLT", uncertainty=False)
+    r2 = predictor.predict_property("nf", "wt", seq1, uncertainty=False)
+    r3 = predictor.predict_property("nf", "wt", seq2, uncertainty=False)
+    r4 = predictor.predict_binding_affinity("wt",  target_seq=seq2, binder_str=smiles, uncertainty=False)
+    r5 = predictor.predict_property("halflife",  "smiles", smiles)
     print(r1)
     print(r2)
     print(r3)
+    print(r4)
+    print(r5)