black formatting

nfqa_model.py

@@ -1,11 +1,11 @@
-from typing import Sequence, Optional, Union, Tuple
+from typing import Optional, Sequence, Tuple, Union
 
 import torch
 from torch import nn
-from torch.nn import functional, CrossEntropyLoss
+from torch.nn import CrossEntropyLoss, functional
 from transformers import RobertaConfig
 from transformers.modeling_outputs import SequenceClassifierOutput
-from transformers.models.roberta.modeling_roberta import RobertaModel, RobertaPreTrainedModel, RobertaPooler
+from transformers.models.roberta.modeling_roberta import RobertaModel, RobertaPooler, RobertaPreTrainedModel
 
 
 class MishActivation(nn.Module):
@@ -15,28 +15,24 @@ class MishActivation(nn.Module):
 
 class NFQAClassificationHead(nn.Module):
     def __init__(
-            self, input_dim: int, num_labels: int, hidden_dims: Sequence[int] = (768, 512), dropout: float = 0.0,
+        self, input_dim: int, num_labels: int, hidden_dims: Sequence[int] = (768, 512), dropout: float = 0.0,
     ) -> None:
         super().__init__()
 
-        self.linear_layers = nn.Sequential(
-            *(nn.Linear(input_dim, dim) for dim in hidden_dims)
-        )
+        self.linear_layers = nn.Sequential(*(nn.Linear(input_dim, dim) for dim in hidden_dims))
         self.classification_layer = torch.nn.Linear(hidden_dims[-1], num_labels)
         self.activations = [MishActivation()] * len(hidden_dims)
         self.dropouts = [torch.nn.Dropout(p=dropout)] * len(hidden_dims)
 
     def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         output = inputs
-        for layer, activation, dropout in zip(
-                self.linear_layers, self.activations, self.dropouts
-        ):
+        for layer, activation, dropout in zip(self.linear_layers, self.activations, self.dropouts):
             output = dropout(activation(layer(output)))
         return self.classification_layer(output)
 
 
 class RobertaNFQAClassification(RobertaPreTrainedModel):
-    _keys_to_ignore_on_load_missing = [r"position_ids"]
+    _keys_to_ignore_on_load_missing = [r'position_ids']
     _DROPOUT = 0.0
 
     def __init__(self, config: RobertaConfig):
@@ -51,19 +47,18 @@ class RobertaNFQAClassification(RobertaPreTrainedModel):
 
         self.init_weights()
 
-
     def forward(
-            self,
-            input_ids: Optional[torch.LongTensor] = None,
-            attention_mask: Optional[torch.FloatTensor] = None,
-            token_type_ids: Optional[torch.LongTensor] = None,
-            position_ids: Optional[torch.LongTensor] = None,
-            head_mask: Optional[torch.FloatTensor] = None,
-            inputs_embeds: Optional[torch.FloatTensor] = None,
-            labels: Optional[torch.LongTensor] = None,
-            output_attentions: Optional[bool] = None,
-            output_hidden_states: Optional[bool] = None,
-            return_dict: Optional[bool] = None,
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
     ) -> Union[Tuple[torch.Tensor, ...], SequenceClassifierOutput]:
         r"""
         labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
@@ -98,8 +93,5 @@ class RobertaNFQAClassification(RobertaPreTrainedModel):
             return ((loss,) + output) if loss is not None else output
 
         return SequenceClassifierOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
+            loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions,
         )