Reformat code

SwissArmyTransformer/model/mixins.py

@@ -18,40 +18,44 @@ from SwissArmyTransformer.mpu.transformer import unscaled_init_method
 from .base_model import BaseMixin
 from .cached_autoregressive_model import CachedAutoregressiveMixin
 
+
 class PositionEmbeddingMixin(BaseMixin):
-    def __init__(self, additional_sequence_length, hidden_size, 
-                init_method_std=0.02, reinit_slice=slice(-1024, None)
-        ):
+    def __init__(self, additional_sequence_length, hidden_size,
+                 init_method_std=0.02, reinit_slice=slice(-1024, None)
+                 ):
         super(PositionEmbeddingMixin, self).__init__()
         self.reinit_slice = reinit_slice
         self.position_embeddings = torch.nn.Embedding(additional_sequence_length, hidden_size)
         torch.nn.init.normal_(self.position_embeddings.weight, mean=0.0, std=init_method_std)
+
     def reinit(self, *pre_mixins):
         old_weights = self.transformer.position_embeddings.weight.data[self.reinit_slice]
         old_len, hidden_size = old_weights.shape
         assert hidden_size == self.position_embeddings.weight.shape[-1]
         self.position_embeddings.weight.data.view(-1, old_len, hidden_size).copy_(old_weights)
 
+
 class AttentionMixin(BaseMixin):
     def __init__(self, num_layers,
-                hidden_size, 
-                init_method=unscaled_init_method(0.02),
-                output_layer_init_method=unscaled_init_method(0.02)
-        ):
+                 hidden_size,
+                 init_method=unscaled_init_method(0.02),
+                 output_layer_init_method=unscaled_init_method(0.02)
+                 ):
         super(AttentionMixin, self).__init__()
-        self.num_layers = num_layers # replace attention in the LAST n layers
+        self.num_layers = num_layers  # replace attention in the LAST n layers
         self.query_key_value = torch.nn.ModuleList(
-            [ColumnParallelLinear(hidden_size, 3*hidden_size,stride=3,
-                gather_output=False,init_method=init_method)
-                for layer_id in range(num_layers)
-            ])
+            [ColumnParallelLinear(hidden_size, 3 * hidden_size, stride=3,
+                                  gather_output=False, init_method=init_method)
+             for layer_id in range(num_layers)
+             ])
         self.dense = torch.nn.ModuleList(
             [RowParallelLinear(hidden_size,
-                hidden_size,
-                input_is_parallel=True,
-                init_method=output_layer_init_method)
-                for layer_id in range(num_layers)
-            ])
+                               hidden_size,
+                               input_is_parallel=True,
+                               init_method=output_layer_init_method)
+             for layer_id in range(num_layers)
+             ])
+
     def reinit(self, *pre_mixins):
         start_layer = len(self.transformer.layers) - self.num_layers
         assert start_layer >= 0