initial commit

ChemPepMTR.py

@@ -0,0 +1,154 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchtune.modules import RotaryPositionalEmbeddings
+from transformers import PreTrainedModel
+from .config import model_config
+from typing import Mapping
+from transformers.tokenization_utils_base import BatchEncoding
+
+class SwiGLU(nn.Module):
+    def __init__(self, input_dim, hidden_dim):
+        super().__init__()
+        self.linear1 = nn.Linear(input_dim, hidden_dim * 2, bias=True)
+        self.linear2 = nn.Linear(hidden_dim, input_dim, bias=True)
+        self.dropout = nn.Dropout(0.1)  # Add dropout for regularization
+    def forward(self, x):
+        # x: (N, input_dim)
+        x1, x2 = self.linear1(x).chunk(2, dim=-1)
+        output = self.linear2(F.silu(x1) * x2)
+        return self.dropout(output)
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads, max_seq_len):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+
+        self.qkv_proj = nn.Linear(embed_dim, embed_dim * 3, bias=False)
+        self.rotary = RotaryPositionalEmbeddings(dim=self.head_dim, max_seq_len=max_seq_len)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.dropout = nn.Dropout(0.1)  # Add dropout for regularization
+
+    def forward(self, x, input_pos=None, mask=None):
+        B, T, C = x.shape  # Batch, sequence, embedding dim
+        
+        # project into queries, keys, and values
+        q, k, v = self.qkv_proj(x).view(B, T, 3, self.num_heads, self.head_dim).unbind(2)  # (B, T, num_heads, head_dim)
+
+        # Apply rotary positional embeddings to queries and keys
+        q, k = self.rotary(q, input_pos=input_pos), self.rotary(k, input_pos=input_pos)
+
+        # Reshape to (B, num_heads, T, head_dim)
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        if mask is not None:
+            # set padding positions to -inf
+            mask = mask.to(dtype=torch.float32)  # Ensure mask is float
+            mask = (1.0 - mask) * -1e9  # Convert to -inf for padding positions
+            
+            # mask: (B, T) -> (B, 1, 1, T)
+            mask = mask.unsqueeze(1).unsqueeze(2)  # (B, 1, 1, T)
+            mask = mask.expand(B, 1, T, T)  # expands to (batch, 1, seqlen, seqlen)
+            
+        # Scaled dot-product attention
+        attn_output = F.scaled_dot_product_attention(query=q, key=k, value=v, attn_mask=mask)
+        attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, C)
+        attn_output = self.out_proj(attn_output)
+        return self.dropout(attn_output)
+
+
+class UnifiedTransformerBlock(nn.Module):
+    def __init__(self, embed_dim, num_heads, ffn_hidden_dim, max_seq_len):
+        super().__init__()
+        self.attn_norm = nn.LayerNorm(embed_dim)
+        self.attn = MultiHeadAttention(embed_dim, num_heads, max_seq_len)
+        self.ffn_norm = nn.LayerNorm(embed_dim)
+        self.ffn = SwiGLU(embed_dim, ffn_hidden_dim)
+
+    def forward(self, x, input_pos=None, mask=None):
+        x = x + self.attn(self.attn_norm(x), input_pos=input_pos, mask=mask)
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+class TransformerStack(nn.Module):
+    def __init__(self, num_blocks, embed_dim, num_heads, ffn_hidden_dim, max_seq_len):
+        super().__init__()
+        self.blocks = nn.ModuleList([
+            UnifiedTransformerBlock(embed_dim, num_heads, ffn_hidden_dim, max_seq_len)
+            for _ in range(num_blocks)
+        ])
+        self.norm = nn.LayerNorm(embed_dim)
+
+    def forward(self, x, input_pos=None, mask=None):
+        for block in self.blocks:
+            x = block(x, input_pos=input_pos, mask=mask)
+        return self.norm(x)
+
+class MLM_core(nn.Module):
+    def __init__(
+        self,
+        vocab_size: int,
+        embed_dim: int,
+        num_blocks: int,
+        num_heads: int,
+        ffn_hidden_dim: int,
+        output_dim: int,
+        max_seq_len: int,
+    ):
+        super().__init__()
+        self.embed = nn.Embedding(vocab_size, embed_dim)
+        self.transformer = TransformerStack(
+            num_blocks, embed_dim, num_heads, ffn_hidden_dim, max_seq_len
+        )
+        self.sequence_head = nn.Linear(embed_dim, output_dim, bias=True)
+
+
+    def forward(self, ids, mask=None, pad_token_id=0, input_pos=None):
+        x = self.embed(ids)
+        x = self.transformer(x, mask=mask, input_pos=input_pos)
+        # generate logits for MLM
+        # print(f"x shape: {x.shape}")  # Debugging line to check the shape of x
+        logits = self.sequence_head(x)
+        # print(f"logits shape: {logits.shape}")  # Debugging line to check the shape of logits
+
+        # mean pool but remove positions that have pad tokens
+        mean_pool = x.masked_fill(ids.unsqueeze(-1) == pad_token_id, 0).mean(dim=1)
+
+        outputs = {
+            'logits': logits,
+            'last_layer': x,
+            'mean_pool': mean_pool
+        }
+
+        return outputs
+
+class MLM_model(PreTrainedModel): # HF-facing class name
+    config_class = model_config
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = MLM_core(
+            vocab_size=config.vocab_size,
+            embed_dim=config.embed_dim,
+            num_blocks=config.num_blocks,
+            num_heads=config.num_heads,
+            ffn_hidden_dim=config.ffn_hidden_dim,
+            output_dim=config.output_dim,
+            max_seq_len=config.max_seq_len,
+        )
+        self.post_init()  # Initialize weights and apply final processing
+    
+    # if inputs are dictionary
+    def forward(self, x=None, **kwargs):
+        if isinstance(x, (BatchEncoding, Mapping)):
+            return self.model(x.get("input_ids"), mask=x.get("attention_mask"))
+        
+        if "input_ids" in kwargs or "attention_mask" in kwargs:
+            return self.model(kwargs.get("input_ids"), mask=kwargs.get("attention_mask"))
+        
+        return self.model(x)

__init__.py

@@ -0,0 +1,2 @@
+from .config import model_config
+from .ChemPepMTR import MLM_model

config.json

@@ -0,0 +1,16 @@
+{
+    "model_type": "MLM_model",
+    "size": "base",
+    "ffn_hidden_dim": 1024, 
+    "embed_dim": 768, 
+    "num_heads": 12, 
+    "num_blocks": 24,
+    "vocab_size": 405,
+    "output_dim": 405,
+    "max_seq_len": 2048,
+    "auto_map": {
+        "AutoConfig": "config.model_config",
+        "AutoModel": "ChemPepMTR.MLM_model"
+    },
+    "architectures": ["MLM_model"]
+}

config.py

@@ -0,0 +1,25 @@
+from transformers import PretrainedConfig
+
+class model_config(PretrainedConfig):
+    model_type = "MLM_model"
+    def __init__(
+        self,
+        ffn_hidden_dim = 1024,
+        embed_dim = 768,
+        num_heads = 12,
+        num_blocks = 24,
+        vocab_size = 405,
+        output_dim = 405,
+        max_seq_len = 2048,
+        size = "base",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.ffn_hidden_dim = ffn_hidden_dim
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.num_blocks = num_blocks
+        self.vocab_size = vocab_size
+        self.output_dim = output_dim
+        self.max_seq_len = max_seq_len
+        self.size = size

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ee9244b84ff636080bbb7ef874d4a2d5159f70a8518f67a5065ff13fcab54e3
+size 456074420