Add files using upload-large-folder tool

7ed266e verified about 1 month ago

12.5 kB

"""ESM-2 encoder and Global branch for LoGoPPI."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Mapping

import torch
from torch import Tensor, nn
from torch.nn import functional as F
from transformers import AutoModel, EsmConfig, EsmModel, PreTrainedModel
from transformers.utils import ModelOutput

from .config import ESM2PPIConfig


@dataclass
class ESM2PPIOutput(ModelOutput):
    """Model output with logits [B] and pooled representations [B, D]."""

    loss: Tensor | None = None
    logits: Tensor | None = None
    logits_ab: Tensor | None = None
    logits_ba: Tensor | None = None
    pooled_a: Tensor | None = None
    pooled_b: Tensor | None = None


class ESM2ForPPI(PreTrainedModel):
    """Encode two proteins with ESM-2 and average their AB and BA logits.
    Evaluation mode computes symmetric scores for the input order.
    In training mode, repeated calls may differ due to dropout.
    """

    config_class = ESM2PPIConfig
    base_model_prefix = "encoder"
    main_input_name = "input_a"
    supports_gradient_checkpointing = True

    def __init__(
        self,
        config: ESM2PPIConfig,
        encoder: nn.Module | None = None,
    ) -> None:
        super().__init__(config)

        # Build the encoder from the stored configuration when none is supplied.
        if encoder is None:
            if config.encoder_config is not None:
                # This creates the architecture without loading pretrained weights.
                esm_config = EsmConfig.from_dict(config.encoder_config)
                encoder = EsmModel(
                    esm_config,
                    add_pooling_layer=False,
                )
            else:
                encoder = AutoModel.from_pretrained(
                    config.base_model_name,
                    revision=config.base_model_revision,
                    add_pooling_layer=False,
                )
                config.encoder_config = encoder.config.to_dict()

        encoder_config = getattr(encoder, "config", None)
        hidden_size = getattr(encoder_config, "hidden_size", None)

        if not isinstance(hidden_size, int) or hidden_size <= 0:
            raise ValueError("encoder.config.hidden_size must be positive")

        self.encoder = encoder

        # Project each residue representation to the shared embedding dimension.
        self.projection = nn.Linear(
            hidden_size,
            config.projection_dim,
        )
        self.projection_dropout = nn.Dropout(config.projection_dropout)

        # Classifier input: [pooled_a, pooled_b, |pooled_a - pooled_b|].
        pair_dim = 3 * config.projection_dim

        self.global_layer_norm = nn.LayerNorm(pair_dim)
        self.global_mlp = nn.Sequential(
            nn.Linear(pair_dim, config.global_hidden_dim),
            nn.ReLU(),
            nn.Dropout(config.classifier_dropout),
            nn.Linear(config.global_hidden_dim, 1),
        )

        # Store the fixed loss weight with the model and move it across devices.
        self.register_buffer(
            "pos_weight",
            torch.tensor([config.pos_weight], dtype=torch.float32),
        )

        self._freeze_unused_encoder_parameters()
        self._initialize_added_modules()

    def _freeze_unused_encoder_parameters(self) -> None:
        """Freeze encoder parameters that are not used for PPI scoring."""
        unused_names = {
            "pooler.dense.weight",
            "pooler.dense.bias",
            "contact_head.regression.weight",
            "contact_head.regression.bias",
        }

        position_type = getattr(
            self.encoder.config,
            "position_embedding_type",
            None,
        )

        if position_type == "rotary":
            unused_names.add("embeddings.position_embeddings.weight")

        for name, parameter in self.encoder.named_parameters():
            if name in unused_names:
                parameter.requires_grad_(False)

    def _initialize_added_modules(self) -> None:
        """Initialize the projection and Global classifier layers."""
        added_modules = (
            self.projection,
            self.global_layer_norm,
            self.global_mlp,
        )

        for module in added_modules:
            module.apply(self._init_weights)

    @staticmethod
    def _validate_input(batch: Mapping[str, Tensor], name: str) -> None:
        """Check that required token tensors exist and have matching shapes."""
        required_keys = {
            "input_ids",
            "attention_mask",
            "special_tokens_mask",
        }

        missing_keys = required_keys.difference(batch.keys())

        if missing_keys:
            raise KeyError(f"{name} missing {sorted(missing_keys)}")

        token_shape = batch["input_ids"].shape
        attention_shape = batch["attention_mask"].shape
        special_shape = batch["special_tokens_mask"].shape

        if token_shape != attention_shape or token_shape != special_shape:
            raise ValueError(f"{name} token/mask shapes differ")

    @staticmethod
    def residue_mask(batch: Mapping[str, Tensor]) -> Tensor:
        """Return residue positions [B, L], excluding padding and special tokens."""
        valid_tokens = batch["attention_mask"].bool()
        special_tokens = batch["special_tokens_mask"].bool()

        mask = valid_tokens & ~special_tokens

        has_residue = mask.any(dim=1)

        if not has_residue.all().item():
            raise ValueError(
                "every protein must contain at least one residue token"
            )

        return mask

    def encode(
        self,
        batch: Mapping[str, Tensor],
    ) -> tuple[Tensor, Tensor]:
        """Return projected residue representations [B, L, D] and their mask."""
        self._validate_input(batch, "protein batch")

        encoder_output = self.encoder(
            input_ids=batch["input_ids"],
            attention_mask=batch["attention_mask"],
        )

        embeddings = encoder_output.last_hidden_state
        projected = self.projection(embeddings)
        projected = self.projection_dropout(projected)

        mask = self.residue_mask(batch)

        return projected, mask

    @staticmethod
    def mean_pool(
        projected: Tensor,
        residue_mask: Tensor,
    ) -> Tensor:
        """Average true residue positions into protein representations [B, D]."""
        if projected.ndim != 3:
            raise ValueError("projected embeddings must have shape [B, L, D]")

        if residue_mask.shape != projected.shape[:2]:
            raise ValueError("projected embeddings and residue mask shapes differ")

        # [B, L] -> [B, L, 1] applies one mask value to every residue dimension.
        weights = residue_mask.to(dtype=projected.dtype)
        weights = weights.unsqueeze(-1)

        residue_sum = (projected * weights).sum(dim=1)
        residue_count = weights.sum(dim=1).clamp_min(1)

        return residue_sum / residue_count

    def _ordered_logit(
        self,
        pooled_a: Tensor,
        pooled_b: Tensor,
    ) -> Tensor:
        """Compute the Global logit for one protein order."""
        difference = (pooled_a - pooled_b).abs()

        features = torch.cat(
            [pooled_a, pooled_b, difference],
            dim=-1,
        )

        features = self.global_layer_norm(features)
        logits = self.global_mlp(features)

        return logits.squeeze(-1)

    def score_projected(
        self,
        projected_a: Tensor,
        residue_mask_a: Tensor,
        projected_b: Tensor,
        residue_mask_b: Tensor,
        labels: Tensor | None = None,
        return_dict: bool = True,
    ) -> ESM2PPIOutput | tuple[Tensor, ...]:
        """Compute Global logits and optional loss from projected residues."""
        pooled_a = self.mean_pool(projected_a, residue_mask_a)
        pooled_b = self.mean_pool(projected_b, residue_mask_b)

        # Apply the shared classifier to both protein orders.
        logits_ab = self._ordered_logit(pooled_a, pooled_b)
        logits_ba = self._ordered_logit(pooled_b, pooled_a)

        logits = (logits_ab + logits_ba) * 0.5

        loss = None

        if labels is not None:
            targets = labels.to(
                device=logits.device,
                dtype=logits.dtype,
            )
            targets = targets.reshape_as(logits)

            pos_weight = self.pos_weight.to(
                device=logits.device,
                dtype=logits.dtype,
            )

            loss = F.binary_cross_entropy_with_logits(
                logits,
                targets,
                pos_weight=pos_weight,
            )

        if return_dict:
            return ESM2PPIOutput(
                loss=loss,
                logits=logits,
                logits_ab=logits_ab,
                logits_ba=logits_ba,
                pooled_a=pooled_a,
                pooled_b=pooled_b,
            )

        outputs = (
            logits,
            logits_ab,
            logits_ba,
            pooled_a,
            pooled_b,
        )

        if loss is not None:
            return (loss,) + outputs

        return outputs

    def forward(
        self,
        input_a: Mapping[str, Tensor],
        input_b: Mapping[str, Tensor],
        labels: Tensor | None = None,
        return_dict: bool | None = None,
    ) -> ESM2PPIOutput | tuple[Tensor, ...]:
        """Encode two proteins and compute their symmetric Global logit."""
        if return_dict is None:
            return_dict = self.config.use_return_dict

        projected_a, mask_a = self.encode(input_a)
        projected_b, mask_b = self.encode(input_b)

        return self.score_projected(
            projected_a=projected_a,
            residue_mask_a=mask_a,
            projected_b=projected_b,
            residue_mask_b=mask_b,
            labels=labels,
            return_dict=return_dict,
        )


class GlobalHead(nn.Module):
    """Compute Global logits from stored projected residue representations."""

    def __init__(
        self,
        dimension: int = 512,
        hidden_dimension: int = 512,
    ) -> None:
        super().__init__()

        pair_dimension = 3 * dimension

        self.layer_norm = nn.LayerNorm(pair_dimension)
        self.mlp = nn.Sequential(
            nn.Linear(pair_dimension, hidden_dimension),
            nn.ReLU(),
            nn.Dropout(0.0),
            nn.Linear(hidden_dimension, 1),
        )

    @classmethod
    def from_step1_state(
        cls,
        state: Mapping[str, Tensor],
    ) -> GlobalHead:
        """Load only the Global classifier from a Step 1 checkpoint."""
        model = cls()
        head_state = {}

        for name, weight in state.items():
            if name.startswith("global_layer_norm."):
                suffix = name.removeprefix("global_layer_norm.")
                head_state[f"layer_norm.{suffix}"] = weight

            elif name.startswith("global_mlp."):
                suffix = name.removeprefix("global_mlp.")
                head_state[f"mlp.{suffix}"] = weight

        model.load_state_dict(head_state, strict=True)

        model.eval()
        model.requires_grad_(False)

        return model

    @staticmethod
    def mean_pool(embeddings: Tensor, mask: Tensor) -> Tensor:
        """Average cached residue representations in FP32."""
        embeddings = embeddings.float()
        weights = mask.to(dtype=torch.float32)
        weights = weights.unsqueeze(-1)

        residue_sum = (embeddings * weights).sum(dim=1)
        residue_count = weights.sum(dim=1).clamp_min(1)

        return residue_sum / residue_count

    def ordered(
        self,
        pooled_a: Tensor,
        pooled_b: Tensor,
    ) -> Tensor:
        """Compute the Global logit for one protein order."""
        difference = (pooled_a - pooled_b).abs()

        features = torch.cat(
            [pooled_a, pooled_b, difference],
            dim=-1,
        )

        features = self.layer_norm(features)
        logits = self.mlp(features)

        return logits.squeeze(-1)

    def forward(
        self,
        a: Tensor,
        mask_a: Tensor,
        b: Tensor,
        mask_b: Tensor,
    ) -> Tensor:
        pooled_a = self.mean_pool(a, mask_a)
        pooled_b = self.mean_pool(b, mask_b)

        logits_ab = self.ordered(pooled_a, pooled_b)
        logits_ba = self.ordered(pooled_b, pooled_a)

        return (logits_ab + logits_ba) * 0.5

	"""ESM-2 encoder and Global branch for LoGoPPI."""

	from __future__ import annotations

	from dataclasses import dataclass
	from typing import Mapping

	import torch
	from torch import Tensor, nn
	from torch.nn import functional as F
	from transformers import AutoModel, EsmConfig, EsmModel, PreTrainedModel
	from transformers.utils import ModelOutput

	from .config import ESM2PPIConfig


	@dataclass
	class ESM2PPIOutput(ModelOutput):
	"""Model output with logits [B] and pooled representations [B, D]."""

	loss: Tensor \| None = None
	logits: Tensor \| None = None
	logits_ab: Tensor \| None = None
	logits_ba: Tensor \| None = None
	pooled_a: Tensor \| None = None
	pooled_b: Tensor \| None = None


	class ESM2ForPPI(PreTrainedModel):
	"""Encode two proteins with ESM-2 and average their AB and BA logits.

	Evaluation mode computes symmetric scores for the input order.
	In training mode, repeated calls may differ due to dropout.
	"""

	config_class = ESM2PPIConfig
	base_model_prefix = "encoder"
	main_input_name = "input_a"
	supports_gradient_checkpointing = True

	def __init__(
	self,
	config: ESM2PPIConfig,
	encoder: nn.Module \| None = None,
	) -> None:
	super().__init__(config)

	# Build the encoder from the stored configuration when none is supplied.
	if encoder is None:
	if config.encoder_config is not None:
	# This creates the architecture without loading pretrained weights.
	esm_config = EsmConfig.from_dict(config.encoder_config)
	encoder = EsmModel(
	esm_config,
	add_pooling_layer=False,
	)
	else:
	encoder = AutoModel.from_pretrained(
	config.base_model_name,
	revision=config.base_model_revision,
	add_pooling_layer=False,
	)
	config.encoder_config = encoder.config.to_dict()

	encoder_config = getattr(encoder, "config", None)
	hidden_size = getattr(encoder_config, "hidden_size", None)

	if not isinstance(hidden_size, int) or hidden_size <= 0:
	raise ValueError("encoder.config.hidden_size must be positive")

	self.encoder = encoder

	# Project each residue representation to the shared embedding dimension.
	self.projection = nn.Linear(
	hidden_size,
	config.projection_dim,
	)
	self.projection_dropout = nn.Dropout(config.projection_dropout)

	# Classifier input: [pooled_a, pooled_b, \|pooled_a - pooled_b\|].
	pair_dim = 3 * config.projection_dim

	self.global_layer_norm = nn.LayerNorm(pair_dim)
	self.global_mlp = nn.Sequential(
	nn.Linear(pair_dim, config.global_hidden_dim),
	nn.ReLU(),
	nn.Dropout(config.classifier_dropout),
	nn.Linear(config.global_hidden_dim, 1),
	)

	# Store the fixed loss weight with the model and move it across devices.
	self.register_buffer(
	"pos_weight",
	torch.tensor([config.pos_weight], dtype=torch.float32),
	)

	self._freeze_unused_encoder_parameters()
	self._initialize_added_modules()

	def _freeze_unused_encoder_parameters(self) -> None:
	"""Freeze encoder parameters that are not used for PPI scoring."""
	unused_names = {
	"pooler.dense.weight",
	"pooler.dense.bias",
	"contact_head.regression.weight",
	"contact_head.regression.bias",
	}

	position_type = getattr(
	self.encoder.config,
	"position_embedding_type",
	None,
	)

	if position_type == "rotary":
	unused_names.add("embeddings.position_embeddings.weight")

	for name, parameter in self.encoder.named_parameters():
	if name in unused_names:
	parameter.requires_grad_(False)

	def _initialize_added_modules(self) -> None:
	"""Initialize the projection and Global classifier layers."""
	added_modules = (
	self.projection,
	self.global_layer_norm,
	self.global_mlp,
	)

	for module in added_modules:
	module.apply(self._init_weights)

	@staticmethod
	def _validate_input(batch: Mapping[str, Tensor], name: str) -> None:
	"""Check that required token tensors exist and have matching shapes."""
	required_keys = {
	"input_ids",
	"attention_mask",
	"special_tokens_mask",
	}

	missing_keys = required_keys.difference(batch.keys())

	if missing_keys:
	raise KeyError(f"{name} missing {sorted(missing_keys)}")

	token_shape = batch["input_ids"].shape
	attention_shape = batch["attention_mask"].shape
	special_shape = batch["special_tokens_mask"].shape

	if token_shape != attention_shape or token_shape != special_shape:
	raise ValueError(f"{name} token/mask shapes differ")

	@staticmethod
	def residue_mask(batch: Mapping[str, Tensor]) -> Tensor:
	"""Return residue positions [B, L], excluding padding and special tokens."""
	valid_tokens = batch["attention_mask"].bool()
	special_tokens = batch["special_tokens_mask"].bool()

	mask = valid_tokens & ~special_tokens

	has_residue = mask.any(dim=1)

	if not has_residue.all().item():
	raise ValueError(
	"every protein must contain at least one residue token"
	)

	return mask

	def encode(
	self,
	batch: Mapping[str, Tensor],
	) -> tuple[Tensor, Tensor]:
	"""Return projected residue representations [B, L, D] and their mask."""
	self._validate_input(batch, "protein batch")

	encoder_output = self.encoder(
	input_ids=batch["input_ids"],
	attention_mask=batch["attention_mask"],
	)

	embeddings = encoder_output.last_hidden_state
	projected = self.projection(embeddings)
	projected = self.projection_dropout(projected)

	mask = self.residue_mask(batch)

	return projected, mask

	@staticmethod
	def mean_pool(
	projected: Tensor,
	residue_mask: Tensor,
	) -> Tensor:
	"""Average true residue positions into protein representations [B, D]."""
	if projected.ndim != 3:
	raise ValueError("projected embeddings must have shape [B, L, D]")

	if residue_mask.shape != projected.shape[:2]:
	raise ValueError("projected embeddings and residue mask shapes differ")

	# [B, L] -> [B, L, 1] applies one mask value to every residue dimension.
	weights = residue_mask.to(dtype=projected.dtype)
	weights = weights.unsqueeze(-1)

	residue_sum = (projected * weights).sum(dim=1)
	residue_count = weights.sum(dim=1).clamp_min(1)

	return residue_sum / residue_count

	def _ordered_logit(
	self,
	pooled_a: Tensor,
	pooled_b: Tensor,
	) -> Tensor:
	"""Compute the Global logit for one protein order."""
	difference = (pooled_a - pooled_b).abs()

	features = torch.cat(
	[pooled_a, pooled_b, difference],
	dim=-1,
	)

	features = self.global_layer_norm(features)
	logits = self.global_mlp(features)

	return logits.squeeze(-1)

	def score_projected(
	self,
	projected_a: Tensor,
	residue_mask_a: Tensor,
	projected_b: Tensor,
	residue_mask_b: Tensor,
	labels: Tensor \| None = None,
	return_dict: bool = True,
	) -> ESM2PPIOutput \| tuple[Tensor, ...]:
	"""Compute Global logits and optional loss from projected residues."""
	pooled_a = self.mean_pool(projected_a, residue_mask_a)
	pooled_b = self.mean_pool(projected_b, residue_mask_b)

	# Apply the shared classifier to both protein orders.
	logits_ab = self._ordered_logit(pooled_a, pooled_b)
	logits_ba = self._ordered_logit(pooled_b, pooled_a)

	logits = (logits_ab + logits_ba) * 0.5

	loss = None

	if labels is not None:
	targets = labels.to(
	device=logits.device,
	dtype=logits.dtype,
	)
	targets = targets.reshape_as(logits)

	pos_weight = self.pos_weight.to(
	device=logits.device,
	dtype=logits.dtype,
	)

	loss = F.binary_cross_entropy_with_logits(
	logits,
	targets,
	pos_weight=pos_weight,
	)

	if return_dict:
	return ESM2PPIOutput(
	loss=loss,
	logits=logits,
	logits_ab=logits_ab,
	logits_ba=logits_ba,
	pooled_a=pooled_a,
	pooled_b=pooled_b,
	)

	outputs = (
	logits,
	logits_ab,
	logits_ba,
	pooled_a,
	pooled_b,
	)

	if loss is not None:
	return (loss,) + outputs

	return outputs

	def forward(
	self,
	input_a: Mapping[str, Tensor],
	input_b: Mapping[str, Tensor],
	labels: Tensor \| None = None,
	return_dict: bool \| None = None,
	) -> ESM2PPIOutput \| tuple[Tensor, ...]:
	"""Encode two proteins and compute their symmetric Global logit."""
	if return_dict is None:
	return_dict = self.config.use_return_dict

	projected_a, mask_a = self.encode(input_a)
	projected_b, mask_b = self.encode(input_b)

	return self.score_projected(
	projected_a=projected_a,
	residue_mask_a=mask_a,
	projected_b=projected_b,
	residue_mask_b=mask_b,
	labels=labels,
	return_dict=return_dict,
	)


	class GlobalHead(nn.Module):
	"""Compute Global logits from stored projected residue representations."""

	def __init__(
	self,
	dimension: int = 512,
	hidden_dimension: int = 512,
	) -> None:
	super().__init__()

	pair_dimension = 3 * dimension

	self.layer_norm = nn.LayerNorm(pair_dimension)
	self.mlp = nn.Sequential(
	nn.Linear(pair_dimension, hidden_dimension),
	nn.ReLU(),
	nn.Dropout(0.0),
	nn.Linear(hidden_dimension, 1),
	)

	@classmethod
	def from_step1_state(
	cls,
	state: Mapping[str, Tensor],
	) -> GlobalHead:
	"""Load only the Global classifier from a Step 1 checkpoint."""
	model = cls()
	head_state = {}

	for name, weight in state.items():
	if name.startswith("global_layer_norm."):
	suffix = name.removeprefix("global_layer_norm.")
	head_state[f"layer_norm.{suffix}"] = weight

	elif name.startswith("global_mlp."):
	suffix = name.removeprefix("global_mlp.")
	head_state[f"mlp.{suffix}"] = weight

	model.load_state_dict(head_state, strict=True)

	model.eval()
	model.requires_grad_(False)

	return model

	@staticmethod
	def mean_pool(embeddings: Tensor, mask: Tensor) -> Tensor:
	"""Average cached residue representations in FP32."""
	embeddings = embeddings.float()
	weights = mask.to(dtype=torch.float32)
	weights = weights.unsqueeze(-1)

	residue_sum = (embeddings * weights).sum(dim=1)
	residue_count = weights.sum(dim=1).clamp_min(1)

	return residue_sum / residue_count

	def ordered(
	self,
	pooled_a: Tensor,
	pooled_b: Tensor,
	) -> Tensor:
	"""Compute the Global logit for one protein order."""
	difference = (pooled_a - pooled_b).abs()

	features = torch.cat(
	[pooled_a, pooled_b, difference],
	dim=-1,
	)

	features = self.layer_norm(features)
	logits = self.mlp(features)

	return logits.squeeze(-1)

	def forward(
	self,
	a: Tensor,
	mask_a: Tensor,
	b: Tensor,
	mask_b: Tensor,
	) -> Tensor:
	pooled_a = self.mean_pool(a, mask_a)
	pooled_b = self.mean_pool(b, mask_b)

	logits_ab = self.ordered(pooled_a, pooled_b)
	logits_ba = self.ordered(pooled_b, pooled_a)

	return (logits_ab + logits_ba) * 0.5