Upload Flowformer

configuration_flowformer.py

@@ -2,7 +2,7 @@ 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
+    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
@@ -22,16 +22,6 @@ class FlowformerConfig(PretrainedConfig):
             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,

model_flowformer.py

@@ -4,7 +4,6 @@ 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
 
 
@@ -61,29 +60,17 @@ class ISAB(nn.Module):
 
         return self.mab1(X, H)
     
-FLOWFORMER_START_DOCSTRING = r"""
+
+class Flowformer(PreTrainedModel):
+    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.
+        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: FlowformerConfig):
         super().__init__(config)
 
@@ -111,8 +98,16 @@ class Flowformer(PreTrainedModel):
     def markers(self):
         return self._pretrained_markers
 
-    @add_start_docstrings_to_model_forward(FLOWFORMER_INPUTS_DOCSTRING)
     def forward(self, tensor: torch.Tensor, labels: torch.Tensor=None, markers: list=None):
+        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.
+        """
         B, L, M = tensor.shape
         if markers is not None:
             assert len(markers) == M, "Number of markers in x and markers must be identical"