script/evo_if_model.py

"""
Evo-IF model components and the Q-Former style fusion bridge.
"""

from __future__ import annotations

import math
from typing import Iterable, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

RNA_SENTINEL_TO_DNA = str.maketrans(
    {
        "b": "A",
        "d": "C",
        "h": "G",
        "u": "T",
        "y": "N",
        "B": "A",
        "D": "C",
        "H": "G",
        "U": "T",
        "Y": "N",
    }
)

POLYMER_CONTEXT_DIM = 8
CONTEXT_ADAPTER_RANK = 16
NAIAD_VOCAB_SIZE = 33


def normalize_naiad_sequence_for_evo2(sequence: str) -> str:
    """Convert NAIAD's DNA/RNA display alphabet into Evo2-style DNA letters."""
    normalized = sequence.translate(RNA_SENTINEL_TO_DNA).upper()
    normalized = normalized.replace("/", "")
    return "".join(ch if ch in "ACGTN" else "N" for ch in normalized)


def polymer_context_features(feature_dict: dict) -> torch.Tensor:
    """One-hot encode structure-level protein/DNA/RNA presence for fusion routing."""
    valid = feature_dict.get("mask")
    if valid is None:
        valid = torch.ones_like(feature_dict["dna_mask"])
    valid = valid > 0
    has_protein = ((feature_dict["protein_mask"] > 0) & valid).any(dim=1).long()
    has_dna = ((feature_dict["dna_mask"] > 0) & valid).any(dim=1).long()
    has_rna = ((feature_dict["rna_mask"] > 0) & valid).any(dim=1).long()
    category = has_protein * 4 + has_dna * 2 + has_rna
    return F.one_hot(category, num_classes=POLYMER_CONTEXT_DIM).to(
        device=feature_dict["dna_mask"].device,
        dtype=feature_dict["dna_mask"].dtype,
    )


def validate_bridge_adapter(payload: dict) -> None:
    """Validate the bridge-only Evo-IF adapter payload."""
    state = payload.get("bridge_state_dict")
    if not isinstance(state, dict) or not state:
        raise ValueError("adapter.pt does not contain bridge weights")
    if not all(isinstance(key, str) and isinstance(value, torch.Tensor) for key, value in state.items()):
        raise TypeError("bridge_state_dict must map string keys to tensors")
    forbidden = {"model_state_dict", "inverse_folding_state_dict", "optimizer_state_dict"}
    present = forbidden.intersection(payload)
    if present:
        raise ValueError(f"adapter.pt must be bridge-only; unexpected keys: {sorted(present)}")


class FullEvo2HiddenEncoder(nn.Module):
    """Frozen official Evo2 runtime wrapper that returns one hidden layer."""

    def __init__(
        self,
        model_name: str,
        checkpoint_path: str,
        layer_name: str,
        use_kernels: bool = False,
    ):
        super().__init__()
        if not checkpoint_path:
            raise ValueError(
                "Download the Evo2 base model separately and pass --evo2-checkpoint."
            )
        try:
            from evo2 import Evo2
        except ImportError as exc:
            raise ImportError(
                "Evo-IF requires the official Evo2 runtime. Install the "
                "dependencies from requirements.txt."
            ) from exc

        evo2 = Evo2(model_name, local_path=checkpoint_path, use_kernels=use_kernels)
        self.model = evo2.model
        self.tokenizer = evo2.tokenizer
        self.layer_name = layer_name
        self.hidden_dim = int(getattr(self.model.config, "hidden_size"))
        self.model.eval()
        for param in self.model.parameters():
            param.requires_grad_(False)

    def tokenize(self, sequences: Iterable[str], device: torch.device) -> torch.Tensor:
        seqs = list(sequences)
        if not seqs:
            raise ValueError("at least one sequence is required")
        tokenized = [self.tokenizer.tokenize(seq) for seq in seqs]
        max_len = max(len(tokens) for tokens in tokenized)
        ids = torch.zeros((len(tokenized), max_len), dtype=torch.int, device=device)
        for row, tokens in enumerate(tokenized):
            values = torch.tensor(tokens, dtype=torch.int, device=device)
            ids[row, : values.numel()] = values
        return ids

    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        embeddings: dict[str, torch.Tensor] = {}

        def hook_fn(_module, _inputs, output):
            if isinstance(output, tuple):
                output = output[0]
            embeddings[self.layer_name] = output.detach()

        layer = self.model.get_submodule(self.layer_name)
        handle = layer.register_forward_hook(hook_fn)
        try:
            with torch.no_grad():
                self.model.forward(token_ids)
        finally:
            handle.remove()

        if self.layer_name not in embeddings:
            raise RuntimeError(f"Evo2 layer did not produce embeddings: {self.layer_name}")
        return embeddings[self.layer_name]


