Upload FupBERT

config.json

@@ -0,0 +1,24 @@
+{
+  "_name_or_path": "../models/bert_saved",
+  "architectures": [
+    "FupBERT"
+  ],
+  "auto_map": {
+    "AutoConfig": "fup_bert_config.FupBERTConfig",
+    "AutoModel": "fup_bert.FupBERT"
+  },
+  "cls_idx": 1,
+  "dropout": 0.1,
+  "edge_idx": 2,
+  "model_type": "FupBERT",
+  "nhead": 12,
+  "nhid": 3072,
+  "ninp": 768,
+  "nlayers": 12,
+  "ntoken": 608,
+  "num_out": 1,
+  "padding_idx": 0,
+  "token_reduction": "mean",
+  "torch_dtype": "float32",
+  "transformers_version": "4.30.2"
+}

fup_bert.py

@@ -0,0 +1,46 @@
+"""
+© Battelle Memorial Institute 2023
+Made available under the GNU General Public License v 2.0
+
+BECAUSE THE PROGRAM IS LICENSED FREE OF CHARGE, THERE IS NO WARRANTY
+FOR THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW.  EXCEPT WHEN
+OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES
+PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED
+OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
+MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.  THE ENTIRE RISK AS
+TO THE QUALITY AND PERFORMANCE OF THE PROGRAM IS WITH YOU.  SHOULD THE
+PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF ALL NECESSARY SERVICING,
+REPAIR OR CORRECTION.
+"""
+
+import torch
+from transformers import PreTrainedModel
+
+from .fup_bert_config import FupBERTConfig
+from .fup_bert_model import FupBERTModel
+
+
+class FupBERT(PreTrainedModel):
+    """Hugging Face Wrapper"""
+    config_class = FupBERTConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = FupBERTModel(ntoken=config.ntoken,
+                                  ninp=config.ninp,
+                                  nhead=config.nhead,
+                                  nhid=config.nhid,
+                                  nlayers=config.nlayers,
+                                  token_reduction=config.token_reduction,
+                                  padding_idx=config.padding_idx,
+                                  cls_idx=config.cls_idx,
+                                  edge_idx=config.edge_idx,
+                                  num_out=config.num_out,
+                                  dropout=config.dropout,
+                                  )
+
+    def forward(self, src):
+        return self.model(src)
+
+    def load_params(self, pt_file):
+        self.model.load_state_dict(torch.load(pt_file))

fup_bert_config.py

@@ -0,0 +1,49 @@
+"""
+© Battelle Memorial Institute 2023
+Made available under the GNU General Public License v 2.0
+
+BECAUSE THE PROGRAM IS LICENSED FREE OF CHARGE, THERE IS NO WARRANTY
+FOR THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW.  EXCEPT WHEN
+OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES
+PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED
+OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
+MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.  THE ENTIRE RISK AS
+TO THE QUALITY AND PERFORMANCE OF THE PROGRAM IS WITH YOU.  SHOULD THE
+PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF ALL NECESSARY SERVICING,
+REPAIR OR CORRECTION.
+"""
+
+from transformers import PretrainedConfig
+
+
+class FupBERTConfig(PretrainedConfig):
+    model_type = "FupBERT"
+
+    def __init__(
+            self,
+            ntoken=608,
+            ninp=768,
+            nhead=12,
+            nhid=3072,
+            nlayers=12,
+            token_reduction='mean',
+            padding_idx=0,
+            cls_idx=1,
+            edge_idx=2,
+            num_out=1,
+            dropout=0.1,
+            **kwargs):
+
+        # Store the input parameters
+        self.ntoken = ntoken
+        self.ninp = ninp
+        self.nhead = nhead
+        self.nhid = nhid
+        self.nlayers = nlayers
+        self.token_reduction = token_reduction
+        self.padding_idx = padding_idx
+        self.cls_idx = cls_idx
+        self.edge_idx = edge_idx
+        self.num_out = num_out
+        self.dropout = dropout
+        super().__init__(**kwargs)

fup_bert_model.py

