cas9/objectives.py

# import warnings
# import logging
# warnings.filterwarnings("ignore", message="to-Python converter.*already registered", category=RuntimeWarning)
# warnings.filterwarnings("ignore", message="pkg_resources is deprecated", category=UserWarning)
# warnings.filterwarnings("ignore", category=FutureWarning)
# warnings.filterwarnings("ignore", category=DeprecationWarning)
# warnings.filterwarnings("ignore", category=UserWarning)
# from rdkit import RDLogger, rdBase
# RDLogger.DisableLog("rdApp.*")
# logging.basicConfig(level=logging.ERROR)
# logging.getLogger("chemprop").setLevel(logging.ERROR)
# logging.getLogger("hyperopt").setLevel(logging.ERROR)
# rdBase.DisableLog('rdApp.error')

import torch
import torch.nn.functional as F
import logging
from rdkit import Chem
# Make chemprop imports conditional to avoid import errors when not needed
try:
    from chemprop import data, models
    # Try to import featurizers separately as it may not exist in all versions
    try:
        from chemprop import featurizers
    except (ImportError, AttributeError):
        featurizers = None
    CHEMPROP_AVAILABLE = True
except (ImportError, AttributeError) as e:
    CHEMPROP_AVAILABLE = False
    # Create dummy objects to avoid NameError if someone tries to use them
    data = None
    featurizers = None
    models = None
from lightning import pytorch as pl

# ---- QUIET MODE (put these lines at the top of your script) ----
import os, warnings, logging
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"

# Make PyTorch stop suggesting Tensor Core settings
torch.set_float32_matmul_precision("high")

# Silence Python warnings (fine-tune as needed)
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning, message=r".*predict_dataloader.*many workers.*")
warnings.filterwarnings("ignore", message=r"Dropping last batch of size .*")

# Quiet RDKit
from rdkit import RDLogger
RDLogger.DisableLog("rdApp.*")

# Quiet common loggers (Lightning, Chemprop, etc.)
logging.basicConfig(level=logging.ERROR, force=True)
for name in [
    "lightning", "pytorch_lightning", "lightning.pytorch",
    "chemprop", "rdkit", "urllib3", "torch"
]:
    logging.getLogger(name).setLevel(logging.ERROR)
# ---------------------------------------------------------------


# from admet_ai import ADMETModel
import sys
import os
import numpy as np
from transformers import AutoModelForMaskedLM
import warnings
import numpy as np
import torch.nn as nn
from rdkit import Chem
from collections import defaultdict
import pdb

import math
import sys
import os
from pathlib import Path

# Resolve cas_predictor / protein2pam through the repo-root path helpers.
# Both are optional: if absent, the dependent objectives raise a clear error
# at construction time rather than at import time.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from paths import cas_predictor_root, protein2pam_root  # noqa: E402

try:
    cas_predictor_path = cas_predictor_root()
    if str(cas_predictor_path) not in sys.path:
        sys.path.insert(0, str(cas_predictor_path))
except FileNotFoundError:
    cas_predictor_path = Path("/nonexistent/cas_predictor")

try:
    protein2pam_path = protein2pam_root()
    if str(protein2pam_path) not in sys.path:
        sys.path.insert(0, str(protein2pam_path))
except FileNotFoundError:
    protein2pam_path = Path("/nonexistent/protein2pam")

try:
    # Import using importlib to avoid conflicts with editflows/model package
    import importlib.util
    import importlib
    
    # Temporarily add cas_predictor to sys.path at the front to ensure it takes precedence
    cas_predictor_str = str(cas_predictor_path)
    original_path = sys.path.copy()
    if cas_predictor_str not in sys.path:
        sys.path.insert(0, cas_predictor_str)
    
    try:
        # Load model.py directly - use cas_predictor.model as the module name
        model_file = cas_predictor_path / "model.py"
        spec = importlib.util.spec_from_file_location("cas_predictor.model", model_file)
        cas_model = importlib.util.module_from_spec(spec)
        # Set __file__ to help with relative imports
        cas_model.__file__ = str(model_file)
        # Register the module in sys.modules before exec_module so imports work
        sys.modules['cas_predictor.model'] = cas_model
        sys.modules['model'] = cas_model  # Also register as 'model' for lightning_module imports
        spec.loader.exec_module(cas_model)
        Cas9Classifier = cas_model.Cas9Classifier
        
        # Load lightning_module.py - it will import from 'model' which should now resolve correctly
        lightning_file = cas_predictor_path / "lightning_module.py"
        spec_lightning = importlib.util.spec_from_file_location("cas_predictor.lightning_module", lightning_file)
        cas_lightning = importlib.util.module_from_spec(spec_lightning)
        # Set __file__ to help with relative imports
        cas_lightning.__file__ = str(lightning_file)
        # Register the module in sys.modules before exec_module
        sys.modules['cas_predictor.lightning_module'] = cas_lightning
        spec_lightning.loader.exec_module(cas_lightning)
        Cas9ClassifierModule = cas_lightning.Cas9ClassifierModule
        
        import yaml
        try:
            from easydict import EasyDict as edict
        except ImportError:
            # Fallback: use regular dict if easydict not available
            class EasyDict(dict):
                def __init__(self, *args, **kwargs):
                    super().__init__(*args, **kwargs)
                def __getattr__(self, key):
                    try:
                        return self[key]
                    except KeyError:
                        raise AttributeError(key)
                def __setattr__(self, key, value):
                    self[key] = value
            edict = EasyDict
        CAS9_PREDICTOR_AVAILABLE = True
    finally:
        # Restore original sys.path
        sys.path[:] = original_path
except (ImportError, Exception) as e:
    print(f"Warning: Could not import Cas9 predictor modules: {e}")
    import traceback
    traceback.print_exc()
    CAS9_PREDICTOR_AVAILABLE = False

# SMARTS patterns
_AMIDE_SMARTS        = Chem.MolFromSmarts("[CX3](=[OX1])[NX3]")                     # C(=O)-N
_CARBONYL_C_SMARTS   = Chem.MolFromSmarts("[CX3](=[OX1])")                          # carbonyl C
_DIPEPTIDE_SMARTS    = Chem.MolFromSmarts("[CX3](=[OX1])N[#6X4][CX3](=[OX1])N")     # amide–C(sp3)–amide

def _amide_bond_indices(mol, ignore_ring_amides=False):
    ids = set()
    for c_idx, _, n_idx in mol.GetSubstructMatches(_AMIDE_SMARTS):
        b = mol.GetBondBetweenAtoms(c_idx, n_idx)
        if b and b.GetBondType() == Chem.rdchem.BondType.SINGLE:
            if ignore_ring_amides and b.IsInRing():
                continue
            ids.add(b.GetIdx())
    return ids

