Add PeptiVerse affinity backend

README.md

@@ -179,6 +179,24 @@ Key flags:
 
 Output: prints the valid yield (valid / `num_samples`) with per-round counts, and writes the valid sequences to `--save_path` (default `results/valid_<strategy>_len<L>.csv`) with columns idx, sequence, n_chars.
 
+### Binding-affinity backends
+
+The original TD3B affinity predictor remains the default. To use PeptiVerse's
+pooled target-sequence/binder-SMILES regressor instead, add:
+
+```bash
+python inference.py \
+    --ckpt_path checkpoints/td3b.ckpt \
+    --val_csv data/test.csv \
+    --affinity_backend peptiverse
+```
+
+The PeptiVerse checkpoint and its ESM-2/ChemBERTa encoders are downloaded from
+Hugging Face on first use. For offline runs, provide
+`--peptiverse_affinity_checkpoint /path/to/best_model.pt` together with
+`--peptiverse_local_files_only`. Training supports the same flags; alternatively,
+set `AFFINITY_BACKEND="peptiverse"` in `launch_multi_target.sh`.
+
 ## Training
 
 ### Multi-target TD3B

baselines/run_mcts_tr2d2.py

@@ -27,7 +27,8 @@ from configs.finetune_config import (
 )
 from training.finetune_utils import load_tokenizer
 from training.distributed_utils import setup_distributed, cleanup_distributed, is_main_process
-from scoring.functions.binding import MultiTargetBindingAffinity, TargetSpecificBindingAffinity
+from scoring.affinity_config import add_affinity_arguments, create_affinity_from_args
+from scoring.functions.binding import TargetSpecificBindingAffinity
 from td3b.direction_oracle import DirectionalOracle
 from finetune_multi_target_tr2d2_ddp import TR2D2GatedReward, TargetDataset, create_tr2d2_mcts
 from utils.app import PeptideAnalyzer
@@ -106,6 +107,18 @@ def _build_args(cfg: Dict[str, Any], cli: argparse.Namespace) -> argparse.Namesp
         merged["exploration"] = cli.exploration
     if cli.max_sequence_length is not None:
         merged["max_sequence_length"] = cli.max_sequence_length
+    for name in (
+        "affinity_backend",
+        "peptiverse_affinity_checkpoint",
+        "peptiverse_repo_id",
+        "peptiverse_revision",
+        "peptiverse_cache_dir",
+        "peptiverse_local_files_only",
+        "peptiverse_batch_size",
+    ):
+        value = getattr(cli, name, None)
+        if value is not None:
+            merged[name] = value
 
     args = SimpleNamespace(**merged)
 
@@ -216,6 +229,7 @@ def _nanstd(values: np.ndarray) -> float:
 
 def main() -> None:
     parser = argparse.ArgumentParser(description="MCTS-based TR2-D2 evaluation.")
+    add_affinity_arguments(parser, default_backend=None)
     parser.add_argument("--ckpt_path", required=True, help="Path to finetuned checkpoint (.ckpt)")
     parser.add_argument("--val_csv", required=True, help="Validation CSV path")
     parser.add_argument("--device", default="cuda", help="Device string (e.g., cuda:0 or cpu)")
@@ -258,7 +272,8 @@ def main() -> None:
 
     policy_model = _build_model(args, payload["state_dict"], device)
 
