MSALab/PerceptionDLM-Base

1import re2import torch3from torch import nn4from torch.nn import functional as F5 6 7def build_projection(projection_type: str, in_dim: int, out_dim: int) -> nn.Module:8    mlp_gelu_match = re.match(r'^mlp(\d+)x_gelu$', projection_type)9    if mlp_gelu_match:10        mlp_depth = int(mlp_gelu_match.group(1))11        modules = [nn.Linear(in_dim, out_dim)]12        for _ in range(1, mlp_depth):13            modules.append(nn.GELU())14            modules.append(nn.Linear(out_dim, out_dim))15        projection = nn.Sequential(*modules)16        return projection17 18    raise ValueError(f'Unknown projector type: {projection_type}')19 20 21class PerceiverProjection(nn.Module):22    def __init__(self, projection_type: str, in_dim: int, out_dim: int):23        super().__init__()24        self.projection = build_projection(projection_type, in_dim, out_dim)25 26    def forward(self, input_embeds: torch.Tensor):27        input_embeds.requires_grad_(True)28        embeds = self.projection(input_embeds)29        embeds.requires_grad_(True)30        return embeds