def _carbonyl_c_indices(mol):
    return {m[0] for m in mol.GetSubstructMatches(_CARBONYL_C_SMARTS)}

def _carbonyl_neighbor_stats(mol, c_indices):
    stats = {"total": 0, "with_N": 0, "with_O": 0, "with_S": 0, "pure_amide": 0}
    for c_idx in c_indices:
        c = mol.GetAtomWithIdx(c_idx)
        stats["total"] += 1
        hasN = hasO = hasS = False
        for b in c.GetBonds():
            if b.GetBondType() != Chem.rdchem.BondType.SINGLE:
                continue
            z = b.GetOtherAtom(c).GetAtomicNum()
            if   z == 7:  hasN = True
            elif z == 8:  hasO = True
            elif z == 16: hasS = True
        stats["with_N"] += int(hasN)
        stats["with_O"] += int(hasO)
        stats["with_S"] += int(hasS)
        if hasN and not (hasO or hasS):
            stats["pure_amide"] += 1
    return stats

def _adjacent_amide_pairs(mol):
    # Count distinct amide–C(sp3)–amide windows (dedup by central carbon)
    centers = set()
    for match in mol.GetSubstructMatches(_DIPEPTIDE_SMARTS):
        centers.add(match[3])  # central sp3 carbon index
    return len(centers)

def analyze_peptide_likeness(smiles: str,
                             ignore_ring_amides: bool = False,
                             amide_density_target: float = 0.12):
    """
    Compute peptide-likeness metrics and a continuous score in [0,1].
    - ignore_ring_amides=False: include macro/cyclic peptides by default.
    - amide_density_target: amide-per-atom density to hit score~1 for peptides
      (~0.10–0.15 works well; default 0.12 ≈ 1 amide per ~8 heavy atoms).
    """
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError(f"Invalid SMILES: {smiles}")

    n_heavy_atoms  = mol.GetNumAtoms()
    n_heavy_bonds  = mol.GetNumBonds()

    amide_bonds    = _amide_bond_indices(mol, ignore_ring_amides=ignore_ring_amides)
    n_amide_bonds  = len(amide_bonds)

    carbonyl_cs    = _carbonyl_c_indices(mol)
    cstats         = _carbonyl_neighbor_stats(mol, carbonyl_cs)
    total_carb     = max(1, cstats["total"])

    # Core features
    f1 = cstats["with_N"] / total_carb  # acyl-N fraction ∈ [0,1]
    amide_per_atom = n_amide_bonds / max(1, n_heavy_atoms)
    f2 = min(1.0, amide_per_atom / max(1e-8, amide_density_target))  # saturate at 1
    n_adjacent = _adjacent_amide_pairs(mol)
    f3 = 1.0 - math.exp(-n_adjacent)  # 0, 0.63, 0.86, 0.95, ... as pairs increase

    # Penalty for non-peptidic carbonyls (carbamates/anhydrides/thioesters)
    pure_amide_fraction = cstats["pure_amide"] / total_carb
    penalty = 1.0 - pure_amide_fraction  # 0 (all pure amide) … 1 (no pure amide)

    # Final heuristic score in [0,1]
    score = 0.55 * f1 + 0.25 * f2 + 0.20 * f3 - 0.25 * penalty
    score = max(0.0, min(1.0, score))

    return {
        "n_heavy_atoms": n_heavy_atoms,
        "n_heavy_bonds": n_heavy_bonds,
        "n_carbonyls": cstats["total"],
        "n_amide_bonds": n_amide_bonds,
        "amide_bond_ratio_all_bonds": n_amide_bonds / max(1, n_heavy_bonds),
        "acyl_N_fraction": f1,
        "pure_amide_fraction": pure_amide_fraction,
        "amide_per_atom": amide_per_atom,
        "n_adjacent_amide_pairs": n_adjacent,
        "peptide_likeness": score,  # <<< continuous score in [0,1]
    }


def score_combination(ratios, scores, admet_scores):
    high_mask = ratios > 0.6
    low_mask  = ratios < 0.1
    mid_mask  = ~(high_mask | low_mask)

    # start with zeros
    final_scores = torch.zeros_like(scores)

    # high-peptide: use peptide scores
    final_scores[high_mask] = scores[high_mask]

    # low-peptide: use admet scores
    final_scores[low_mask] = admet_scores[low_mask]

    # middle band: linear blend
    if mid_mask.any():
        r_mid = ratios[mid_mask]
        alpha = (r_mid - 0.1) / 0.5          # in [0, 1]
        blended = alpha * scores[mid_mask] + (1 - alpha) * admet_scores[mid_mask]
        final_scores[mid_mask] = blended

def detokenize_output(x, cfg, tokenizer, bos_id, eos_id, pad_id):
    """
    Convert a single generated sequence (1, L) back to string.
    """
    seq = x[0].tolist()
    # strip padding
    seq = [tok for tok in seq if tok != pad_id]
    # strip BOS/EOS
    if len(seq) > 0 and seq[0] == bos_id:
        seq = seq[1:]
    if len(seq) > 0 and seq[-1] == eos_id:
        seq = seq[:-1]

    if cfg.task == 'protein':
        # esm tokenizer has batch_decode
        return tokenizer.batch_decode([seq], skip_special_tokens=True)[0]
    elif cfg.task in ('smiles', 'selfies'):
        return tokenizer.decode(seq)
    else:
        return " ".join(map(str, seq))



