File size: 5,837 Bytes
4074358 adff1ee 4074358 adff1ee 4074358 adff1ee 4074358 adff1ee 4074358 adff1ee 4074358 adff1ee 4074358 adff1ee 4074358 adff1ee 4074358 adff1ee 4074358 adff1ee 4074358 adff1ee 4074358 adff1ee 4074358 adff1ee 4074358 adff1ee 4074358 adff1ee 4074358 adff1ee 4074358 adff1ee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | #!/usr/bin/env python3
from __future__ import annotations
from pathlib import Path
from typing import List, Optional, Tuple
from huggingface_hub import snapshot_download
from inference import (
PeptiVersePredictor,
read_best_manifest_csv,
_parse_model_and_emb,
EMB_TAG_TO_FOLDER_SUFFIX,
)
# -----------------------------
# Config
# -----------------------------
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"}
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, 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}"
folder_suffix = EMB_TAG_TO_FOLDER_SUFFIX.get(emb_tag or "wt", emb_tag or "wt")
if prop_key == "binding_affinity":
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 in {"xgb", "xgb_reg"}:
return f"{base}/xgb_wt_log"
# ------------------------------------------------------------------
# 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]:
best = read_best_manifest_csv(manifest_path)
allow: List[str] = []
for prop_key, row in best.items():
for col, parsed in [("wt", row.best_wt), ("smiles", row.best_smiles)]:
if parsed is None:
continue
model_name, emb_tag = parsed
# 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)
allow += [
f"{model_dir}/best_model.json",
f"{model_dir}/best_model.pt",
f"{model_dir}/best_model*.joblib",
f"{model_dir}/best_model*.json",
]
# Deduplicate while preserving order
seen = set()
out = []
for p in allow:
if p not in seen:
out.append(p)
seen.add(p)
return out
def download_assets(
repo_id: str,
manifest_path: Path,
out_dir: Path,
) -> Path:
out_dir = out_dir.resolve()
out_dir.mkdir(parents=True, exist_ok=True)
allow_patterns = build_allow_patterns_from_manifest(manifest_path)
snapshot_download(
repo_id=repo_id,
local_dir=str(out_dir),
local_dir_use_symlinks=False,
allow_patterns=allow_patterns,
)
return out_dir
# -----------------------------
# Main
# -----------------------------
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("--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}")
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)
)
"""TEST CODE
predictor = PeptiVersePredictor(
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() |