scoring/affinity_config.py

"""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
        ),
    )