class Cas9Classification:
    """
    Cas9 binary classification objective.
    
    This objective works with protein sequences (not SMILES).
    It uses ESM-2 embeddings and a trained Cas9Classifier to predict
    whether a protein sequence is Cas9-like.
    
    Note: This objective expects protein sequences, not SMILES sequences.
    If used with SMILES-based generation, you'll need to convert SMILES to proteins first.
    """
    def __init__(self, device, checkpoint_path=None, config_path=None, shared_esm_model=None):
        if not CAS9_PREDICTOR_AVAILABLE:
            raise ImportError(
                "Cas9 predictor modules not available. "
                "Ensure cas_predictor is in the correct location."
            )
        
        self.device = device
        
        # Default checkpoint path if not provided
        if checkpoint_path is None:
            from paths import cas9_classifier_ckpt
            checkpoint_path = str(cas9_classifier_ckpt())
        
        # Load config
        config = None
        if config_path:
            with open(config_path, "r") as f:
                config_dict = yaml.safe_load(f)
                config = edict(config_dict)
        else:
            # Try to load from checkpoint
            try:
                checkpoint = torch.load(checkpoint_path, map_location="cpu")
                if "hyper_parameters" in checkpoint:
                    config_dict = checkpoint["hyper_parameters"]
                    if isinstance(config_dict, dict):
                        config = edict(config_dict)
            except Exception as e:
                print(f"Warning: Could not load config from checkpoint: {e}")
        
        # If config still not found, create a default config with reasonable defaults
        if config is None:
            print("Warning: Could not load config from checkpoint or config_path. Using default config.")
            config = edict({
                "model": edict({
                    "esm_model_name": "facebook/esm2_t33_650M_UR50D",
                    "freeze_esm": True,
                    "hidden_size": 512,
                    "dropout": 0.1,
                    "pooling_type": "mean"
                }),
                "optim": edict({
                    "lr": 5e-5,
                    "beta1": 0.9,
                    "beta2": 0.999,
                    "eps": 1e-8,
                    "weight_decay": 0.01,
                    "warmup_ratio": 0.1
                })
            })
        
        # Load model from checkpoint
        # If shared_esm_model is provided, we need to manually instantiate and load
        # to avoid creating a duplicate ESM model
        if shared_esm_model is not None:
            # Manually instantiate with shared ESM model
            self.model = Cas9ClassifierModule(config, shared_esm_model=shared_esm_model)
            # Load checkpoint state dict (excluding ESM weights since we're using shared model)
            checkpoint = torch.load(checkpoint_path, map_location="cpu")
            state_dict = checkpoint["state_dict"]
            # Filter out ESM weights to avoid loading them (we're using shared model)
            filtered_state_dict = {k: v for k, v in state_dict.items() 
                                 if not k.startswith("model.esm_emb")}
            self.model.load_state_dict(filtered_state_dict, strict=False)
            print(f"Loaded Cas9Classifier checkpoint with shared ESM model. "
                  f"Filtered out {len(state_dict) - len(filtered_state_dict)} ESM parameters.")
        else:
            # Original behavior: use load_from_checkpoint (creates its own ESM model)
            self.model = Cas9ClassifierModule.load_from_checkpoint(
                checkpoint_path,
                config=config,
                strict=False,
            )
        
        self.model.eval()
        self.model.to(device)
        
        # Load ESM tokenizer (same approach as BindingAffinity)
        # We only need the alphabet for tokenization, not the full model
        # If we have a shared ESM model, we can get the tokenizer from it
        # Otherwise, load the ESM model just to get the alphabet
        if shared_esm_model is not None:
            # Use the shared ESM model's tokenizer if available
            # For transformers.EsmModel, we need to use EsmTokenizer separately
            from transformers import EsmTokenizer
            esm_model_name = getattr(config.model, "esm_model_name", "facebook/esm2_t33_650M_UR50D")
            tokenizer = EsmTokenizer.from_pretrained(esm_model_name)
            # Create a simple batch converter function
            def batch_converter(data):
                # data is list of (name, sequence) tuples
                sequences = [seq for _, seq in data]
                encoded = tokenizer(sequences, padding=True, return_tensors="pt", add_special_tokens=True)
                # Return format: (names, sequences, tokens)
                names = [name for name, _ in data]
                return names, sequences, encoded["input_ids"]
            self.batch_converter = batch_converter
        else:
            # Use HuggingFace transformers API instead of old esm.pretrained API
            from transformers import EsmTokenizer
            esm_model_name = getattr(config.model, "esm_model_name", "facebook/esm2_t33_650M_UR50D")
            tokenizer = EsmTokenizer.from_pretrained(esm_model_name)
            # Create a simple batch converter function
            def batch_converter(data):
                # data is list of (name, sequence) tuples
                sequences = [seq for _, seq in data]
                encoded = tokenizer(sequences, padding=True, return_tensors="pt", add_special_tokens=True)
                # Return format: (names, sequences, tokens)
                names = [name for name, _ in data]
                return names, sequences, encoded["input_ids"]
            self.batch_converter = batch_converter
    
    def get_scores(self, protein_seqs):
        """
        Get Cas9 classification scores for protein sequences.
        
        Args:
            protein_seqs: List of protein sequence strings (amino acid sequences) or single string
        
        Returns:
            scores: List of probabilities (0-1) that each sequence is Cas9-like
        """
        # Handle single string input
        if isinstance(protein_seqs, str):
            protein_seqs = [protein_seqs]
        
        if not protein_seqs:
            return []
        
        scores = []
        
        with torch.no_grad():
            # Tokenize sequences using ESM batch converter
            data = [(f"seq_{i}", seq) for i, seq in enumerate(protein_seqs)]
            _, _, batch_tokens = self.batch_converter(data)
            batch_tokens = batch_tokens.to(self.device)
            
            # Create attention mask (True for valid tokens, False for padding)
            # ESM uses 1 for padding, so we check for non-padding tokens
            # Convert to bool tensor explicitly (model expects BoolTensor)
            attention_mask = (batch_tokens != 1).bool().to(self.device)  # 1 is ESM's padding token
            
            # Debug: Print shapes and values for first sequence if debugging
            if len(protein_seqs) == 1 and len(protein_seqs[0]) < 2000:  # Only debug for single sequences and reasonable lengths
                import os
                if os.environ.get("DEBUG_CAS9", "0") == "1":
                    print(f"[DEBUG] Cas9 classifier - sequence length: {len(protein_seqs[0])}")
                    print(f"[DEBUG] Cas9 classifier - batch_tokens shape: {batch_tokens.shape}")
                    print(f"[DEBUG] Cas9 classifier - attention_mask shape: {attention_mask.shape}")
                    print(f"[DEBUG] Cas9 classifier - attention_mask sum (valid tokens): {attention_mask.sum().item()}")
                    print(f"[DEBUG] Cas9 classifier - batch_tokens min/max: {batch_tokens.min().item()}/{batch_tokens.max().item()}")
            
            # Forward pass
            logits = self.model(batch_tokens, attention_mask)
            
            # Debug: Print logits if debugging
            if len(protein_seqs) == 1 and len(protein_seqs[0]) < 2000:
                import os
                if os.environ.get("DEBUG_CAS9", "0") == "1":
                    print(f"[DEBUG] Cas9 classifier - logits shape: {logits.shape}")
                    print(f"[DEBUG] Cas9 classifier - logits value: {logits.item()}")
            
            probs = torch.sigmoid(logits).cpu()
            
            # Debug: Print probabilities if debugging
            if len(protein_seqs) == 1 and len(protein_seqs[0]) < 2000:
                import os
                if os.environ.get("DEBUG_CAS9", "0") == "1":
                    print(f"[DEBUG] Cas9 classifier - probabilities: {probs.tolist()}")
            
            scores = probs.tolist()
        
        return scores
    
    def __call__(self, protein_tokens, protein_seqs):
        """
        Objective call interface.
        
        Args:
            protein_tokens: Unused (kept for interface compatibility)
            protein_seqs: List of protein sequence strings
        
        Returns:
            Tuple of ('cas9', scores) where scores is a list of probabilities
        """
        scores = self.get_scores(protein_seqs)
        return 'cas9', scores


