modeling_lemon.py

"""
Protein sequence encoder: transformer with RoPE, SwiGLU FFN, attention pooling,
and a bottleneck projector for contrastive sequence similarity.

No external dependencies beyond PyTorch.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint as grad_checkpoint


# ── Positional Embeddings ─────────────────────────────────────────────────

class RotaryPositionalEmbedding(nn.Module):
    def __init__(self, dim: int, max_tokens: int = 512, base: int = 10000,
                 scaling_type=None):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self.max_tokens = max_tokens
        self.scaling_type = scaling_type

    def _scaled_positions(self, seq_len, device, dtype):
        positions = torch.arange(seq_len, device=device, dtype=dtype)
        if seq_len <= self.max_tokens or not self.scaling_type:
            return positions
        if self.scaling_type == "linear":
            return positions * (self.max_tokens / seq_len)
        raise ValueError(f"Unknown scaling_type: {self.scaling_type}")

    def forward(self, x):
        t = self._scaled_positions(x.shape[1], x.device, self.inv_freq.dtype)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        return emb[None, :, :]


def rotate_half(x):
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    return torch.cat([-x2, x1], dim=-1)

def apply_rotary_pos_emb(q, k, pos_emb):
    cos, sin = pos_emb.cos(), pos_emb.sin()
    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)


class ScaledPositionalEmbedding(nn.Module):
    def __init__(self, num_embeddings, embedding_dim, extend_strategy="alibi",
                 padding_idx=None):
        super().__init__()
        self.embedding = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx)
        self.extend_strategy = extend_strategy
        self.num_embeddings = num_embeddings
        self.padding_idx = padding_idx
        max_index = num_embeddings - 1
        if padding_idx is not None:
            content_max_index = max(0, min(max_index - 1, padding_idx - 1))
        else:
            content_max_index = max_index
        self.register_buffer("max_index", torch.tensor(max_index), persistent=False)
        self.register_buffer("content_max_index", torch.tensor(content_max_index), persistent=False)
        if extend_strategy == "alibi":
            slopes = torch.logspace(0, -3, steps=embedding_dim, base=2.0)
            self.register_buffer("alibi_slopes", slopes, persistent=False)
        else:
            self.register_buffer("alibi_slopes", None, persistent=False)

    def forward(self, positions):
        positions = positions.long()
        if self.padding_idx is not None:
            pad_mask = positions == self.padding_idx
            clamped = positions.clamp(0, int(self.content_max_index.item()))
            clamped = torch.where(pad_mask, torch.full_like(clamped, self.padding_idx), clamped)
        else:
            pad_mask = torch.zeros_like(positions, dtype=torch.bool)
            clamped = positions.clamp(0, int(self.max_index.item()))
        embeddings = self.embedding(clamped)
        if self.extend_strategy == "alibi":
            excess = (positions - self.content_max_index).clamp_min(0)
            if self.padding_idx is not None:
                excess = torch.where(pad_mask, torch.zeros_like(excess), excess)
            embeddings = embeddings + excess.unsqueeze(-1) * self.alibi_slopes
        if self.padding_idx is not None:
            embeddings = embeddings.masked_fill(pad_mask.unsqueeze(-1), 0.0)
        return embeddings


# ── Transformer Block ─────────────────────────────────────────────────────

class AttentionBlock(nn.Module):
    def __init__(self, embed_dim, nhead, ff_mult=4, dropout=0.1):
        super().__init__()
        self.embed_dim = embed_dim
        self.nhead = nhead
        self.head_dim = embed_dim // nhead
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.o_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        ff_dim = int(embed_dim * ff_mult * 2 / 3)
        self.w1 = nn.Linear(embed_dim, ff_dim, bias=False)
        self.w2 = nn.Linear(embed_dim, ff_dim, bias=False)
        self.w3 = nn.Linear(ff_dim, embed_dim, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None, rope_emb=None):
        B, L, _ = x.shape
        if mask is not None:
            mask = mask.to(torch.bool)
        h = self.norm1(x)
        if mask is not None:
            h = h.masked_fill(~mask.unsqueeze(-1), 0.0)
        q = self.q_proj(h).view(B, L, self.nhead, self.head_dim).transpose(1, 2)
        k = self.k_proj(h).view(B, L, self.nhead, self.head_dim).transpose(1, 2)
        v = self.v_proj(h).view(B, L, self.nhead, self.head_dim).transpose(1, 2)
        if mask is not None:
            m = mask.unsqueeze(1).unsqueeze(-1)
            k = k.masked_fill(~m, 0.0)
            v = v.masked_fill(~m, 0.0)
        if rope_emb is not None:
            q, k = apply_rotary_pos_emb(q, k, rope_emb.unsqueeze(1))
        out = F.scaled_dot_product_attention(
            q, k, v, attn_mask=None,
            dropout_p=self.dropout.p if self.training else 0.0,
            is_causal=False,
        )
        out = out.transpose(1, 2).contiguous().view(B, L, self.embed_dim)
        out = self.o_proj(out)
        if mask is not None:
            out = out.masked_fill(~mask.unsqueeze(-1), 0.0)
        x = x + self.dropout(out)
        h = self.norm2(x)
        if mask is not None:
            h = h.masked_fill(~mask.unsqueeze(-1), 0.0)
        ff = self.w3(F.silu(self.w1(h)) * self.w2(h))
        if mask is not None:
            ff = ff.masked_fill(~mask.unsqueeze(-1), 0.0)
        x = x + self.dropout(ff)
        return x


# ── Encoder Stack ─────────────────────────────────────────────────────────

class Encoder(nn.Module):
    def __init__(self, vocab_size=32000, embed_dim=256, num_layers=4, nhead=8,
                 ff_mult=4, dropout=0.1, max_tokens=1024, padding_idx=0,
                 rope_scaling_type=None):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_layers = num_layers
        self.gradient_checkpointing = False
        self.token_emb = nn.Embedding(vocab_size, embed_dim, padding_idx=padding_idx)
        self.emb_dropout = nn.Dropout(dropout)
        self.rope = RotaryPositionalEmbedding(
            embed_dim // nhead, max_tokens, scaling_type=rope_scaling_type)
        self.layers = nn.ModuleList([
            AttentionBlock(embed_dim, nhead, ff_mult, dropout) for _ in range(num_layers)
        ])
        self.final_norm = nn.LayerNorm(embed_dim)
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Embedding):
                nn.init.normal_(m.weight, std=0.02)
                if m.padding_idx is not None:
                    m.weight.data[m.padding_idx].zero_()

    def forward(self, tokens, mask=None):
        x = self.emb_dropout(self.token_emb(tokens))
        padding_mask = mask.to(torch.bool) if mask is not None else None
        rope_emb = self.rope(x)
        for layer in self.layers:
            if self.gradient_checkpointing and self.training:
                x = grad_checkpoint(layer, x, padding_mask, rope_emb, use_reentrant=False)
            else:
                x = layer(x, padding_mask, rope_emb)
        x = self.final_norm(x)
        if padding_mask is not None:
            x = x.masked_fill(~padding_mask.unsqueeze(-1), 0.0)
        return x


# ── Attention Pooling ─────────────────────────────────────────────────────

class AttentionAggregator(nn.Module):
    def __init__(self, embed_dim, num_heads=4, dropout=0.1):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.dropout = dropout
        self.pool_query = nn.Parameter(torch.randn(1, 1, embed_dim) * 0.02)
        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.o_proj = nn.Linear(embed_dim, embed_dim, bias=False)
        self.norm = nn.LayerNorm(embed_dim)
        self.align = nn.Linear(embed_dim, embed_dim, bias=False)
        nn.init.eye_(self.align.weight)

    def forward(self, token_embs, mask=None):
        B, L, _ = token_embs.shape
        H, D = self.num_heads, self.head_dim
        query = self.pool_query.expand(B, -1, -1)
        q = self.q_proj(query).view(B, 1, H, D).transpose(1, 2)
        k = self.k_proj(token_embs).view(B, L, H, D).transpose(1, 2)
        v = self.v_proj(token_embs).view(B, L, H, D).transpose(1, 2)
        attn_mask = None
        if mask is not None:
            bool_mask = mask.to(torch.bool)
            attn_mask = torch.zeros(B, 1, 1, L, device=token_embs.device, dtype=token_embs.dtype)
            attn_mask = attn_mask.masked_fill(~bool_mask.unsqueeze(1).unsqueeze(2), float("-inf"))
            m = bool_mask.unsqueeze(1).unsqueeze(-1)
            k = k.masked_fill(~m, 0.0)
            v = v.masked_fill(~m, 0.0)
        out = F.scaled_dot_product_attention(
            q, k, v, attn_mask=attn_mask,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=False,
        )
        out = out.transpose(1, 2).contiguous().view(B, 1, self.embed_dim)
        pooled = self.norm(self.o_proj(out)).squeeze(1)
        return self.align(pooled)


# ── Residue Expansion Head ────────────────────────────────────────────────

class ExpansionHead(nn.Module):
    def __init__(self, embed_dim, max_residues, vocab_size=20, padding_index=31,
                 local_extend_strategy="alibi"):
        super().__init__()
        self.max_residues = max_residues
        self.PADDING_INDEX = padding_index
        self.local_pos_emb = ScaledPositionalEmbedding(
            num_embeddings=self.PADDING_INDEX + 1, embedding_dim=embed_dim,
            extend_strategy=local_extend_strategy, padding_idx=self.PADDING_INDEX,
        )
        self.residue_expansion = self._bottleneck_mlp(embed_dim, embed_dim, 0.1, 2)
        self.aa_logits_proj = nn.Linear(embed_dim, vocab_size)

    def _bottleneck_mlp(self, input_dim, hidden, dropout, num_layers, output_dim=None):
        def lng(i, h, d):
            return [nn.Linear(i, h), nn.LayerNorm(h), nn.GELU(), nn.Dropout(d)]
        layers = []
        for _ in range(num_layers):
            layers.extend(lng(input_dim, hidden, dropout) + lng(hidden, input_dim, dropout))
        if output_dim is not None:
            layers.append(nn.Linear(input_dim, output_dim))
        return nn.Sequential(*layers)

    def _expand_tokens(self, z, token_lengths, target_len):
        B, T, D = z.shape
        device = z.device
        if not torch.is_tensor(token_lengths):
            token_lengths = torch.tensor(token_lengths, device=device, dtype=torch.long)
        token_lengths = torch.clamp(token_lengths.clone(), min=0)
        cumsum = torch.cumsum(token_lengths, dim=1)
        total_residues = cumsum[:, -1]
        positions = torch.arange(target_len, device=device).unsqueeze(0).expand(B, -1)
        token_indices = torch.clamp(torch.searchsorted(cumsum, positions + 1), 0, T - 1)
        z_expanded = torch.gather(z, 1, token_indices.unsqueeze(-1).expand(-1, -1, D))
        cumsum_shifted = F.pad(cumsum[:, :-1], (1, 0), value=0)
        local_indices = positions - torch.gather(cumsum_shifted, 1, token_indices)
        padding_mask = positions >= total_residues.unsqueeze(1)
        z_expanded = z_expanded.masked_fill(padding_mask.unsqueeze(-1), 0.0)
        local_indices = torch.where(
            padding_mask, torch.full_like(local_indices, self.PADDING_INDEX),
            local_indices.clamp(min=0))
        return z_expanded, local_indices

    def forward(self, z, token_lengths, global_indices, global_pos_emb):
        target_len = global_indices.shape[1]
        z_exp, local_idx = self._expand_tokens(z, token_lengths, target_len)
        pos_global = global_pos_emb(torch.clamp(global_indices, min=0)).to(z_exp.dtype)
        pos_local = self.local_pos_emb(local_idx).to(z_exp.dtype)
        z_final = z_exp + pos_global + pos_local
        return self.aa_logits_proj(self.residue_expansion(z_final))


# ── Main Encoder Model ────────────────────────────────────────────────────

class LemonEncoder(nn.Module):
    """
    Protein sequence encoder: trie-tokenised input → per-token embeddings →
    attention-pooled sequence embedding → L2-normalised projector output.

    Suitable for contrastive similarity search (family / fold retrieval).

    Loading
    -------
    >>> from modeling_protein_encoder import LemonEncoder
    >>> model = LemonEncoder.from_pretrained("model.safetensors",
    ...                                                 config_path="config.json")
    >>> model.eval()
    """

    def __init__(
        self,
        vocab_size: int = 32000,
        embed_dim: int = 256,
        num_layers: int = 8,
        nhead: int = 8,
        ff_mult: int = 4,
        dropout: float = 0.1,
        max_tokens: int = 1024,
        proj_dim: int = 128,
        padding_idx: int = 0,
        rope_scaling_type=None,
        num_proj_layers: int = 2,
        proj_ff_mult: int = 2,
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.max_tokens = max_tokens
        self.embed_dim = embed_dim
        self.proj_dim = proj_dim or embed_dim
        self.num_proj_layers = num_proj_layers
        self.proj_ff_mult = proj_ff_mult

        self.core = Encoder(
            vocab_size=vocab_size, embed_dim=embed_dim, num_layers=num_layers,
            nhead=nhead, ff_mult=ff_mult, dropout=dropout, max_tokens=max_tokens,
            padding_idx=padding_idx, rope_scaling_type=rope_scaling_type,
        )
        self.log_temperature = nn.Parameter(torch.tensor(0.07).log())
        self.aggregator = AttentionAggregator(embed_dim, dropout=dropout)
        self.profile_expansion_head = ExpansionHead(embed_dim, max_residues=3 * max_tokens)
        self.global_pos_emb = nn.Embedding(3 * max_tokens, embed_dim)

        hidden = embed_dim * proj_ff_mult
        self.projector = self._bottleneck_mlp(embed_dim, hidden, dropout,
                                              num_proj_layers, output_dim=proj_dim)
        self.config = dict(
            vocab_size=vocab_size, embed_dim=embed_dim, num_layers=num_layers,
            nhead=nhead, ff_mult=ff_mult, dropout=dropout, max_tokens=max_tokens,
            proj_dim=proj_dim, padding_idx=padding_idx,
            rope_scaling_type=rope_scaling_type,
            num_proj_layers=num_proj_layers, proj_ff_mult=proj_ff_mult,
        )

    def _bottleneck_mlp(self, input_dim, hidden, dropout, num_layers, output_dim=None):
        def lng(i, h, d):
            return [nn.Linear(i, h), nn.LayerNorm(h), nn.GELU(), nn.Dropout(d)]
        layers = []
        for _ in range(num_layers):
            layers.extend(lng(input_dim, hidden, dropout) + lng(hidden, input_dim, dropout))
        if output_dim is not None:
            layers.append(nn.Linear(input_dim, output_dim))
        return nn.Sequential(*layers)

    @property
    def temperature(self):
        return self.log_temperature.exp().clamp(min=0.01, max=1.0)

    def encode_tokens(self, tokens, mask=None):
        return self.core(tokens, mask)

    def embed(self, tokens, mask=None, normalize=True):
        """Encode tokens → L2-normalised sequence embedding [B, proj_dim]."""
        token_embs = self.encode_tokens(tokens, mask)
        agg = self.aggregator(token_embs, mask)
        proj = self.projector(agg)
        if normalize:
            proj = F.normalize(proj, p=2, dim=-1)
        return proj

    def similarity(self, emb_a, emb_b):
        """Temperature-scaled cosine similarity matrix [B_a, B_b]."""
        return (emb_a @ emb_b.T) / self.temperature

    def forward(self, tokens, mask=None):
        token_embs = self.encode_tokens(tokens, mask)
        return F.linear(token_embs, self.core.token_emb.weight)

    @classmethod
    def from_pretrained(cls, weights_path: str, config_path: str = None,
                        device: str = "cpu"):
        """Load from safetensors weights + JSON config."""
        import json, os
        from safetensors.torch import load_file

        if config_path is None:
            config_path = os.path.join(os.path.dirname(weights_path), "config.json")

        with open(config_path) as f:
            cfg = json.load(f)

        arch = cfg.get("architecture", cfg)
        model = cls(**arch)
        state = load_file(weights_path, device=device)
        model.load_state_dict(state)
        return model