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