class DeletionCount:
    """
    Objective that maximizes the number of deletions (length reduction) in protein sequences.
    
    This objective computes the difference between the original sequence length and
    the current sequence length, normalized by a maximum deletion percentage to produce
    a value between 0 and 1.
    
    Args:
        original_seq: The original protein sequence string (before any edits)
        max_deletion_percentage: Maximum deletion percentage (0-1). The objective value will be computed as
                                 (deletion_count / (original_length * max_deletion_percentage)), clamped to [0, 1].
                                 If None, defaults to 1.0 (100% deletion possible).
    """
    def __init__(self, original_seq, max_deletion_percentage=None):
        # Store original sequence length (excluding spaces)
        self.original_length = len(original_seq.replace(' ', ''))
        # Set max_deletion_percentage to 1.0 (100%) if not provided
        if max_deletion_percentage is None:
            self.max_deletion_percentage = 1.0
        else:
            if max_deletion_percentage <= 0 or max_deletion_percentage > 1.0:
                raise ValueError(f"max_deletion_percentage must be in (0, 1], got {max_deletion_percentage}")
            self.max_deletion_percentage = float(max_deletion_percentage)
        
        # Calculate max_deletion as absolute value for backward compatibility
        self.max_deletion = self.original_length * self.max_deletion_percentage
    
    def __call__(self, protein_tokens, protein_seqs):
        """
        Objective call interface.
        
        Args:
            protein_tokens: Unused (kept for interface compatibility)
            protein_seqs: List of protein sequence strings
        
        Returns:
            Tuple of ('deletion_count', scores) where scores is a list of normalized deletion
            percentages (deletion_count / (original_length * max_deletion_percentage)), clamped to [0, 1]
        """
        # Handle single string input
        if isinstance(protein_seqs, str):
            protein_seqs = [protein_seqs]
        
        if not protein_seqs:
            return 'deletion_count', []
        
        scores = []
        for seq in protein_seqs:
            # Remove spaces and compute current length
            current_length = len(seq.replace(' ', ''))
            # Number of deletions = original_length - current_length
            deletion_count = self.original_length - current_length
            # Normalize by max_deletion and clamp to [0, 1]
            normalized_score = max(0.0, min(1.0, deletion_count / self.max_deletion))
            scores.append(float(normalized_score))
        
        return 'deletion_count', scores


