File size: 2,914 Bytes
ee96220 | 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 | """Shared CLI and construction helpers for binding-affinity backends."""
from scoring.functions.binding import create_multi_target_affinity_predictor
def add_affinity_arguments(parser, default_backend="original"):
"""Add the common affinity-backend options to an argument parser."""
group = parser.add_argument_group("Binding Affinity")
group.add_argument(
"--affinity_backend",
choices=("original", "peptiverse"),
default=default_backend,
help="Affinity predictor to use; the original TD3B model remains the default.",
)
group.add_argument(
"--peptiverse_affinity_checkpoint",
default=None,
help=(
"Optional local PeptiVerse pooled SMILES-affinity checkpoint. "
"If omitted, it is downloaded from the Hub."
),
)
group.add_argument(
"--peptiverse_repo_id",
default=(
"ChatterjeeLab/PeptiVerse" if default_backend is not None else None
),
help="Hugging Face repository used to download the PeptiVerse checkpoint.",
)
group.add_argument(
"--peptiverse_revision",
default=None,
help="Optional Hugging Face revision for reproducible model downloads.",
)
group.add_argument(
"--peptiverse_cache_dir",
default=None,
help="Optional cache directory for PeptiVerse and encoder artifacts.",
)
group.add_argument(
"--peptiverse_local_files_only",
action="store_true",
default=False if default_backend is not None else None,
help="Require all PeptiVerse and encoder artifacts to exist locally.",
)
group.add_argument(
"--peptiverse_batch_size",
type=int,
default=32 if default_backend is not None else None,
help="Binder-SMILES embedding batch size for PeptiVerse scoring.",
)
return parser
def create_affinity_from_args(args, tokenizer, base_path, device, emb_model=None):
"""Build the selected multi-target affinity predictor from CLI/config values."""
return create_multi_target_affinity_predictor(
backend=getattr(args, "affinity_backend", None) or "original",
tokenizer=tokenizer,
base_path=base_path,
device=device,
emb_model=emb_model,
peptiverse_checkpoint=getattr(
args, "peptiverse_affinity_checkpoint", None
),
peptiverse_repo_id=(
getattr(args, "peptiverse_repo_id", None)
or "ChatterjeeLab/PeptiVerse"
),
peptiverse_revision=getattr(args, "peptiverse_revision", None),
peptiverse_cache_dir=getattr(args, "peptiverse_cache_dir", None),
peptiverse_local_files_only=bool(
getattr(args, "peptiverse_local_files_only", False)
),
peptiverse_batch_size=(
getattr(args, "peptiverse_batch_size", None) or 32
),
)
|