@@ -0,0 +1,175 @@
+"""
+© Battelle Memorial Institute 2023
+Made available under the GNU General Public License v 2.0
+
+BECAUSE THE PROGRAM IS LICENSED FREE OF CHARGE, THERE IS NO WARRANTY
+FOR THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW.  EXCEPT WHEN
+OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES
+PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED
+OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
+MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.  THE ENTIRE RISK AS
+TO THE QUALITY AND PERFORMANCE OF THE PROGRAM IS WITH YOU.  SHOULD THE
+PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF ALL NECESSARY SERVICING,
+REPAIR OR CORRECTION.
+"""
+
+import torch
+import torch.nn as nn
+
+from .positional_encoding import PositionalEncoding
+
+
+class FupBERTModel(nn.Module):
+    """
+    A class that extends torch.nn.Module that implements a custom Transformer
+    encoder model to create a single embedding for Fup prediction.
+    """
+
+    def __init__(
+        self,
+        ntoken,
+        ninp,
+        nhead,
+        nhid,
+        nlayers,
+        token_reduction,
+        padding_idx,
+        cls_idx,
+        edge_idx,
+        num_out,
+        dropout=0.1,
+    ):
+        """
+        Initializes a FubBERT object.
+
+        Parameters
+        ----------
+        ntoken : int
+            The maximum number of tokens the embedding layer should expect. This
+            is the same as the size of the vocabulary.
+        ninp : int
+            The hidden dimension that should be used for embedding and input
+            to the Transformer encoder.
+        nhead : int
+            The number of heads to use in the Transformer encoder.
+        nhid : int
+            The size of the hidden dimension to use throughout the Transformer
+            encoder.
+        nlayers : int
+            The number of layers to use in a single head of the Transformer
+            encoder.
+        token_reduction : str
+            The type of token reduction to use. This can be either 'mean' or
+            'cls'.
+        padding_idx : int
+            The index used as padding for the input sequences.
+        cls_idx : int
+            The index used as the cls token for the input sequences.
+        edge_idx : int
+            The index used as the edge token for the input sequences.
+        num_out : int
+            The number of outputs to predict with the model.
+        dropout : float, optional
+            The fractional dropout to apply to the model. The default is 0.1.
+
+        Returns
+        -------
+        None.
+
+        """
+        super(FupBERTModel, self).__init__()
+        # Store the input parameters
+        self.ntoken = ntoken
+        self.ninp = ninp
+        self.nhead = nhead
+        self.nhid = nhid
+        self.nlayers = nlayers
+        self.token_reduction = token_reduction
+        self.padding_idx = padding_idx
+        self.cls_idx = cls_idx
+        self.edge_idx = edge_idx
+        self.num_out = num_out
+        self.dropout = dropout
+        # Set the model parameters
+        self.model_type = "Transformer Encoder"
+        self.embedding = nn.Embedding(
+            self.ntoken, self.ninp, padding_idx=self.padding_idx
+        )
+        self.pos_encoder = PositionalEncoding(self.ninp, self.dropout)
+        encoder_layers = nn.TransformerEncoderLayer(
+            self.ninp,
+            self.nhead,
+            self.nhid,
+            self.dropout,
+            activation="gelu",
+            batch_first=True,
+        )
+        self.transformer_encoder = nn.TransformerEncoder(encoder_layers, self.nlayers)
+        self.pred_head = nn.Linear(self.ninp, self.num_out)
+
+    def _generate_src_key_mask(self, src):
+        mask = src == self.padding_idx
+        mask = mask.type(torch.bool)
+
+        return mask
+
+    def forward(self, src):
+        """
+        Perform a forward pass of the module.
+
+        Parameters
+        ----------
+        src : tensor
+            The input tensor. The shape should be (batch size, sequence length).
+
+        Returns
+        -------
+        output : tensor
+            The output tensor. The shape will be (batch size, num_out).
+
+        """
+        src = self.get_embeddings(src)
+        output = self.pred_head(src)
+
+        return output
+
+    def get_embeddings(self, src):
+        """
+        Perform a forward pass of the module excluding the classification layers. This
+        will return the embeddings from the encoder.
+
+        Parameters
+        ----------
+        src : tensor
+            The input tensor. The shape should be (batch size, sequence length).
+
+        Returns
+        -------
+        embeds : tensor
+            The output tensor of sequence embeddings. The shape should be
+            (batch size, self.ninp)
+        """
+        src_mask = self._generate_src_key_mask(src)
+        x = self.embedding(src)
+        x = self.pos_encoder(x)
+        x = self.transformer_encoder(x, src_key_padding_mask=src_mask)
+        # Mask the data based on the token reduction strategy
+        if self.token_reduction == "mean":
+            pad_mask = src == self.padding_idx
+            cls_mask = src == self.cls_idx
+            edge_mask = src == self.edge_idx
+            mask = torch.logical_or(pad_mask, cls_mask)
+            mask = torch.logical_or(mask, edge_mask)
+            # Apply the mask
+            x[mask[:, : x.shape[1]]] = torch.nan
+            # Take the mean of the embeddings
+            embeds = torch.nanmean(x, dim=1)
+        elif self.token_reduction == "cls":
+            embeds = x[:, 0, :]
+        else:
+            raise ValueError(
+                "Token reduction must be mean or cls. "
+                "Recieved {}".format(self.token_reduction)
+            )
+
+        return embeds

positional_encoding.py

@@ -0,0 +1,77 @@
+"""
+© Battelle Memorial Institute 2023
+Made available under the GNU General Public License v 2.0
+
+BECAUSE THE PROGRAM IS LICENSED FREE OF CHARGE, THERE IS NO WARRANTY
+FOR THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW.  EXCEPT WHEN
+OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES
+PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED
+OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
+MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.  THE ENTIRE RISK AS
+TO THE QUALITY AND PERFORMANCE OF THE PROGRAM IS WITH YOU.  SHOULD THE
+PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF ALL NECESSARY SERVICING,
+REPAIR OR CORRECTION.
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class PositionalEncoding(nn.Module):
+    """
+    A class that extends torch.nn.Module that applies positional encoding
+    for use in the Transformer architecture.
+    """
+
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        """
+        Initializes a PositionalEncoding object.
+
+        Parameters
+        ----------
+        d_model : int
+            The size of the model's embedding dimension.
+        dropout : float, optional
+            The fractional dropout to apply to the embedding. The default is 0.1.
+        max_len : int, optional
+            The maximum potential input sequnce length. The default is 5000.
+
+        Returns
+        -------
+        None.
+
+        """
+        super(PositionalEncoding, self).__init__()
+        # Create the dropout
+        self.dropout = nn.Dropout(p=dropout)
+        # Create the encoding
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        """
+        Perform a forward pass of the module.
+
+        Parameters
+        ----------
+        x : tensor
+            The input tensor to apply the positional encoding to.
+
+        Returns
+        -------
+        tensor
+            The resulting tensor after applying the positional encoding to the
+            input.
+
+        """
+        x = x + self.pe[:, : x.size(1)]
+
+        return self.dropout(x)

pytorch_model.bin

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