class ContextLowRankAdapter(nn.Module):
    """Zero-output low-rank residual used by one polymer context."""

    def __init__(self, hidden_dim: int, rank: int = CONTEXT_ADAPTER_RANK):
        super().__init__()
        positions = torch.arange(hidden_dim, dtype=torch.float32).unsqueeze(0) + 0.5
        frequencies = torch.arange(1, rank + 1, dtype=torch.float32).unsqueeze(1)
        basis = torch.cos(math.pi * frequencies * positions / hidden_dim)
        basis = basis * math.sqrt(2.0 / hidden_dim)
        self.down_weight = nn.Parameter(basis)
        self.up_weight = nn.Parameter(torch.zeros(hidden_dim, rank))

    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        lowrank = F.gelu(F.linear(hidden, self.down_weight))
        return F.linear(lowrank, self.up_weight)


class ContextLowRankLogitAdapter(nn.Module):
    """Zero-output route-specific residual applied directly to NAIAD logits."""

    def __init__(
        self,
        hidden_dim: int,
        output_dim: int = NAIAD_VOCAB_SIZE,
        rank: int = CONTEXT_ADAPTER_RANK,
    ):
        super().__init__()
        positions = torch.arange(hidden_dim, dtype=torch.float32).unsqueeze(0) + 0.5
        frequencies = torch.arange(1, rank + 1, dtype=torch.float32).unsqueeze(1)
        basis = torch.cos(math.pi * frequencies * positions / hidden_dim)
        basis = basis * math.sqrt(2.0 / hidden_dim)
        self.down_weight = nn.Parameter(basis)
        self.up_weight = nn.Parameter(torch.zeros(output_dim, rank))

    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        lowrank = F.gelu(F.linear(hidden, self.down_weight))
        return F.linear(lowrank, self.up_weight)


def context_lowrank_delta(
    adapters: nn.ModuleList,
    hidden: torch.Tensor,
    context_features: torch.Tensor,
) -> torch.Tensor:
    weights = context_features.to(device=hidden.device, dtype=hidden.dtype)
    return sum(
        weights[:, index].reshape(-1, 1, 1) * adapter(hidden)
        for index, adapter in enumerate(adapters)
    )


def context_lowrank_logit_delta(
    adapters: nn.ModuleList,
    hidden: torch.Tensor,
    context_features: torch.Tensor,
) -> torch.Tensor:
    weights = context_features.to(device=hidden.device, dtype=hidden.dtype)
    return sum(
        weights[:, index].reshape(-1, 1, 1) * adapter(hidden)
        for index, adapter in enumerate(adapters)
    )