-    multi_target_affinity = MultiTargetBindingAffinity(
+    multi_target_affinity = create_affinity_from_args(
+        args=args,
         tokenizer=tokenizer,
         base_path=args.base_path,
         device=device,

baselines/run_validation_td3b.py

@@ -28,7 +28,8 @@ from configs.finetune_config import (
 from training.finetune_utils import load_tokenizer, create_reward_function
 from finetune_multi_target import TargetDataset
 from training.distributed_utils import setup_distributed, cleanup_distributed, is_main_process
-from scoring.functions.binding import MultiTargetBindingAffinity, TargetSpecificBindingAffinity
+from scoring.affinity_config import add_affinity_arguments, create_affinity_from_args
+from scoring.functions.binding import TargetSpecificBindingAffinity
 from td3b.direction_oracle import DirectionalOracle
 from utils.app import PeptideAnalyzer
 
@@ -94,6 +95,18 @@ def _build_args(cfg: Dict[str, Any], cli: argparse.Namespace) -> argparse.Namesp
         merged["sampling_eps"] = cli.sampling_eps
     if cli.seed is not None:
         merged["seed"] = cli.seed
+    for name in (
+        "affinity_backend",
+        "peptiverse_affinity_checkpoint",
+        "peptiverse_repo_id",
+        "peptiverse_revision",
+        "peptiverse_cache_dir",
+        "peptiverse_local_files_only",
+        "peptiverse_batch_size",
+    ):
+        value = getattr(cli, name, None)
+        if value is not None:
+            merged[name] = value
 
     args = SimpleNamespace(**merged)
 
@@ -270,6 +283,7 @@ def _nanstd(values: np.ndarray) -> float:
 
 def main() -> None:
     parser = argparse.ArgumentParser(description="Run TD3B validation from a saved checkpoint.")
+    add_affinity_arguments(parser, default_backend=None)
     parser.add_argument("--ckpt_path", required=True, help="Path to saved checkpoint (.ckpt)")
     parser.add_argument("--val_csv", required=True, help="Validation CSV path")
     parser.add_argument("--device", default="cuda", help="Device string (e.g., cuda:0 or cpu)")
@@ -312,7 +326,8 @@ def main() -> None:
 
     policy_model = _build_model(args, payload["state_dict"], device)
 
-    multi_target_affinity = MultiTargetBindingAffinity(
+    multi_target_affinity = create_affinity_from_args(
+        args=args,
         tokenizer=tokenizer,
         base_path=args.base_path,
         device=device,

baselines/sampling_setup.py

@@ -15,8 +15,8 @@ from hydra import compose, initialize_config_dir
 from hydra.core.global_hydra import GlobalHydra
 
 from models.diffusion import Diffusion
+from scoring.affinity_config import add_affinity_arguments, create_affinity_from_args
 from scoring.scoring_functions import ScoringFunctions
-from scoring.functions.binding import MultiTargetBindingAffinity
 from td3b.direction_oracle import DirectionalOracle, resolve_device
 from td3b.data_utils import peptide_seq_to_smiles, smiles_token_length
 
@@ -104,11 +104,14 @@ def load_reward_models(
     base_path: Optional[str] = None,
     multi_target: bool = False,
     score_func_names: Optional[List[str]] = None,
+    args=None,
 ):
-    if multi_target:
+    affinity_backend = getattr(args, "affinity_backend", "original")
+    if multi_target or affinity_backend == "peptiverse":
         if base_model is None or base_path is None:
-            raise ValueError("base_model and base_path are required for multi-target affinity.")
-        return MultiTargetBindingAffinity(
+            raise ValueError("base_model and base_path are required for affinity scoring.")
+        return create_affinity_from_args(
+            args=args,
             tokenizer=base_model.tokenizer,
             base_path=base_path,
             device=device,
@@ -222,6 +225,7 @@ def run_baseline(
 
 def main():
     parser = argparse.ArgumentParser()
+    add_affinity_arguments(parser)
     parser.add_argument("--ckpt_path", type=str, required=True)
     parser.add_argument("--device", type=str, default="cuda:0")
     parser.add_argument("--baseline", type=str, default="cg", choices=["cg", "smc", "tds", "unguided", "peptune"])
@@ -308,6 +312,7 @@ def main():
         base_model=base_model,
         base_path=base_path,
         multi_target=multi_target,
+        args=args,
     )
     direction_oracle = load_direction_oracle(args, args.device)
     reward_alpha = args.reward_alpha if args.reward_alpha is not None else args.alpha

env.yml

@@ -23,6 +23,7 @@ dependencies:
       - lightning==2.5.5
       - fair-esm==2.0.0
       - transformers==4.56.2
+      - huggingface-hub
       - SmilesPE==0.0.3
       - scipy==1.13.1
       - wandb==0.22.0
@@ -34,4 +35,4 @@ dependencies:
       - seaborn==0.13.2
       - timm==1.0.20
       - xgboost==3.0.5
-      - loguru==0.7.3
+      - loguru==0.7.3

finetune_multi_target.py

@@ -36,7 +36,8 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 from models.diffusion import Diffusion
 from tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
 from utils.app import PeptideAnalyzer
-from scoring.functions.binding import MultiTargetBindingAffinity, TargetSpecificBindingAffinity
+from scoring.affinity_config import add_affinity_arguments, create_affinity_from_args
+from scoring.functions.binding import TargetSpecificBindingAffinity
 from td3b.data_utils import peptide_seq_to_smiles, smiles_token_length
 
 # TD3B imports
@@ -193,7 +194,7 @@ class TargetDataset:
 
 def run_validation(
     policy_model: Diffusion,
-    multi_target_affinity: MultiTargetBindingAffinity,
+    multi_target_affinity,
     directional_oracle: DirectionalOracle,
     tokenizer: SMILES_SPE_Tokenizer,
     val_dataset: TargetDataset,
@@ -394,6 +395,7 @@ def run_validation(
 def parse_args():
     """Parse command-line arguments."""
     parser = argparse.ArgumentParser(description='Multi-Target TD3B Fine-Tuning')
+    add_affinity_arguments(parser)
 
     # Paths
     path_group = parser.add_argument_group('Paths')
@@ -681,13 +683,14 @@ def main():
     policy_model = add_td3b_sampling_to_model(policy_model) 
 
     # Multi-target affinity predictor
-    multi_target_affinity = MultiTargetBindingAffinity(
+    multi_target_affinity = create_affinity_from_args(
+        args=args,
         tokenizer=tokenizer,
         base_path=args.base_path,
         device=device,
-        emb_model=policy_model.backbone  # Use backbone Roformer model (matches v1)
+        emb_model=policy_model.backbone,
     )
-    logger.info("Created multi-target binding affinity predictor")
+    logger.info("Created %s binding affinity predictor", args.affinity_backend)
 
     # Directional oracle (GPCR classifier)
     for path_label, path in [

inference.py

@@ -29,10 +29,11 @@ from configs.finetune_config import (
     DiffusionConfig, RoFormerConfig, NoiseConfig,
     TrainingConfig, SamplingConfig, EvalConfig, OptimConfig, MCTSConfig,
 )
+from scoring.affinity_config import add_affinity_arguments, create_affinity_from_args
+from scoring.functions.binding import TargetSpecificBindingAffinity
 from training.finetune_utils import load_tokenizer
 from td3b.direction_oracle import DirectionalOracle
 from td3b.td3b_scoring import create_td3b_reward_function
-from scoring.functions.binding import MultiTargetBindingAffinity, TargetSpecificBindingAffinity
 from td3b.data_utils import peptide_seq_to_smiles, smiles_token_length
 from utils.app import PeptideAnalyzer
 import sampling_strategies
@@ -138,6 +139,7 @@ def score_sequences(reward_model, sequences: List[str]):
 
 def main():
     parser = argparse.ArgumentParser(description="TD3B Inference")
+    add_affinity_arguments(parser)
     parser.add_argument("--ckpt_path", type=str, required=True, help="Path to TD3B checkpoint")
     parser.add_argument("--val_csv", type=str, required=True, help="CSV with Target_Sequence, Ligand_Sequence, label columns")
     parser.add_argument("--save_path", type=str, default="results", help="Output directory")
@@ -185,7 +187,6 @@ def main():
     # Load model
     logger.info(f"Loading model from {args.ckpt_path}")
     model, tokenizer = load_model(args.ckpt_path, device)
-
     # Load targets
     logger.info(f"Loading targets from {args.val_csv}")
     df = pd.read_csv(args.val_csv)
@@ -219,12 +220,14 @@ def main():
     vocab_path = os.path.join(ROOT_DIR, "tokenizer", "new_vocab.txt")
     splits_path = os.path.join(ROOT_DIR, "tokenizer", "new_splits.txt")
 
-    multi_affinity = MultiTargetBindingAffinity(
+    multi_affinity = create_affinity_from_args(
+        args=args,
         tokenizer=tokenizer,
         base_path=ROOT_DIR,
         device=device,
         emb_model=model.backbone,
     )
+    logger.info("Using %s binding affinity predictor", args.affinity_backend)
     directional_oracle = DirectionalOracle(
         model_ckpt=oracle_ckpt,
         tr2d2_checkpoint=oracle_tr2d2,

launch_multi_target.sh

@@ -17,6 +17,12 @@ VAL_CSV="${BASE_PATH}/data/test.csv"  # Optional: create validation split
 # Run configuration
 RUN_NAME="multi_target_td3b"  # Timestamp will be added automatically
 DEVICE="cuda:0"
+
+# Binding-affinity backend ("original" preserves the published TD3B behavior)
+AFFINITY_BACKEND="original"  # original or peptiverse
+PEPTIVERSE_AFFINITY_CHECKPOINT=""  # Optional local best_model.pt
+PEPTIVERSE_CACHE_DIR=""             # Optional Hugging Face cache directory
+PEPTIVERSE_LOCAL_FILES_ONLY=0       # Set to 1 for offline-only loading
 # Multi-target sampling
 TARGETS_PER_MCTS=2        # Number of targets sampled per MCTS round (K)
 RESAMPLE_TARGETS_EVERY=1  # Resample targets every N epochs
@@ -70,6 +76,17 @@ EXTRA_ORACLE_ARGS=""
 if [ -n "$ORACLE_ESM_CACHE_DIR" ]; then
     EXTRA_ORACLE_ARGS="$EXTRA_ORACLE_ARGS --direction_oracle_esm_cache_dir $ORACLE_ESM_CACHE_DIR"
 fi
+
+EXTRA_AFFINITY_ARGS=""
+if [ -n "$PEPTIVERSE_AFFINITY_CHECKPOINT" ]; then
+    EXTRA_AFFINITY_ARGS="$EXTRA_AFFINITY_ARGS --peptiverse_affinity_checkpoint $PEPTIVERSE_AFFINITY_CHECKPOINT"
+fi
+if [ -n "$PEPTIVERSE_CACHE_DIR" ]; then
+    EXTRA_AFFINITY_ARGS="$EXTRA_AFFINITY_ARGS --peptiverse_cache_dir $PEPTIVERSE_CACHE_DIR"
+fi
+if [ "$PEPTIVERSE_LOCAL_FILES_ONLY" -eq 1 ]; then
+    EXTRA_AFFINITY_ARGS="$EXTRA_AFFINITY_ARGS --peptiverse_local_files_only"
+fi
 if [ "$ORACLE_ESM_LOCAL_FILES_ONLY" -eq 1 ]; then
     EXTRA_ORACLE_ARGS="$EXTRA_ORACLE_ARGS --direction_oracle_esm_local_files_only"
 fi
@@ -107,6 +124,8 @@ CMD="python finetune_multi_target.py \
     --pretrained_checkpoint ${PRETRAINED_CHECKPOINT} \
     --run_name ${RUN_NAME} \
     --device ${DEVICE} \
+    --affinity_backend ${AFFINITY_BACKEND} \
+    ${EXTRA_AFFINITY_ARGS} \
     \
     --targets_per_mcts ${TARGETS_PER_MCTS} \
     --resample_targets_every ${RESAMPLE_TARGETS_EVERY} \

scoring/affinity_config.py

@@ -0,0 +1,78 @@
+"""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
+        ),
+    )

scoring/functions/binding.py

@@ -1,5 +1,6 @@
 import sys
 import os, torch
+from pathlib import Path
 import numpy as np
 import torch
 import pandas as pd
@@ -7,6 +8,15 @@ import torch.nn as nn
 import esm
 from transformers import AutoModelForMaskedLM
 
+from scoring.functions.peptiverse_binding import PeptiVerseBindingAffinity
+
+
+def _resolve_project_root(base_path):
+    base_path = Path(base_path)
+    if (base_path / "scoring" / "functions" / "classifiers").is_dir():
+        return base_path
+    return base_path / "tr2d2-pep"
+
 
 def _sanitize_token_ids(input_ids: torch.Tensor, vocab_size: int, unk_id: int) -> torch.Tensor:
     if vocab_size <= 0 or input_ids.numel() == 0:
@@ -154,7 +164,8 @@ class BindingAffinity:
             self.max_pep_len = getattr(self.pep_model.config, "max_position_embeddings", None)
 
         self.model = ImprovedBindingPredictor().to(self.device)
-        checkpoint = torch.load(f'{base_path}/scoring/functions/classifiers/binding-affinity.pt',
+        project_root = _resolve_project_root(base_path)
+        checkpoint = torch.load(project_root / 'scoring/functions/classifiers/binding-affinity.pt',
                                 map_location=self.device,
                                 weights_only=False)
         _bind_sd = checkpoint.get('model_state_dict', checkpoint.get('state_dict', checkpoint)) \
@@ -270,7 +281,8 @@ class MultiTargetBindingAffinity:
 
         # Binding affinity prediction model
         self.model = ImprovedBindingPredictor().to(self.device)
-        checkpoint = torch.load(f'{base_path}/scoring/functions/classifiers/binding-affinity.pt',
+        project_root = _resolve_project_root(base_path)
+        checkpoint = torch.load(project_root / 'scoring/functions/classifiers/binding-affinity.pt',
                                 map_location=self.device,
                                 weights_only=False)
         _bind_sd = checkpoint.get('model_state_dict', checkpoint.get('state_dict', checkpoint)) \
@@ -450,7 +462,7 @@ class TargetSpecificBindingAffinity:
     where only peptide sequences need to be provided.
     """
 
-    def __init__(self, multi_target_predictor: MultiTargetBindingAffinity, prot_seq: str):
+    def __init__(self, multi_target_predictor, prot_seq: str):
         """
         Create a target-specific binding affinity predictor.
 
@@ -484,3 +496,42 @@ class TargetSpecificBindingAffinity:
             List of binding affinity scores
         """
         return self.forward(input_seqs)
+
+
+def create_multi_target_affinity_predictor(
+    backend="original",
+    tokenizer=None,
+    base_path=None,
+    device=None,
+    emb_model=None,
+    peptiverse_checkpoint=None,
+    peptiverse_repo_id="ChatterjeeLab/PeptiVerse",
+    peptiverse_revision=None,
+    peptiverse_cache_dir=None,
+    peptiverse_local_files_only=False,
+    peptiverse_batch_size=32,
+):
+    """Create the original TD3B or PeptiVerse affinity backend."""
+    backend = str(backend).lower()
+    if backend == "original":
+        if tokenizer is None or base_path is None:
+            raise ValueError("The original affinity backend requires tokenizer and base_path.")
+        return MultiTargetBindingAffinity(
+            tokenizer=tokenizer,
+            base_path=base_path,
+            device=device,
+            emb_model=emb_model,
+        )
+    if backend == "peptiverse":
+        return PeptiVerseBindingAffinity(
+            device=device,
+            checkpoint_path=peptiverse_checkpoint,
+            repo_id=peptiverse_repo_id,
+            revision=peptiverse_revision,
+            cache_dir=peptiverse_cache_dir,
+            local_files_only=peptiverse_local_files_only,
+            batch_size=peptiverse_batch_size,
+        )
+    raise ValueError(
+        f"Unknown affinity backend {backend!r}; choose 'original' or 'peptiverse'."
+    )

scoring/functions/peptiverse_binding.py

@@ -0,0 +1,274 @@
+"""PeptiVerse binding-affinity adapter for TD3B.
+
+This module implements the pooled target-sequence/binder-SMILES model published
+in ChatterjeeLab/PeptiVerse without loading PeptiVerse's unrelated predictors.
+"""
+
+import logging
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_REPO_ID = "ChatterjeeLab/PeptiVerse"
+DEFAULT_CHECKPOINT_FILE = (
+    "training_classifiers/binding_affinity/"
+    "chemberta_smiles_pooled/best_model.pt"
+)
+DEFAULT_ESM_MODEL = "facebook/esm2_t33_650M_UR50D"
+DEFAULT_CHEMBERTA_MODEL = "DeepChem/ChemBERTa-77M-MLM"
+
+
+class PeptiVersePooledAffinityModel(nn.Module):
+    """PeptiVerse's bidirectional cross-attention affinity head."""
+
+    def __init__(
+        self,
+        target_dim: int,
+        binder_dim: int,
+        hidden_dim: int,
+        n_heads: int,
+        n_layers: int,
+        dropout: float,
+    ) -> None:
+        super().__init__()
+        self.t_proj = nn.Sequential(
+            nn.Linear(target_dim, hidden_dim), nn.LayerNorm(hidden_dim)
+        )
+        self.b_proj = nn.Sequential(
+            nn.Linear(binder_dim, hidden_dim), nn.LayerNorm(hidden_dim)
+        )
+        self.layers = nn.ModuleList()
+        for _ in range(n_layers):
+            self.layers.append(
+                nn.ModuleDict(
+                    {
+                        "attn_tb": nn.MultiheadAttention(
+                            hidden_dim, n_heads, dropout=dropout
+                        ),
+                        "attn_bt": nn.MultiheadAttention(
+                            hidden_dim, n_heads, dropout=dropout
+                        ),
+                        "n1t": nn.LayerNorm(hidden_dim),
+                        "n2t": nn.LayerNorm(hidden_dim),
+                        "n1b": nn.LayerNorm(hidden_dim),
+                        "n2b": nn.LayerNorm(hidden_dim),
+                        "fft": nn.Sequential(
+                            nn.Linear(hidden_dim, 4 * hidden_dim),
+                            nn.GELU(),
+                            nn.Dropout(dropout),
+                            nn.Linear(4 * hidden_dim, hidden_dim),
+                        ),
+                        "ffb": nn.Sequential(
+                            nn.Linear(hidden_dim, 4 * hidden_dim),
+                            nn.GELU(),
+                            nn.Dropout(dropout),
+                            nn.Linear(4 * hidden_dim, hidden_dim),
+                        ),
+                    }
+                )
+            )
+
+        self.shared = nn.Sequential(
+            nn.Linear(2 * hidden_dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+        )
+        self.reg = nn.Linear(hidden_dim, 1)
+        self.cls = nn.Linear(hidden_dim, 3)
+
+    def forward(self, target: torch.Tensor, binder: torch.Tensor):
+        target = self.t_proj(target).unsqueeze(0)
+        binder = self.b_proj(binder).unsqueeze(0)
+        for layer in self.layers:
+            target_attn, _ = layer["attn_tb"](target, binder, binder)
+            target = layer["n1t"]((target + target_attn).transpose(0, 1)).transpose(0, 1)
+            target = layer["n2t"](
+                (target + layer["fft"](target)).transpose(0, 1)
+            ).transpose(0, 1)
+
+            binder_attn, _ = layer["attn_bt"](binder, target, target)
+            binder = layer["n1b"]((binder + binder_attn).transpose(0, 1)).transpose(0, 1)
+            binder = layer["n2b"](
+                (binder + layer["ffb"](binder)).transpose(0, 1)
+            ).transpose(0, 1)
+
+        hidden = self.shared(torch.cat([target[0], binder[0]], dim=-1))
+        return self.reg(hidden).squeeze(-1), self.cls(hidden)
+
+
+class PeptiVerseBindingAffinity:
+    """Score target/peptide-SMILES pairs with PeptiVerse's pK regressor."""
+
+    backend_name = "peptiverse"
+
+    def __init__(
+        self,
+        device=None,
+        checkpoint_path: Optional[str] = None,
+        repo_id: str = DEFAULT_REPO_ID,
+        revision: Optional[str] = None,
+        cache_dir: Optional[str] = None,
+        local_files_only: bool = False,
+        esm_name: str = DEFAULT_ESM_MODEL,
+        chemberta_name: str = DEFAULT_CHEMBERTA_MODEL,
+        batch_size: int = 32,
+        max_protein_length: int = 1022,
+        max_smiles_length: int = 512,
+    ) -> None:
+        from transformers import AutoModel, AutoTokenizer, EsmModel, EsmTokenizer
+
+        self.device = torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu"
+        ) if device is None else torch.device(device)
+        self.batch_size = max(1, int(batch_size))
+        self.max_protein_length = max_protein_length
+        self.max_smiles_length = max_smiles_length
+
+        resolved_checkpoint = self._resolve_checkpoint(
+            checkpoint_path=checkpoint_path,
+            repo_id=repo_id,
+            revision=revision,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+        checkpoint = torch.load(
+            resolved_checkpoint, map_location=self.device, weights_only=False
+        )
+        if checkpoint.get("mode") != "pooled":
+            raise ValueError(
+                f"Expected a pooled PeptiVerse checkpoint, got {checkpoint.get('mode')!r}"
+            )
+
+        state_dict = checkpoint["state_dict"]
+        params = checkpoint.get("best_params", {})
+        model = PeptiVersePooledAffinityModel(
+            target_dim=int(state_dict["t_proj.0.weight"].shape[1]),
+            binder_dim=int(state_dict["b_proj.0.weight"].shape[1]),
+            hidden_dim=int(params.get("hidden_dim", state_dict["t_proj.0.weight"].shape[0])),
+            n_heads=int(params.get("n_heads", 4)),
+            n_layers=int(params.get("n_layers", self._infer_layers(state_dict))),
+            dropout=float(params.get("dropout", 0.0)),
+        )
+        model.load_state_dict(state_dict, strict=True)
+        self.model = model.to(self.device).eval()
+
+        model_kwargs = {
+            "cache_dir": cache_dir,
+            "local_files_only": local_files_only,
+        }
+        self.target_tokenizer = EsmTokenizer.from_pretrained(esm_name, **model_kwargs)
+        self.target_encoder = EsmModel.from_pretrained(
+            esm_name, add_pooling_layer=False, **model_kwargs
+        ).to(self.device).eval()
+        self.binder_tokenizer = AutoTokenizer.from_pretrained(
+            chemberta_name, **model_kwargs
+        )
+        self.binder_encoder = AutoModel.from_pretrained(
+            chemberta_name, **model_kwargs
+        ).to(self.device).eval()
+        self.target_cache: Dict[str, torch.Tensor] = {}
+        logger.info("Loaded PeptiVerse affinity checkpoint: %s", resolved_checkpoint)
+
+    @staticmethod
+    def _resolve_checkpoint(
+        checkpoint_path: Optional[str],
+        repo_id: str,
+        revision: Optional[str],
+        cache_dir: Optional[str],
+        local_files_only: bool,
+    ) -> str:
+        if checkpoint_path is not None:
+            path = Path(checkpoint_path).expanduser()
+            if not path.is_file():
+                raise FileNotFoundError(f"PeptiVerse checkpoint not found: {path}")
+            return str(path)
+
+        from huggingface_hub import hf_hub_download
+
+        return hf_hub_download(
+            repo_id=repo_id,
+            filename=DEFAULT_CHECKPOINT_FILE,
+            revision=revision,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+
+    @staticmethod
+    def _infer_layers(state_dict) -> int:
+        layer_ids = {
+            int(key.split(".")[1])
+            for key in state_dict
+            if key.startswith("layers.")
+        }
+        return max(layer_ids) + 1
+
+    @staticmethod
+    def _special_token_ids(tokenizer) -> List[int]:
+        values = [
+            getattr(tokenizer, f"{name}_token_id", None)
+            for name in ("pad", "cls", "sep", "bos", "eos", "mask")
+        ]
+        return sorted({int(value) for value in values if value is not None})
+
+    @torch.no_grad()
+    def _pool(self, texts, tokenizer, encoder, max_length: int) -> torch.Tensor:
+        tokens = tokenizer(
+            list(texts),
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=max_length,
+        )
+        tokens = {name: value.to(self.device) for name, value in tokens.items()}
+        attention_mask = tokens.get(
+            "attention_mask", torch.ones_like(tokens["input_ids"])
+        ).bool()
+        valid_mask = attention_mask
+        special_ids = self._special_token_ids(tokenizer)
+        if special_ids:
+            special = torch.tensor(special_ids, device=self.device)
+            valid_mask = valid_mask & ~torch.isin(tokens["input_ids"], special)
+        hidden = encoder(**tokens).last_hidden_state
+        weights = valid_mask.unsqueeze(-1).to(hidden.dtype)
+        return (hidden * weights).sum(dim=1) / weights.sum(dim=1).clamp_min(1.0)
+
+    def get_protein_embedding(self, prot_seq: str) -> torch.Tensor:
+        prot_seq = prot_seq.strip()
+        if prot_seq not in self.target_cache:
+            self.target_cache[prot_seq] = self._pool(
+                [prot_seq],
+                self.target_tokenizer,
+                self.target_encoder,
+                self.max_protein_length,
+            )
+        return self.target_cache[prot_seq]
+
+    @torch.no_grad()
+    def forward(self, input_seqs, prot_seq: str):
+        input_seqs = list(input_seqs)
+        if not input_seqs:
+            return []
+
+        target = self.get_protein_embedding(prot_seq)
+        scores = []
+        for start in range(0, len(input_seqs), self.batch_size):
+            batch = input_seqs[start:start + self.batch_size]
+            binder = self._pool(
+                batch,
+                self.binder_tokenizer,
+                self.binder_encoder,
+                self.max_smiles_length,
+            )
+            affinity, _ = self.model(target.expand(len(batch), -1), binder)
+            scores.extend(affinity.detach().cpu().tolist())
+        return scores
+
+    def __call__(self, input_seqs, prot_seq: str):
+        return self.forward(input_seqs, prot_seq)
+
+    def clear_cache(self) -> None:
+        self.target_cache.clear()

td3b/__init__.py

@@ -7,7 +7,6 @@ from .direction_oracle import DirectionalOracle
 from .td3b_scoring import TD3BRewardFunction, TD3BConfidenceWeighting, create_td3b_reward_function
 from .td3b_losses import ContrastiveLoss, InfoNCELoss, TD3BTotalLoss, extract_embeddings_from_mdlm
 from .td3b_mcts import TD3B_MCTS, create_td3b_mcts
-from .td3b_finetune import td3b_finetune, add_td3b_sampling_to_model
 from .data_utils import TD3BDataset, load_td3b_data
 
 __all__ = [
@@ -28,3 +27,15 @@ __all__ = [
 ]
 
 __version__ = '0.1.0'
+
+
+def __getattr__(name):
+    """Load fine-tuning helpers lazily to avoid a package import cycle."""
+    if name in {'td3b_finetune', 'add_td3b_sampling_to_model'}:
+        from .td3b_finetune import add_td3b_sampling_to_model, td3b_finetune
+
+        return {
+            'td3b_finetune': td3b_finetune,
+            'add_td3b_sampling_to_model': add_td3b_sampling_to_model,
+        }[name]
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

tests/test_peptiverse_binding.py

@@ -0,0 +1,124 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from scoring.functions.peptiverse_binding import (
+    PeptiVerseBindingAffinity,
+    PeptiVersePooledAffinityModel,
+)
+
+
+class DummyTokenizer:
+    pad_token_id = 0
+    cls_token_id = 1
+    eos_token_id = 2
+    sep_token_id = None
+    bos_token_id = None
+    mask_token_id = None
+
+    def __call__(self, texts, **kwargs):
+        del texts, kwargs
+        return {
+            "input_ids": torch.tensor([[1, 3, 4, 2, 0]]),
+            "attention_mask": torch.tensor([[1, 1, 1, 1, 0]]),
+        }
+
+
+class DummyEncoder(torch.nn.Module):
+    def forward(self, input_ids, attention_mask):
+        del attention_mask
+        hidden = input_ids.float().unsqueeze(-1).repeat(1, 1, 2)
+        return SimpleNamespace(last_hidden_state=hidden)
+
+
+class DummyAffinityHead(torch.nn.Module):
+    def forward(self, target, binder):
+        del target
+        return binder.sum(dim=-1), torch.zeros(len(binder), 3)
+
+
+class PeptiVerseBindingTests(unittest.TestCase):
+    def test_affinity_head_shapes(self):
+        model = PeptiVersePooledAffinityModel(
+            target_dim=8,
+            binder_dim=6,
+            hidden_dim=12,
+            n_heads=3,
+            n_layers=2,
+            dropout=0.0,
+        ).eval()
+        affinity, classes = model(torch.randn(4, 8), torch.randn(4, 6))
+        self.assertEqual(tuple(affinity.shape), (4,))
+        self.assertEqual(tuple(classes.shape), (4, 3))
+
+    def test_pool_excludes_special_tokens(self):
+        predictor = object.__new__(PeptiVerseBindingAffinity)
+        predictor.device = torch.device("cpu")
+        pooled = predictor._pool(
+            ["unused"], DummyTokenizer(), DummyEncoder(), max_length=8
+        )
+        expected = torch.tensor([[3.5, 3.5]])
+        self.assertTrue(torch.equal(pooled, expected))
+
+    def test_factory_keeps_original_as_default(self):
+        from scoring.functions import binding
+
+        original = object()
+        with patch.object(
+            binding, "MultiTargetBindingAffinity", return_value=original
+        ) as constructor:
+            result = binding.create_multi_target_affinity_predictor(
+                tokenizer=object(), base_path="/tmp", device="cpu"
+            )
+        self.assertIs(result, original)
+        constructor.assert_called_once()
+
+    def test_factory_selects_peptiverse(self):
+        from scoring.functions import binding
+
+        peptiverse = object()
+        with patch.object(
+            binding, "PeptiVerseBindingAffinity", return_value=peptiverse
+        ) as constructor:
+            result = binding.create_multi_target_affinity_predictor(
+                backend="peptiverse",
+                device="cpu",
+                peptiverse_checkpoint="model.pt",
+            )
+        self.assertIs(result, peptiverse)
+        constructor.assert_called_once_with(
+            device="cpu",
+            checkpoint_path="model.pt",
+            repo_id="ChatterjeeLab/PeptiVerse",
+            revision=None,
+            cache_dir=None,
+            local_files_only=False,
+            batch_size=32,
+        )
+
+    def test_forward_batches_binder_smiles(self):
+        predictor = object.__new__(PeptiVerseBindingAffinity)
+        predictor.batch_size = 2
+        predictor.binder_tokenizer = object()
+        predictor.binder_encoder = object()
+        predictor.max_smiles_length = 16
+        predictor.model = DummyAffinityHead()
+        predictor.get_protein_embedding = lambda _: torch.zeros(1, 3)
+
+        batch_sizes = []
+
+        def fake_pool(texts, tokenizer, encoder, max_length):
+            del tokenizer, encoder, max_length
+            batch_sizes.append(len(texts))
+            return torch.ones(len(texts), 2)
+
+        predictor._pool = fake_pool
+        scores = predictor.forward(["a", "b", "c", "d", "e"], "TARGET")
+        self.assertEqual(batch_sizes, [2, 2, 1])
+        self.assertEqual(scores, [2.0] * 5)
+
+
+if __name__ == "__main__":
+    unittest.main()

training/finetune_utils.py

@@ -192,8 +192,13 @@ def load_tokenizer(base_path: str) -> SMILES_SPE_Tokenizer:
         >>> tokenizer = load_tokenizer('To Be Added')
     """
     base_path = Path(base_path)
-    vocab_path = base_path / "tokenizer" / "new_vocab.txt"
-    spe_path = base_path / "tokenizer" / "new_splits.txt"
+    project_root = (
+        base_path
+        if (base_path / "tokenizer").is_dir()
+        else base_path / "tr2d2-pep"
+    )
+    vocab_path = project_root / "tokenizer" / "new_vocab.txt"
+    spe_path = project_root / "tokenizer" / "new_splits.txt"
 
     if not vocab_path.exists():
         raise FileNotFoundError(f"Vocabulary file not found: {vocab_path}")