Upload model

config.json

@@ -2,6 +2,10 @@
   "architectures": [
     "PhyloGPNModel"
   ],
+  "auto_map": {
+    "AutoConfig": "configuration_phylogpn.PhyloGPNConfig",
+    "AutoModel": "modeling_phylogpn.PhyloGPNModel"
+  },
   "inner_dim": 480,
   "kernel_size": 5,
   "model_type": "phylogpn",

configuration_phylogpn.py

@@ -0,0 +1,12 @@
+from transformers import PretrainedConfig
+
+class PhyloGPNConfig(PretrainedConfig):
+    model_type = "phylogpn"
+
+    def __init__(self, outer_dim: int = 960, inner_dim: int = 480, kernel_size: int = 5, stack_size: int = 2, num_stacks: int = 20, **kwargs):
+        self.outer_dim = outer_dim
+        self.inner_dim = inner_dim
+        self.kernel_size = kernel_size
+        self.stack_size = stack_size
+        self.num_stacks = num_stacks
+        super().__init__(**kwargs)

modeling_phylogpn.py

@@ -0,0 +1,251 @@
+from typing import List, Optional
+
+import torch
+from torch import nn
+from torch.nn.utils import parametrize
+
+
+def check_if_involution(indices: List[int]) -> bool:
+    return all(indices[indices[idx]] == idx for idx in range(len(indices)))
+
+
+def get_conv1d_output_length(
+    input_length: int, kernel_size: int, stride_size: int = 1, pad_size: int = 0, dilation_rate: int = 1
+) -> int:
+    return (input_length + 2 * pad_size - dilation_rate * (kernel_size - 1) - 1) // stride_size + 1
+
+
+def get_involution_indices(size: int) -> List[int]:
+    return list(reversed(range(size)))
+
+
+class RCEWeight(nn.Module):
+    def __init__(
+        self, input_involution_indices: List[int], output_involution_indices: List[int]
+    ):
+        if not check_if_involution(input_involution_indices) or not check_if_involution(
+                output_involution_indices):
+            raise ValueError(
+                "`input_involution_indices` and `output_involution_indices` must be involutions"
+            )
+
+        super().__init__()
+        self.input_involution_indices = input_involution_indices
+        self.output_involution_indices = output_involution_indices
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        output_involution_indices = torch.tensor(self.output_involution_indices, device=x.device)
+        input_involution_indices = torch.tensor(self.input_involution_indices, device=x.device)
+        return (x + x[output_involution_indices][:, input_involution_indices].flip(2)) / 2
+
+
+class IEBias(nn.Module):
+    def __init__(self, involution_indices: List[int]):
+        if not check_if_involution(involution_indices):
+            raise ValueError("`involution_indices` must be an involution")
+
+        super().__init__()
+        self.involution_indices = involution_indices
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        involution_indices = torch.tensor(self.involution_indices, device=x.device)
+        return (x + x[involution_indices]) / 2
+
+
+class IEWeight(nn.Module):
+    def __init__(
+        self, input_involution_indices: List[int], output_involution_indices: List[int]
+    ):
+        if not check_if_involution(input_involution_indices) or not check_if_involution(
+                output_involution_indices):
+            raise ValueError(
+                "`input_involution_indices` and `output_involution_indices` must be involutions"
+            )
+
+        super().__init__()
+        self.input_involution_indices = input_involution_indices
+        self.output_involution_indices = output_involution_indices
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_involution_indices = torch.tensor(self.input_involution_indices, device=x.device)
+        output_involution_indices = torch.tensor(self.output_involution_indices, device=x.device)
+        return (x + x[input_involution_indices][:, output_involution_indices]) / 2
+    
+
+class RCEByteNetBlock(nn.Module):
+    def __init__(
+        self,
+        outer_involution_indices: List[int],
+        inner_dim: int,
+        kernel_size: int,
+        dilation_rate: int = 1
+    ):
+        outer_dim = len(outer_involution_indices)
+
+        if outer_dim % 2 != 0:
+            raise ValueError("`outer_involution_indices` must have an even length")
+
+        if inner_dim % 2 != 0:
+            raise ValueError("`inner_dim` must be even")
+
+        if kernel_size % 2 == 0:
+            raise ValueError("`kernel_size` must be odd")
+
+        super().__init__()
+        inner_involution_indices = get_involution_indices(inner_dim)
+
+        layers = [
+            nn.GroupNorm(1, outer_dim),
+            nn.GELU(),
+            nn.Conv1d(outer_dim, inner_dim, kernel_size=1),
+            nn.GroupNorm(1, inner_dim),
+            nn.GELU(),
+            nn.Conv1d(inner_dim, inner_dim, kernel_size, dilation=dilation_rate),
+            nn.GroupNorm(1, inner_dim),
+            nn.GELU(),
+            nn.Conv1d(inner_dim, outer_dim, kernel_size=1)
+        ]
+        parametrize.register_parametrization(
+            layers[2], "weight",
+            RCEWeight(outer_involution_indices, inner_involution_indices)
+        )
+        parametrize.register_parametrization(
+            layers[2], "bias",
+            IEBias(inner_involution_indices)
+        )
+        parametrize.register_parametrization(
+            layers[5], "weight",
+            RCEWeight(inner_involution_indices, inner_involution_indices)
+        )
+        parametrize.register_parametrization(
+            layers[5], "bias",
+            IEBias(inner_involution_indices)
+        )
+        parametrize.register_parametrization(
+            layers[8], "weight",
+            RCEWeight(inner_involution_indices, outer_involution_indices)
+        )
+        parametrize.register_parametrization(
+            layers[8], "bias",
+            IEBias(outer_involution_indices)
+        )
+
+        self.layers = nn.Sequential(*layers)
+        self._kernel_size = kernel_size
+        self._dilation_rate = dilation_rate
+
+    @property
+    def kernel_size(self):
+        return self._kernel_size
+
+    @property
+    def dilation_rate(self):
+        return self._dilation_rate
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_length = x.shape[2]
+        output_length = get_conv1d_output_length(input_length, self.kernel_size, dilation_rate=self.dilation_rate)
+        a = (input_length - output_length) // 2
+        
+        if a == 0:
+            return self.layers(x) + x
+        
+        return self.layers(x) + x[:, :, a:-a]
+    
+class RCEByteNet(nn.Module):
+    def __init__(
+        self,
+        input_involution_indices: List[int],
+        output_involution_indices: List[int],
+        dilation_rates: List[int],
+        outer_dim: int,
+        inner_dim: int,
+        kernel_size: int,
+        pad_token_idx: Optional[int] = None,
+    ):
+        if pad_token_idx is not None and input_involution_indices[pad_token_idx] != pad_token_idx:
+            raise ValueError("`input_involution_indices[pad_token_idx]` must be equal to `pad_token_idx`")
+
+        super().__init__()
+        vocab_size = len(input_involution_indices)
+        outer_involution_indices = get_involution_indices(outer_dim)
+
+        self.embedding = nn.Embedding(vocab_size, outer_dim, padding_idx=pad_token_idx)
+        parametrize.register_parametrization(
+            self.embedding, "weight",
+            IEWeight(input_involution_indices, outer_involution_indices)
+        )
+        nn.init.normal_(self.embedding.weight, std=2**0.5)
+        self.embedding.weight.data[self.embedding.padding_idx].zero_()
+        self.embedding.requires_grad = False
+
+        blocks = []
+        receptive_field_size = 1
+
+        for r in dilation_rates:
+            blocks.append(RCEByteNetBlock(outer_involution_indices, inner_dim, kernel_size, dilation_rate=r))
+            receptive_field_size += (kernel_size - 1) * r
+
+        self.blocks = nn.Sequential(*blocks)
+
+        output_dim = len(output_involution_indices)
+        self.output_layers = nn.Sequential(
+            nn.GroupNorm(1, outer_dim), nn.GELU(), nn.Conv1d(outer_dim, output_dim, kernel_size=1)
+        )
+        parametrize.register_parametrization(
+            self.output_layers[-1], "weight",
+            RCEWeight(outer_involution_indices, output_involution_indices)
+        )
+        parametrize.register_parametrization(
+            self.output_layers[-1], "bias", IEBias(output_involution_indices)
+        )
+
+        self._embedding_involution_indices = outer_involution_indices
+
+    @property
+    def embedding_involution_indices(self):
+        return self._embedding_involution_indices
+
+    def get_embeddings(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        x = self.embedding(input_tensor).swapaxes(1, 2)
+        return self.output_layers[0](self.blocks(x)).swapaxes(1, 2)
+
+    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        x = self.get_embeddings(input_tensor).swapaxes(1, 2)
+        return self.output_layers[1:](x).swapaxes(1, 2)
+    
+
+from transformers import PreTrainedModel
+from .configuration_phylogpn import PhyloGPNConfig
+
+class PhyloGPNModel(PreTrainedModel):
+    config_class = PhyloGPNConfig
+
+    def __init__(self, config, **kwargs):
+        super().__init__(config, **kwargs)
+
+        dilation_rates = config.num_stacks * [config.kernel_size**i for i in range(0, config.stack_size)]
+
+        self._model = RCEByteNet(
+            input_involution_indices = [3, 2, 1, 0, 4, 5],
+            output_involution_indices=[3, 2, 1, 0],
+            dilation_rates=dilation_rates,
+            outer_dim = config.outer_dim,
+            inner_dim = config.inner_dim,
+            kernel_size=config.kernel_size,
+            pad_token_idx=5
+        )
+
+    def get_embeddings(self, input_ids: torch.Tensor):
+        return self._model.get_embeddings(input_ids)
+
+    def forward(self, input_ids: torch.Tensor):
+        output_tensor = self._model(input_ids)
+        output_array = output_tensor.numpy(force=True)
+
+        results = {}
+
+        for idx, key in enumerate("ACGT"):
+            results[key] = output_array[:, :, idx]
+
+        return results