Upload Flowformer

configuration_flowformer.py

@@ -1,8 +1,40 @@
 from transformers import PretrainedConfig
 
 class FlowformerConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Flowformer`]. It is used to instantiate an
+    Flowformer model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of out model for ALL data (https://arxiv.org/abs/2108.10072).
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+    Args:
+        dim_hidden (`int`, *optional*, defaults to 32):
+            The dimensionality of the hidden states. dim_hidden must be divisible by num_heads i.e. dim_hidden%num_heads = 0.
+        num_heads (`int`, *optional*, defaults to 4):
+            The number of attention heads.
+        num_inds (`int`, *optional*, defaults to 32):
+            The number of inducing points.
+        hidden_layers (`int`, *optional*, defaults to 3):
+            The number of hidden layers.
+        layer_norm (`bool`, *optional*, defaults to True):
+            Whether to apply layer normalization.
+        dim_input (`int`, *optional*, defaults to 11):
+            The dimensionality of the input.
+        markers (`list`, *optional*, defaults to ["TIME", "FSC-A", "FSC-W", "SSC-A", "CD20", "CD10", "CD45", "CD34", "CD19", "CD38", "SY41"]):
+            The list of markers.
+    Example:
+    ```python
+    >>> from transformers import FlowformerConfig, FlowformerModel
+    >>> # Initializing a Flowformer configuration
+    >>> configuration = FlowformerConfig()
+    >>> # Initializing a model (with random weights) from the Flowformer configuration
+    >>> model = FlowformerModel(configuration)
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
+    """
     def __init__(self,
-                 dim_hidden: int=32,        # dim_hidden must be divisible by num_heads i.e. dim_hidden%num_heads = 0
+                 dim_hidden: int=32,
                  num_heads: int=4,
                  num_inds: int=16,
                  hidden_layers: int=3,

model_flowformer.py

@@ -4,6 +4,7 @@ import torch.nn.functional as F
 from torch.nn.functional import binary_cross_entropy_with_logits
 import math
 from transformers import PreTrainedModel
+from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward
 from .configuration_flowformer import FlowformerConfig
 
 
@@ -11,7 +12,7 @@ class MAB(nn.Module):
     """
     Multihead attention Block (MAB) from https://arxiv.org/abs/1810.00825.
     """
-    def __init__(self, dim_Q, dim_K, dim_V, num_heads, ln=False):
+    def __init__(self, dim_Q: int, dim_K: int, dim_V: int, num_heads: int, ln: int=False):
         super(MAB, self).__init__()
 
         self.dim_V = dim_V
@@ -47,7 +48,7 @@ class ISAB(nn.Module):
     """
     The Induced Set Attention Block (ISAB) from https://arxiv.org/abs/1810.00825.
     """
-    def __init__(self, dim_in, dim_out, num_heads, num_inds, ln=False):
+    def __init__(self, dim_in: int, dim_out: int, num_heads: int, num_inds: int, ln: bool=False):
         super(ISAB, self).__init__()
 
         self.I = nn.Parameter(torch.Tensor(1, num_inds, dim_out))
@@ -60,8 +61,30 @@ class ISAB(nn.Module):
 
         return self.mab1(X, H)
     
+FLOWFORMER_START_DOCSTRING = r"""
+    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
+    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
+    behavior.
+    Parameters:
+        config ([`FlowformerConfig`]): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+FLOWFORMER_INPUTS_DOCSTRING = r"""
+    Args:
+        tensor (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_markers)`):
+            The sample used as a basis for the prediction.
+        labels (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Optional ground truth lables for computing the loss.
+        markers (`list` of length `num_markers`):
+            The list of markers in the same order as the last dimension of the input tensor.
+"""
+
+
+@add_start_docstrings(FLOWFORMER_START_DOCSTRING)
 class Flowformer(PreTrainedModel):
-    def __init__(self, config):
+    def __init__(self, config: FlowformerConfig):
         super().__init__(config)
 
         # Load config
@@ -72,7 +95,7 @@ class Flowformer(PreTrainedModel):
         hidden_layers = config.hidden_layers
         layer_norm = config.layer_norm
         dim_output = 1  
-        self._pretrained_markers = config.markers or ["TIME", "FSC-A", "FSC-W", "SSC-A", "CD20", "CD10", "CD45", "CD34", "CD19", "CD38", "SY41"]
+        self._markers = config.markers
 
         # Define encoder
         enc_layers = [ISAB(dim_input, dim_hidden, num_heads, num_inds, ln=layer_norm)]
@@ -85,17 +108,18 @@ class Flowformer(PreTrainedModel):
         dec_layers = [nn.Linear(dim_input, dim_output)]
         self.dec = nn.Sequential(*dec_layers)
 
-    def pretrained_markers(self):
+    def markers(self):
         return self._pretrained_markers
 
-    def forward(self, tensor, labels=None, markers: list=None):
+    @add_start_docstrings_to_model_forward(FLOWFORMER_INPUTS_DOCSTRING)
+    def forward(self, tensor: torch.Tensor, labels: torch.Tensor=None, markers: list=None):
         B, L, M = tensor.shape
         if markers is not None:
             assert len(markers) == M, "Number of markers in x and markers must be identical"
 
-            zeros = torch.zeros((B, L, len(self._pretrained_markers)), device=tensor.device)
-            valid_markers = [m for m in markers if m in set(self._pretrained_markers).intersection(markers)]
-            idx = [self._pretrained_markers.index(m) for m in valid_markers]
+            zeros = torch.zeros((B, L, len(self.markers())), device=tensor.device)
+            valid_markers = [m for m in markers if m in set(self.markers()).intersection(markers)]
+            idx = [self.markers().index(m) for m in valid_markers]
             zeros[:, :, idx] = tensor # select only the markers that are in the pretrained model
             tensor = zeros
 
@@ -105,10 +129,12 @@ class Flowformer(PreTrainedModel):
         if labels is not None:
             return {
                 'loss': binary_cross_entropy_with_logits(output, labels),
-                'logits': output
+                'logits': output,
+                'prediction': torch.where(output > 0, 1, 0)
             }
         else:
             return {
-                'logits': output
+                'logits': output,
+                'prediction': torch.where(output > 0, 1, 0)
             }