class Evo2QFormerBridge(nn.Module):
    """
    Small Q-Former-style bridge.

    Learnable queries attend to Evo2 token features. NAIAD residue states then
    cross-attend to those query outputs and receive a gated residual update.
    """

    def __init__(
        self,
        evo_dim: int,
        naiad_dim: int,
        num_queries: int = 16,
        num_heads: int = 4,
        num_layers: int = 2,
        dropout: float = 0.0,
    ):
        super().__init__()
        self.evo_to_naiad = nn.Linear(evo_dim, naiad_dim)
        self.query_tokens = nn.Parameter(torch.randn(num_queries, naiad_dim) * 0.02)
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=naiad_dim,
            nhead=num_heads,
            dim_feedforward=naiad_dim * 4,
            dropout=dropout,
            batch_first=True,
            activation="gelu",
            norm_first=True,
        )
        self.qformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        self.residue_cross_attn = nn.MultiheadAttention(
            embed_dim=naiad_dim,
            num_heads=num_heads,
            dropout=dropout,
            batch_first=True,
        )
        self.norm = nn.LayerNorm(naiad_dim)
        self.gate = nn.Parameter(torch.tensor(-4.0))
        self.context_gate = nn.Linear(POLYMER_CONTEXT_DIM, 1, bias=False)
        self.context_scale = nn.Linear(POLYMER_CONTEXT_DIM, naiad_dim, bias=False)
        self.context_adapters = nn.ModuleList(
            ContextLowRankAdapter(naiad_dim) for _ in range(POLYMER_CONTEXT_DIM)
        )
        self.context_logit_adapters = nn.ModuleList(
            ContextLowRankLogitAdapter(naiad_dim) for _ in range(POLYMER_CONTEXT_DIM)
        )
        nn.init.zeros_(self.context_gate.weight)
        nn.init.zeros_(self.context_scale.weight)
        self.num_summary_tokens = int(num_queries)
        self.context_dim = POLYMER_CONTEXT_DIM
        self.context_adapter_rank = CONTEXT_ADAPTER_RANK

    def logit_delta(
        self,
        hidden: torch.Tensor,
        context_features: Optional[torch.Tensor],
    ) -> torch.Tensor:
        if context_features is None:
            return hidden.new_zeros((*hidden.shape[:-1], NAIAD_VOCAB_SIZE))
        return context_lowrank_logit_delta(
            self.context_logit_adapters,
            hidden,
            context_features,
        )

    def forward(
        self,
        residue_hidden: torch.Tensor,
        evo_hidden: torch.Tensor,
        residue_mask: Optional[torch.Tensor] = None,
        evo_padding_mask: Optional[torch.Tensor] = None,
        context_features: Optional[torch.Tensor] = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        evo_hidden = evo_hidden.to(dtype=self.evo_to_naiad.weight.dtype)
        evo_memory = self.evo_to_naiad(evo_hidden)
        queries = self.query_tokens.unsqueeze(0).expand(evo_hidden.shape[0], -1, -1)
        query_hidden = self.qformer(
            tgt=queries,
            memory=evo_memory,
            memory_key_padding_mask=evo_padding_mask,
        )
        update, _ = self.residue_cross_attn(
            query=residue_hidden,
            key=query_hidden,
            value=query_hidden,
            need_weights=False,
        )
        # Keep the residual path identity-safe: when the gate is near zero, the
        # wrapped model should reduce to the frozen NAIAD hidden state. Applying
        # LayerNorm after the residual changes NAIAD even with a closed gate.
        gate_logit = self.gate
        channel_scale = 1.0
        normalized_update = self.norm(update)
        adapter_delta = 0.0
        if context_features is not None:
            context_features = context_features.to(
                device=update.device,
                dtype=self.context_gate.weight.dtype,
            )
            gate_logit = gate_logit + self.context_gate(context_features).unsqueeze(-1)
            channel_scale = 1.0 + torch.tanh(self.context_scale(context_features)).unsqueeze(1)
            adapter_delta = context_lowrank_delta(
                self.context_adapters,
                normalized_update,
                context_features,
            )
        gate_value = torch.sigmoid(gate_logit)
        fused = residue_hidden + gate_value * channel_scale * normalized_update
        fused = fused + gate_value * adapter_delta
        if residue_mask is not None:
            fused = fused * residue_mask.unsqueeze(-1).to(dtype=fused.dtype)
        return fused, query_hidden


def load_bridge_state_compat(
    bridge: nn.Module,
    state_dict: dict[str, torch.Tensor],
    strict: bool = True,
):
    """Load legacy adapters while allowing only zero-init context routing keys to be absent."""
    result = bridge.load_state_dict(state_dict, strict=False)
    allowed_missing = {"context_gate.weight", "context_scale.weight"}
    disallowed_missing = {
        key
        for key in result.missing_keys
        if key not in allowed_missing
        and not key.startswith("context_adapters.")
        and not key.startswith("context_logit_adapters.")
    }
    if strict and (disallowed_missing or result.unexpected_keys):
        raise RuntimeError(
            "Incompatible bridge state: "
            f"missing={sorted(disallowed_missing)} unexpected={sorted(result.unexpected_keys)}"
        )
    return result