class PAMMatching:
    """
    PAM matching objective using temperature-scaled log probabilities.
    
    This objective uses the HuggingFace protein2pam model (cas9_full) to predict PAM logits
    for protein sequences. The score is computed using temperature-scaled log probabilities:
    at each target position (non-N), we compute log P(target_nucleotide | position) after
    applying temperature scaling to the logits. The final score is the geometric mean of
    probabilities (exp of mean log prob) over target positions.
    
    This approach amplifies small improvements, making them more visible in the weighted sum
    during optimization, which helps when candidates show only marginal improvements.
    
    Args:
        device: torch device
        target_pam: Target PAM sequence string of length 10 (e.g., "NGGNNNNNNN")
        model_name: HuggingFace model identifier (default: "Profluent-Bio/protein2pam-cas9_full")
        sigmoid_temperature: Temperature for scaling logits before softmax. Lower values (0.1-0.3)
                            sharpen the distribution, making small improvements more visible.
                            Default 0.2. (Note: parameter name kept for backward compatibility)
        pam_prediction_temperature: Temperature for PAM prediction (in predict_pam method).
                                    Values < 1.0 make distributions sharper (more confident),
                                    > 1.0 make them softer. Default 1.0 (no temperature scaling).
                                    Note: This only affects prediction, not scoring.
        shared_esm_model: Optional shared ESM model (transformers.EsmModel) to use instead of
                         loading a separate ESM backbone. This reduces memory usage when multiple
                         models need ESM embeddings. Default None (loads its own ESM model).
    """
    def __init__(self, device, target_pam, model_name="Profluent-Bio/protein2pam-cas9_full", sigmoid_temperature=0.2, use_entropy_for_n_positions=True, use_ce_loss=False, pam_min_confidence=0.55, pam_prediction_temperature=1.0, shared_esm_model=None):
        self.device = device
        self.target_pam = target_pam.upper()
        
        # Validate target PAM length
        if len(self.target_pam) != 10:
            raise ValueError(f"Target PAM must be exactly 10 nucleotides, got {len(self.target_pam)}")
        
        # Validate nucleotides
        valid_nucleotides = set('ACGTN')
        if not all(nuc in valid_nucleotides for nuc in self.target_pam):
            raise ValueError(f"Target PAM contains invalid nucleotides. Only ACGTN allowed.")
        
        # Convert target PAM to class indices (ACGT = 0,1,2,3, N = -1 for masking)
        nucleotides = ['A', 'C', 'G', 'T']
        self.target_indices = []
        self.target_mask = []
        for nuc in self.target_pam:
            if nuc == 'N':
                # N means any nucleotide - we'll mask this position in loss calculation
                self.target_indices.append(-1)  # Placeholder, will be masked
                self.target_mask.append(False)
            else:
                self.target_indices.append(nucleotides.index(nuc))
                self.target_mask.append(True)
        
        self.target_indices = torch.tensor(self.target_indices, device=device, dtype=torch.long)
        self.target_mask = torch.tensor(self.target_mask, device=device, dtype=torch.bool)
        self.sigmoid_temperature = float(sigmoid_temperature)
        self.use_entropy_for_n_positions = bool(use_entropy_for_n_positions)
        self.use_ce_loss = bool(use_ce_loss)
        self.pam_min_confidence = float(pam_min_confidence)
        self.pam_prediction_temperature = float(pam_prediction_temperature)
        
        # Load HuggingFace model. The protein2pam package root is resolved by
        # paths.protein2pam_root() (honours $PROTEIN2PAM_ROOT).
        
        try:
            protein2pam_path_resolved = protein2pam_root()
        except FileNotFoundError as exc:
            raise ImportError(str(exc)) from exc
        
        # Ensure the path is in sys.path
        protein2pam_path_str = str(protein2pam_path_resolved)
        if protein2pam_path_str not in sys.path:
            sys.path.insert(0, protein2pam_path_str)
            print(f"[DEBUG] Added protein2pam path to sys.path: {protein2pam_path_str}")
        
        try:
            # Try importing directly from huggingface submodule to avoid protein2pam/__init__.py
            # which imports PAMOracle (requires torch at import time)
            import importlib.util
            
            # Import the modules directly without going through __init__.py
            huggingface_dir = protein2pam_path_resolved / "protein2pam" / "huggingface"
            
            # Load configuration_esm
            config_file = huggingface_dir / "configuration_esm.py"
            if config_file.exists():
                spec_config = importlib.util.spec_from_file_location("protein2pam.huggingface.configuration_esm", config_file)
                if spec_config and spec_config.loader:
                    # Create package structure
                    if 'protein2pam' not in sys.modules:
                        sys.modules['protein2pam'] = type(sys)('protein2pam')
                    if 'protein2pam.huggingface' not in sys.modules:
                        sys.modules['protein2pam.huggingface'] = type(sys)('protein2pam.huggingface')
                    
                    config_module = importlib.util.module_from_spec(spec_config)
                    sys.modules['protein2pam.huggingface.configuration_esm'] = config_module
                    spec_config.loader.exec_module(config_module)
            
            # Load modeling_esm (depends on configuration_esm)
            modeling_file = huggingface_dir / "modeling_esm.py"
            if modeling_file.exists():
                spec_modeling = importlib.util.spec_from_file_location("protein2pam.huggingface.modeling_esm", modeling_file)
                if spec_modeling and spec_modeling.loader:
                    modeling_module = importlib.util.module_from_spec(spec_modeling)
                    sys.modules['protein2pam.huggingface.modeling_esm'] = modeling_module
                    spec_modeling.loader.exec_module(modeling_module)
            
            # Load tokenizer
            tokenizer_file = huggingface_dir / "tokenizer.py"
            if tokenizer_file.exists():
                spec_tokenizer = importlib.util.spec_from_file_location("protein2pam.huggingface.tokenizer", tokenizer_file)
                if spec_tokenizer and spec_tokenizer.loader:
                    tokenizer_module = importlib.util.module_from_spec(spec_tokenizer)
                    sys.modules['protein2pam.huggingface.tokenizer'] = tokenizer_module
                    spec_tokenizer.loader.exec_module(tokenizer_module)
            
            # Now import from the loaded modules
            EsmForSequenceClassification = sys.modules['protein2pam.huggingface.modeling_esm'].EsmForSequenceClassification
            get_tokenizer = sys.modules['protein2pam.huggingface.tokenizer'].get_tokenizer
            
            self.tokenizer = get_tokenizer()
            print(f"Loading PAM prediction model: {model_name}...")
            # Load model without device_map (device_map expects string/dict, not torch.device)
            # Then move to device manually for compatibility with different transformers versions
            self.model = EsmForSequenceClassification.from_pretrained(model_name)
            
            # If shared_esm_model is provided, replace the ESM backbone to share parameters
            if shared_esm_model is not None:
                print(f"Replacing PAM model's ESM backbone with shared ESM model...")
                # EsmForSequenceClassification typically has an 'esm' attribute containing the backbone
                if hasattr(self.model, 'esm'):
                    self.model.esm = shared_esm_model
                    print(f"Successfully replaced PAM model's ESM backbone with shared model.")
                elif hasattr(self.model, 'model') and hasattr(self.model.model, 'esm'):
                    # Some models wrap ESM in a 'model' attribute
                    self.model.model.esm = shared_esm_model
                    print(f"Successfully replaced PAM model's ESM backbone (via .model.esm) with shared model.")
                else:
                    print(f"Warning: Could not find ESM backbone in PAM model structure. Available attributes: {[attr for attr in dir(self.model) if not attr.startswith('_')]}")
                    # Try to find it recursively
                    for attr_name in ['esm', 'backbone', 'encoder']:
                        if hasattr(self.model, attr_name):
                            setattr(self.model, attr_name, shared_esm_model)
                            print(f"Replaced {attr_name} with shared ESM model.")
                            break
            
            self.model = self.model.to(device)
            self.model.eval()
            print(f"PAM prediction model loaded successfully.")
        except ImportError as e:
            # Show more details about the import error
            import traceback
            print(f"[DEBUG] protein2pam import error: {e}")
            print(f"[DEBUG] sys.path entries containing 'protein2pam': {[p for p in sys.path if 'protein2pam' in p]}")
            print(f"[DEBUG] Expected protein2pam path: {protein2pam_path_resolved}")
            print(f"[DEBUG] Path exists: {protein2pam_path_resolved.exists() if protein2pam_path_resolved else False}")
            traceback.print_exc()
            raise ImportError(
                f"protein2pam package not available: {e}. "
                "Please ensure protein2pam is installed and in the correct location."
            )
        except Exception as e:
            import traceback
            print(f"[DEBUG] Failed to load PAM prediction model: {e}")
            traceback.print_exc()
            raise RuntimeError(f"Failed to load PAM prediction model: {e}")
    
    def get_pam_probability_distributions(self, protein_seqs):
        """
        Get detailed probability distributions for PAM predictions.
        
        Args:
            protein_seqs: List of protein sequence strings
        
        Returns:
            List of dictionaries, each containing:
            - 'probabilities': (10, 4) tensor of probabilities for each position and nucleotide
            - 'predicted_pam': Predicted PAM string
            - 'per_position': List of dicts with position info (nucleotide probs, entropy, etc.)
        """
        if not protein_seqs:
            return []
        
        # Handle single string input
        if isinstance(protein_seqs, str):
            protein_seqs = [protein_seqs]
        
        nucleotides = ['A', 'C', 'G', 'T']
        results = []
        
        with torch.no_grad():
            # Tokenize sequences
            encodings = self.tokenizer.encode_batch(protein_seqs)
            input_batch = dict(
                input_ids=torch.tensor([encoding.ids for encoding in encodings], device=self.device),
                attention_mask=torch.tensor([encoding.attention_mask for encoding in encodings], device=self.device),
            )
            
            # Get PAM predictions
            output = self.model(**input_batch)
            logits = output.logits  # (batch_size, 10, 4)
            probabilities = F.softmax(logits, dim=-1)  # (batch_size, 10, 4)
            log_probs = F.log_softmax(logits, dim=-1)  # (batch_size, 10, 4)
            
            # Process each sequence
            for i in range(probabilities.shape[0]):
                probs = probabilities[i]  # (10, 4)
                log_probs_seq = log_probs[i]  # (10, 4)
                pam_seq = []
                per_position = []
                
                for pos_idx in range(probs.shape[0]):
                    pos_probs = probs[pos_idx, :]  # (4,)
                    pos_log_probs = log_probs_seq[pos_idx, :]  # (4,)
                    
                    # Compute entropy for this position
                    entropy = -(pos_probs * pos_log_probs).sum().item()
                    max_entropy = math.log(4)  # Maximum entropy for uniform distribution
                    normalized_entropy = entropy / max_entropy
                    
                    max_prob = pos_probs.max().item()
                    max_idx = pos_probs.argmax().item()
                    predicted_nuc = nucleotides[max_idx]
                    
                    # Create position info
                    pos_info = {
                        'position': pos_idx,
                        'probabilities': {nuc: pos_probs[j].item() for j, nuc in enumerate(nucleotides)},
                        'predicted': predicted_nuc,
                        'max_probability': max_prob,
                        'entropy': entropy,
                        'normalized_entropy': normalized_entropy,
                    }
                    per_position.append(pos_info)
                    
                    # Predict nucleotide (use instance min_confidence)
                    if max_prob < self.pam_min_confidence:
                        pam_seq.append('N')
                    else:
                        pam_seq.append(predicted_nuc)
                
                results.append({
                    'probabilities': probs.cpu(),
                    'predicted_pam': ''.join(pam_seq),
                    'per_position': per_position,
                })
        
        return results
    
    def predict_pam(self, protein_seqs, min_confidence=None):
        """
        Predict PAM sequences for protein sequences, with support for 'N' when confidence is low.
        
        Uses temperature scaling (pam_prediction_temperature) to control distribution sharpness.
        Lower temperatures (< 1.0) produce sharper distributions and more confident predictions.
        
        Args:
            protein_seqs: List of protein sequence strings
            min_confidence: Minimum probability threshold for predicting a specific nucleotide.
                           If max probability < min_confidence, predict 'N'.
                           If None, uses self.pam_min_confidence (default: None)
        
        Returns:
            predicted_pams: List of predicted PAM strings (10 nucleotides each)
        """
        if min_confidence is None:
            min_confidence = self.pam_min_confidence
        if not protein_seqs:
            return []
        
        # Handle single string input
        if isinstance(protein_seqs, str):
            protein_seqs = [protein_seqs]
        
        # Check if model and tokenizer are initialized
        if not hasattr(self, 'model') or self.model is None:
            print(f"[ERROR] PAMMatching.predict_pam: model is not initialized. Cannot predict PAM.")
            return []
        
        if not hasattr(self, 'tokenizer') or self.tokenizer is None:
            print(f"[ERROR] PAMMatching.predict_pam: tokenizer is not initialized. Cannot predict PAM.")
            return []
        
        predicted_pams = []
        nucleotides = ['A', 'C', 'G', 'T']
        
        try:
            with torch.no_grad():
                # Tokenize sequences
                encodings = self.tokenizer.encode_batch(protein_seqs)
                input_batch = dict(
                    input_ids=torch.tensor([encoding.ids for encoding in encodings], device=self.device),
                    attention_mask=torch.tensor([encoding.attention_mask for encoding in encodings], device=self.device),
                )
                
                # Get PAM predictions
                output = self.model(**input_batch)
                logits = output.logits  # (batch_size, 10, 4) - 10 positions, 4 nucleotides (ACGT)
                # Apply temperature scaling for prediction (sharper distributions when < 1.0)
                scaled_logits = logits / self.pam_prediction_temperature
                probabilities = F.softmax(scaled_logits, dim=-1)  # (batch_size, 10, 4)
                
                # Predict PAM for each sequence
                for i in range(probabilities.shape[0]):
                    probs = probabilities[i]  # (10, 4)
                    pam_seq = []
                    
                    for pos_idx in range(probs.shape[0]):
                        pos_probs = probs[pos_idx, :]  # (4,)
                        max_prob = pos_probs.max().item()
                        max_idx = pos_probs.argmax().item()
                        
                        # If confidence is low, predict 'N'
                        if max_prob < min_confidence:
                            pam_seq.append('N')
                        else:
                            pam_seq.append(nucleotides[max_idx])
                    
                    predicted_pams.append(''.join(pam_seq))
        except Exception as e:
            import traceback
            print(f"[ERROR] PAMMatching.predict_pam failed: {e}")
            traceback.print_exc()
            return []
        
        return predicted_pams
    
    def get_scores(self, protein_seqs):
        """
        Get PAM matching scores using either cross-entropy loss or log probability approach.
        
        If use_ce_loss=True:
            Uses cross-entropy loss between predicted logits and target PAM indices.
            Score = exp(-mean_ce_loss) to convert to [0, 1] range (higher is better).
        
        If use_ce_loss=False (default):
            Uses raw probabilities (no temperature scaling) to maintain high probability values.
            For non-N positions: score is the geometric mean of probabilities (exp of mean log prob).
            For N positions: if use_entropy_for_n_positions=True, score is based on entropy.
        
        Args:
            protein_seqs: List of protein sequence strings
        
        Returns:
            scores: List of scores in [0, 1] (higher = better match to target PAM)
        """
        if not protein_seqs:
            return []
        
        # Handle single string input
        if isinstance(protein_seqs, str):
            protein_seqs = [protein_seqs]
        
        scores = []
        max_entropy = math.log(4)  # Maximum entropy for uniform distribution over 4 nucleotides
        
        with torch.no_grad():
            # Tokenize sequences
            encodings = self.tokenizer.encode_batch(protein_seqs)
            input_batch = dict(
                input_ids=torch.tensor([encoding.ids for encoding in encodings], device=self.device),
                attention_mask=torch.tensor([encoding.attention_mask for encoding in encodings], device=self.device),
            )
            
            # Get PAM predictions (logits)
            output = self.model(**input_batch)
            logits = output.logits  # (batch_size, 10, 4) - 10 positions, 4 nucleotides (ACGT)
            
            for i in range(logits.shape[0]):
                seq_logits = logits[i]  # (10, 4)
                
                if self.use_ce_loss:
                    # Cross-entropy loss approach
                    # Separate non-N and N positions
                    non_n_mask = self.target_mask  # True for non-N positions
                    n_mask = ~non_n_mask  # True for N positions
                    
                    score_components = []
                    
                    # For non-N positions: compute cross-entropy loss against target nucleotides
                    if non_n_mask.any():
                        masked_logits = seq_logits[non_n_mask]  # (K, 4)
                        masked_targets = self.target_indices[non_n_mask]  # (K,) in {0,1,2,3}
                        
                        # Compute cross-entropy loss per position (reduction='none')
                        ce_loss_per_pos = F.cross_entropy(
                            masked_logits, masked_targets, reduction='none'
                        )  # (K,)
                        
                        # Mean cross-entropy loss over non-N positions
                        mean_ce_loss = ce_loss_per_pos.mean().item()
                        
                        # Convert to score in [0, 1] range using exp(-ce_loss)
                        # Lower CE loss = higher score (better match)
                        non_n_score = math.exp(-mean_ce_loss)
                        score_components.append(non_n_score)
                    
                    # For N positions: handle based on entropy flag
                    if n_mask.any():
                        if self.use_entropy_for_n_positions:
                            # Use entropy scoring (encourage uniform distribution)
                            log_probs = F.log_softmax(seq_logits, dim=-1)  # (10, 4)
                            probs = torch.exp(log_probs)  # (10, 4)
                            n_probs = probs[n_mask]  # (M, 4) where M is number of N positions
                            n_log_probs = log_probs[n_mask]  # (M, 4)
                            entropy_per_pos = -(n_probs * n_log_probs).sum(dim=-1)  # (M,)
                            mean_entropy = entropy_per_pos.mean().item()
                            n_score = mean_entropy / max_entropy
                            score_components.append(n_score)
                        else:
                            # Use CE loss against uniform distribution (consistent with CE loss approach)
                            # This encourages uniform distribution over all 4 nucleotides
                            n_logits = seq_logits[n_mask]  # (M, 4) where M is number of N positions
                            
                            # Compute cross-entropy loss against uniform distribution [0.25, 0.25, 0.25, 0.25]
                            # CE_uniform = -sum(0.25 * log(p_i)) = -0.25 * sum(log(p_i))
                            # Lower CE loss (closer to uniform) = higher score
                            log_probs = F.log_softmax(n_logits, dim=-1)  # (M, 4)
                            uniform_target = 0.25  # Uniform probability for each nucleotide
                            ce_loss_uniform_per_pos = -(uniform_target * log_probs).sum(dim=-1)  # (M,)
                            
                            # Mean cross-entropy loss over N positions
                            mean_ce_loss_uniform = ce_loss_uniform_per_pos.mean().item()
                            
                            # Convert to score in [0, 1] range using exp(-ce_loss)
                            # Lower CE loss (more uniform) = higher score
                            n_score = math.exp(-mean_ce_loss_uniform)
                            score_components.append(n_score)
                    
                    # Combine scores: if both non-N and N positions exist, take weighted average
                    # Weight by number of positions of each type
                    if len(score_components) == 2:
                        num_non_n = non_n_mask.sum().item()
                        num_n = n_mask.sum().item()
                        total_positions = num_non_n + num_n
                        score = (score_components[0] * num_non_n + score_components[1] * num_n) / total_positions
                    elif len(score_components) == 1:
                        score = score_components[0]
                    else:
                        score = 0.0
                    
                    scores.append(score)
                else:
                    # Original log probability approach (default)
                    # Use raw probabilities (no temperature scaling) to maintain high probabilities
                    # This treats all probability values equally, appropriate for maintaining
                    # high probability rather than optimizing low probabilities upward
                    log_probs = F.log_softmax(seq_logits, dim=-1)  # (10, 4)
                    probs = torch.exp(log_probs)  # (10, 4) - needed for entropy calculation
                    
                    # Separate non-N and N positions
                    non_n_mask = self.target_mask  # True for non-N positions
                    n_mask = ~non_n_mask  # True for N positions
                    
                    score_components = []
                    
                    # For non-N positions: maximize log probability of target nucleotide
                    if non_n_mask.any():
                        masked_log_probs = log_probs[non_n_mask]  # (K, 4)
                        masked_targets = self.target_indices[non_n_mask]  # (K,) in {0,1,2,3}
                        
                        # Get log probability of target nucleotide at each position
                        target_log_probs = masked_log_probs.gather(
                            1, masked_targets.unsqueeze(1)
                        ).squeeze(1)  # (K,)
                        
                        # Mean log probability (geometric mean in probability space)
                        mean_log_prob = target_log_probs.mean().item()
                        
                        # Convert to score in [0, 1] range using geometric mean probability
                        # This gives equal weight to maintaining high probabilities
                        non_n_score = math.exp(mean_log_prob)
                        score_components.append(non_n_score)
                    
                    # For N positions: maximize entropy (encourage uniform distribution) if enabled
                    if n_mask.any() and self.use_entropy_for_n_positions:
                        n_probs = probs[n_mask]  # (M, 4) where M is number of N positions
                        
                        # Compute entropy for each N position: H = -sum(p_i * log(p_i))
                        # Use log_probs for numerical stability: H = -sum(p_i * log_p_i)
                        n_log_probs = log_probs[n_mask]  # (M, 4)
                        entropy_per_pos = -(n_probs * n_log_probs).sum(dim=-1)  # (M,)
                        
                        # Mean entropy across N positions
                        mean_entropy = entropy_per_pos.mean().item()
                        
                        # Normalize by max entropy to get score in [0, 1]
                        # Higher entropy (more uniform) = higher score
                        n_score = mean_entropy / max_entropy
                        score_components.append(n_score)
                    
                    # Combine scores: if both non-N and N positions exist, take weighted average
                    # Weight by number of positions of each type
                    # If entropy is disabled, only use non-N score
                    if len(score_components) == 2:
                        # Weighted average based on number of positions
                        num_non_n = non_n_mask.sum().item()
                        num_n = n_mask.sum().item()
                        total_positions = num_non_n + num_n
                        score = (score_components[0] * num_non_n + score_components[1] * num_n) / total_positions
                    elif len(score_components) == 1:
                        score = score_components[0]
                    else:
                        # No positions to evaluate (shouldn't happen, but handle gracefully)
                        score = 0.0
                    
                    scores.append(score)
        
        return scores
    
    def get_score_for_pam(self, protein_seqs, pam_sequence, use_temperature_scaling=False):
        """
        Get probability score for a specific PAM sequence (not necessarily the target).
        
        This computes the log probability score for an arbitrary PAM sequence.
        By default, uses raw probabilities (no temperature scaling) for more
        interpretable results. Set use_temperature_scaling=True to match the
        optimization objective.
        
        Args:
            protein_seqs: List of protein sequence strings
            pam_sequence: PAM sequence string of length 10 (e.g., "NGGNNNNNNN")
            use_temperature_scaling: If True, apply temperature scaling (default: False for display)
        
        Returns:
            scores: List of scores in [0, 1] (probability of the given PAM sequence)
        """
        if not protein_seqs:
            return []
        
        # Handle single string input
        if isinstance(protein_seqs, str):
            protein_seqs = [protein_seqs]
        
        # Validate and convert PAM sequence to indices
        pam_sequence = pam_sequence.upper()
        if len(pam_sequence) != 10:
            raise ValueError(f"PAM sequence must be exactly 10 nucleotides, got {len(pam_sequence)}")
        
        nucleotides = ['A', 'C', 'G', 'T']
        pam_indices = []
        pam_mask = []
        for nuc in pam_sequence:
            if nuc == 'N':
                pam_indices.append(-1)
                pam_mask.append(False)
            else:
                if nuc not in nucleotides:
                    raise ValueError(f"Invalid nucleotide in PAM sequence: {nuc}")
                pam_indices.append(nucleotides.index(nuc))
                pam_mask.append(True)
        
        pam_indices = torch.tensor(pam_indices, device=self.device, dtype=torch.long)
        pam_mask = torch.tensor(pam_mask, device=self.device, dtype=torch.bool)
        
        scores = []
        
        with torch.no_grad():
            # Tokenize sequences
            encodings = self.tokenizer.encode_batch(protein_seqs)
            input_batch = dict(
                input_ids=torch.tensor([encoding.ids for encoding in encodings], device=self.device),
                attention_mask=torch.tensor([encoding.attention_mask for encoding in encodings], device=self.device),
            )
            
            # Get PAM predictions (logits)
            output = self.model(**input_batch)
            logits = output.logits  # (batch_size, 10, 4) - 10 positions, 4 nucleotides (ACGT)
            
            for i in range(logits.shape[0]):
                seq_logits = logits[i]  # (10, 4)
                
                # Apply temperature scaling only if requested (for display, use raw probabilities)
                if use_temperature_scaling:
                    scaled_logits = seq_logits / self.sigmoid_temperature
                    log_probs = F.log_softmax(scaled_logits, dim=-1)  # (10, 4)
                else:
                    # Use raw probabilities (no temperature scaling) for more interpretable results
                    log_probs = F.log_softmax(seq_logits, dim=-1)  # (10, 4)
                
                if pam_mask.any():
                    masked_log_probs = log_probs[pam_mask]  # (K, 4)
                    masked_pam_indices = pam_indices[pam_mask]  # (K,) in {0,1,2,3}
                    
                    # Get log probability of PAM nucleotide at each position
                    pam_log_probs = masked_log_probs.gather(
                        1, masked_pam_indices.unsqueeze(1)
                    ).squeeze(1)  # (K,)
                    
                    # Mean log probability (linear in log space)
                    mean_log_prob = pam_log_probs.mean().item()
                    
                    # Convert to score in [0, 1] range using geometric mean probability
                    score = math.exp(mean_log_prob)
                else:
                    # All positions are N: use neutral score
                    score = 0.0
                
                scores.append(score)
        
        return scores
        
        # ========================================================================
        # OLD IMPLEMENTATION (logit-margin-based with sigmoid) - COMMENTED OUT
        # ========================================================================
        # This was the previous implementation that struggled to show small improvements
        # because sigmoid saturates quickly. Kept for reference.
        #
        # def get_scores_old(self, protein_seqs):
        #     """
        #     Get PAM matching scores (logit-margin-based, in [0, 1]) for protein sequences.
        #     
        #     At each target (non-N) position: margin = (target PAM logit) - (largest non-target logit),
        #     then score_pos = sigmoid(margin / sigmoid_temperature). The sequence score is the mean
        #     over those positions. Higher is better.
        #     
        #     Args:
        #         protein_seqs: List of protein sequence strings
        #     
        #     Returns:
        #         scores: List of scores in [0, 1] (higher = target PAM preferred over alternatives)
        #     """
        #     if not protein_seqs:
        #         return []
        #     
        #     # Handle single string input
        #     if isinstance(protein_seqs, str):
        #         protein_seqs = [protein_seqs]
        #     
        #     scores = []
        #     
        #     with torch.no_grad():
        #         # Tokenize sequences
        #         encodings = self.tokenizer.encode_batch(protein_seqs)
        #         input_batch = dict(
        #             input_ids=torch.tensor([encoding.ids for encoding in encodings], device=self.device),
        #             attention_mask=torch.tensor([encoding.attention_mask for encoding in encodings], device=self.device),
        #         )
        #         
        #         # Get PAM predictions (logits)
        #         output = self.model(**input_batch)
        #         logits = output.logits  # (batch_size, 10, 4) - 10 positions, 4 nucleotides (ACGT)
        #         
        #         for i in range(logits.shape[0]):
        #             seq_logits = logits[i]  # (10, 4)
        #             masked_positions = self.target_mask
        #             if masked_positions.any():
        #                 masked_logits = seq_logits[masked_positions]  # (K, 4)
        #                 masked_targets = self.target_indices[masked_positions]  # (K,) in {0,1,2,3}
        #                 # Target logit at each position
        #                 target_logits = masked_logits.gather(1, masked_targets.unsqueeze(1)).squeeze(1)  # (K,)
        #                 # Largest non-target logit: mask out target class then max over dim=-1
        #                 logits_copy = masked_logits.clone()
        #                 logits_copy.scatter_(1, masked_targets.unsqueeze(1), -1e9)
        #                 max_non_target = logits_copy.max(dim=1).values  # (K,)
        #                 margin = target_logits - max_non_target  # (K,)
        #                 margin = margin / self.sigmoid_temperature
        #                 score_per_pos = torch.sigmoid(margin)
        #                 score = score_per_pos.mean().item()
        #             else:
        #                 # All positions are N: no target to match, use neutral score
        #                 score = 0.5
        #             scores.append(score)
        #     
        #     return scores
    
    def __call__(self, protein_tokens, protein_seqs):
        """
        Objective call interface.
        
        Args:
            protein_tokens: Unused (kept for interface compatibility)
            protein_seqs: List of protein sequence strings
        
        Returns:
            Tuple of ('pam_matching', scores) where scores is a list of values in [0, 1]:
            - For non-N positions: temperature-scaled log probability (geometric mean over positions)
            - For N positions: if use_entropy_for_n_positions=True, normalized entropy (encouraging uniform distribution)
            - If use_entropy_for_n_positions=False, N positions are ignored (score only considers non-N positions)
            - Combined score is weighted average if both types exist and entropy is enabled
            Higher scores indicate better match to target PAM.
        """
        scores = self.get_scores(protein_seqs)
        return 'pam_matching', scores