code/qwenvl/modalities/mol/encoder.py

"""
Molecular encoder: GNN backbone + Perceiver Resampler.

Wraps the pretrained GNN from mol/graph_NN.py and compresses variable-length
graph node embeddings into a fixed number of latent tokens via cross-attention.

Forward signature is compatible with ModalityRouter:
    (input_ids, attention_mask, **kwargs) -> (latent [B,K,D], mask [B,K])

Graph data (edge_index, edge_attr, edge_index_all, batch_idx) is passed
through **kwargs. When kwargs are empty (dummy forward), returns zero tensors
connected to trainable parameters for gradient sync.
"""

import torch
import torch.nn as nn
from torch import Tensor

from .graph_NN import GNN


class _LatentResampler(nn.Module):
    """Perceiver-style cross-attention: variable-length nodes -> K latent tokens."""

    def __init__(self, dim: int, num_latent_tokens: int = 16, heads: int = 8, dropout: float = 0.1):
        super().__init__()
        self.latent_queries = nn.Parameter(torch.randn(1, num_latent_tokens, dim) * 0.02)
        self.norm_q = nn.LayerNorm(dim)
        self.norm_kv = nn.LayerNorm(dim)
        self.cross_attn = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
        ffn_dim = dim * 4
        self.ffn = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, ffn_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ffn_dim, dim),
            nn.Dropout(dropout),
        )

    def forward(self, encoder_out: Tensor, key_padding_mask: Tensor | None = None):
        """
        Args:
            encoder_out:     [B, max_nodes, D]
            key_padding_mask: [B, max_nodes]  True = ignore (padding)
        Returns:
            latent: [B, K, D]
        """
        B = encoder_out.shape[0]
        q = self.norm_q(self.latent_queries.expand(B, -1, -1))
        kv = self.norm_kv(encoder_out)
        h, _ = self.cross_attn(q, kv, kv, key_padding_mask=key_padding_mask)
        h = q + h
        h = h + self.ffn(h)
        return h


class MolEncoder(nn.Module):
    """GNN + Perceiver Resampler -> fixed K latent tokens.

    The GNN processes molecular graphs with variable numbers of atoms.
    The resampler compresses the per-atom embeddings into K fixed-length
    latent tokens suitable for scattering into the LLM embedding space.
    """

    def __init__(self, config):
        super().__init__()
        self.hidden_size = config.mol_encoder_hidden_size
        self.num_latent_tokens = config.num_latent_tokens
        self._gnn_frozen_cached: bool | None = None

        self.gnn = GNN(
            num_layer=config.num_gnn_layers,
            emb_dim=config.mol_encoder_hidden_size,
            drop_ratio=config.gnn_drop_ratio,
            output_type="last",
        )

        self.resampler = _LatentResampler(
            dim=config.mol_encoder_hidden_size,
            num_latent_tokens=config.num_latent_tokens,
            heads=config.num_resampler_heads,
            dropout=getattr(config, "dropout", 0.1),
        )

        # Initialize resampler weights
        self.resampler.apply(self._init_weights)

    @staticmethod
    def _init_weights(module: nn.Module):
        if isinstance(module, nn.Linear):
            nn.init.trunc_normal_(module.weight, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.LayerNorm):
            nn.init.ones_(module.weight)
            nn.init.zeros_(module.bias)

    def _is_gnn_frozen(self) -> bool:
        """Check if GNN has no trainable params (cached for performance)."""
        if self._gnn_frozen_cached is None:
            self._gnn_frozen_cached = not any(
                p.requires_grad for p in self.gnn.parameters()
            )
        return self._gnn_frozen_cached

    def forward(
        self,
        input_ids: Tensor,
        attention_mask: Tensor,
        **kwargs,
    ):
        """
        Args:
            input_ids:      Graph node features [total_nodes, 5] (real forward)
                            or [B, dummy_len] (dummy forward from ModalityRouter)
            attention_mask: Dummy (not used for graph encoder)
            **kwargs:       Graph-specific tensors:
                edge_index:     [2, num_edges]
                edge_attr:      [num_edges, 3]
                edge_index_all: [2, num_fc_edges]
                batch_idx:      [total_nodes] — molecule assignment per node

        Returns:
            latent:      [B, K, D]  — compressed latent tokens
            latent_mask: [B, K]     — all-ones mask (fixed length)
        """
        edge_index = kwargs.get("edge_index")

        if edge_index is None:
            # Fallback when no graph is provided.
            # Must produce a tensor connected to ALL trainable params.
            B = input_ids.shape[0]
            K = self.num_latent_tokens
            device = input_ids.device
            dtype = next(self.resampler.parameters()).dtype
            # Use latent_queries as KV so cross-attention weights also get gradients
            dummy_kv = self.resampler.latent_queries[:, :1, :].expand(B, -1, -1)
            latent = self.resampler(dummy_kv)  # [B, K, D] — all resampler params in graph

            # If GNN has trainable params (freeze_policy may have unfrozen them
            # after register_modality set freeze_mol_gnn=True), connect them too.
            gnn_trainable = [p for p in self.gnn.parameters() if p.requires_grad]
            if gnn_trainable:
                gnn_dummy = sum(p.sum() for p in gnn_trainable) * 0
                latent = latent + gnn_dummy

            latent_mask = torch.ones(B, K, device=device, dtype=torch.long)
            return latent, latent_mask

        # Real graph forward
        edge_attr = kwargs["edge_attr"]
        edge_index_all = kwargs["edge_index_all"]
        batch_idx = kwargs["batch_idx"]

        # GNN forward: [total_nodes, 5] -> [total_nodes, D]
        # Memory optimization: skip autograd graph for frozen GNN
        if self._is_gnn_frozen():
            with torch.no_grad():
                node_embeds = self.gnn(
                    input_ids, edge_index, edge_index_all, edge_attr, batch_idx
                )
                if isinstance(node_embeds, list):
                    node_embeds = node_embeds[-1]
                node_embeds = node_embeds.detach()
        else:
            node_embeds = self.gnn(
                input_ids, edge_index, edge_index_all, edge_attr, batch_idx
            )
            if isinstance(node_embeds, list):
                node_embeds = node_embeds[-1]

        # Group by molecule + pad -> [B, max_nodes, D]
        # Uses PyG's vectorized to_dense_batch (avoids Python for-loop)
        from torch_geometric.utils import to_dense_batch
        padded, real_mask = to_dense_batch(node_embeds, batch_idx)
        # to_dense_batch returns real_mask where True = real node
        # MultiheadAttention key_padding_mask expects True = ignore (padding)
        pad_mask = ~real_mask
        num_mols = padded.shape[0]

        # Resampler: [B, max_nodes, D] -> [B, K, D]
        latent = self.resampler(padded, key_padding_mask=pad_mask)
        K = latent.shape[1]
        latent_mask = torch.ones(num_mols, K, device=latent.device, dtype=torch.long)
        return latent, latent_mask